Assume chunk is last and re-write otherwise

neo4j · Jun 11, 2018 · ad022d6 · ad022d6
1 parent cb951f4
commit ad022d6
Showing 1 changed file with 9 additions and 7 deletions.
diff --git a/.../main/java/org/neo4j/causalclustering/messaging/marshalling/ChunkedReplicatedContent.java b/.../main/java/org/neo4j/causalclustering/messaging/marshalling/ChunkedReplicatedContent.java
@@ -39,7 +39,7 @@ public class ChunkedReplicatedContent implements Marshal, ChunkedInput<ByteBuf>
     private final byte contentType;
     private final ByteBufAwareMarshal byteBufAwareMarshal;
     private final int chunkSize;
-    private boolean lastByteWasWritten;
+    private boolean endOfInput;
     private int progress;
 
     public ChunkedReplicatedContent( byte contentType, ByteBufAwareMarshal byteBufAwareMarshal, int chunkSize )
@@ -68,7 +68,7 @@ public void marshal( WritableChannel channel ) throws IOException
     @Override
     public boolean isEndOfInput()
     {
-        return lastByteWasWritten;
+        return endOfInput;
     }
 
     @Override
@@ -86,15 +86,16 @@ public ByteBuf readChunk( ChannelHandlerContext ctx ) throws IOException
     @Override
     public ByteBuf readChunk( ByteBufAllocator allocator ) throws IOException
     {
-        boolean endOfInput = isEndOfInput();
-        if ( endOfInput )
+        if ( isEndOfInput() )
         {
             return null;
         }
+        // assume this is the last chunk
+        boolean lastChunk = true;
         ByteBuf buffer = allocator.buffer( chunkSize );
         try
         {
-            buffer.writeBoolean( endOfInput );
+            buffer.writeBoolean( lastChunk );
             if ( progress() == 0 )
             {
                 // extra metadata on first chunk
@@ -103,10 +104,11 @@ public ByteBuf readChunk( ByteBufAllocator allocator ) throws IOException
             }
             if ( !byteBufAwareMarshal.encode( buffer ) )
             {
-                lastByteWasWritten = true;
+                this.endOfInput = true;
             }
-            if ( isEndOfInput() != endOfInput )
+            if ( isEndOfInput() != lastChunk )
             {
+                // status changed after writing to buffer.
                 buffer.setBoolean( 0, isEndOfInput() );
             }
             progress += buffer.readableBytes();