validate actual bytes written in MultiSetSend

quipo · May 23, 2011 · ffa2cf2 · ffa2cf2
1 parent 0a54ad4
commit ffa2cf2
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 6 deletions.
diff --git a/core/src/main/scala/kafka/network/Transmission.scala b/core/src/main/scala/kafka/network/Transmission.scala
@@ -84,18 +84,27 @@ private[kafka] trait Send extends Transmission {
 /**
  * A set of composite sends, sent one after another
  */
-class MultiSend[S <: Send](val sends: List[S]) extends Send {
-
+abstract class MultiSend[S <: Send](val sends: List[S]) extends Send {
+  val expectedBytesToWrite: Int
   private var current = sends
-
+  var totalWritten = 0
+
   def writeTo(channel: WritableByteChannel): Int = {
 	  expectIncomplete
     val written = current.head.writeTo(channel)
+    totalWritten += written
     if(current.head.complete)
       current = current.tail
     written
   }
 
-  def complete = current == Nil
-
+  def complete: Boolean = {
+    if (current == Nil) {
+      if (totalWritten != expectedBytesToWrite)
+        logger.error("mismatch in sending bytes over socket; expected: " + expectedBytesToWrite + " actual: " + totalWritten)
+      return true
+    }
+    else
+      return false
+  }
 }
diff --git a/core/src/main/scala/kafka/server/MultiMessageSetSend.scala b/core/src/main/scala/kafka/server/MultiMessageSetSend.scala
@@ -29,7 +29,9 @@ import kafka.utils._
 private[server] class MultiMessageSetSend(val sets: List[MessageSetSend]) extends MultiSend(new ByteBufferSend(6) :: sets) {
 
   val buffer = this.sends.head.asInstanceOf[ByteBufferSend].buffer
-  buffer.putInt(2 + sets.foldLeft(0)(_ + _.sendSize))
+  val allMessageSetSize: Int = sets.foldLeft(0)(_ + _.sendSize)
+  val expectedBytesToWrite: Int = 4 + 2 + allMessageSetSize
+  buffer.putInt(2 + allMessageSetSize)
   buffer.putShort(0)
   buffer.rewind()