Improve chunk content decoder

Simplifies some code around the decoding and reduced the risk of leaks
by avoiding retains.

*Removed the ChunkHandler, instead set a custom Cumulator in the decoder

*Make content type dispatcher an InboundHanlderAdapter. The previous
implementation had weird sideffects and used retain which increases the
risk of a leak.

enterprise/causal-clustering/src/main/java/org/neo4j/causalclustering/core/consensus/protocol/v2/RaftProtocolServerInstallerV2.java

@@ -31,7 +31,6 @@
 import java.util.stream.Collectors;
 
 import org.neo4j.causalclustering.messaging.marshalling.v2.ContentTypeProtocol;
-import org.neo4j.causalclustering.messaging.marshalling.v2.decoding.ChunkHandler;
 import org.neo4j.causalclustering.messaging.marshalling.v2.decoding.ContentTypeDispatcher;
 import org.neo4j.causalclustering.messaging.marshalling.v2.decoding.DecodingDispatcher;
 import org.neo4j.causalclustering.messaging.marshalling.v2.decoding.RaftMessageComposer;
@@ -77,13 +76,11 @@ public void install( Channel channel ) throws Exception
     {
 
         ContentTypeProtocol contentTypeProtocol = new ContentTypeProtocol();
-        ChunkHandler chunkHandler = new ChunkHandler();
-        DecodingDispatcher decodingDispatcher = new DecodingDispatcher( contentTypeProtocol, logProvider, chunkHandler );
+        DecodingDispatcher decodingDispatcher = new DecodingDispatcher( contentTypeProtocol, logProvider );
         pipelineBuilderFactory
                 .server( channel, log )
                 .modify( modifiers )
                 .addFraming()
-                .onClose( chunkHandler::close )
                 .add( "raft_content_type_dispatcher", new ContentTypeDispatcher( contentTypeProtocol ) )
                 .add( "raft_component_decoder", decodingDispatcher )
                 .add( "raft_content_decoder", new ReplicatedContentDecoder( contentTypeProtocol ) )

enterprise/causal-clustering/src/main/java/org/neo4j/causalclustering/messaging/marshalling/ChunkedReplicatedContent.java

@@ -34,7 +34,7 @@
 
 public class ChunkedReplicatedContent implements Marshal, ChunkedInput<ByteBuf>
 {
-    private static final int METADATA_SIZE = Integer.BYTES + 1;
+    private static final int METADATA_SIZE = 1;
 
     private final byte contentType;
     private final ByteBufChunkHandler byteBufChunkHandler;
@@ -117,11 +117,9 @@ private int metadataSize( boolean isFirstChunk )
 
     private ByteBuf writeMetadata( boolean isFirstChunk, ByteBufAllocator allocator, ByteBuf data )
     {
-        int length = data.writerIndex();
         int capacity = metadataSize( isFirstChunk );
         ByteBuf metaData = allocator.buffer( capacity, capacity );
         metaData.writeBoolean( byteBufAwareMarshal.isEndOfInput() );
-        metaData.writeInt( length );
         if ( isFirstChunk )
         {
             metaData.writeByte( contentType );

enterprise/causal-clustering/src/main/java/org/neo4j/causalclustering/messaging/marshalling/CoreReplicatedContentMarshal.java

@@ -22,6 +22,8 @@
  */
 package org.neo4j.causalclustering.messaging.marshalling;
 
+import io.netty.buffer.ByteBuf;
+
 import java.io.IOException;
 import java.util.Collection;
 import java.util.LinkedList;
@@ -43,7 +45,6 @@
 import org.neo4j.causalclustering.core.state.storage.SafeChannelMarshal;
 import org.neo4j.causalclustering.messaging.EndOfStreamException;
 import org.neo4j.causalclustering.messaging.NetworkReadableClosableChannelNetty4;
-import org.neo4j.causalclustering.messaging.marshalling.v2.decoding.ChunkHandler;
 import org.neo4j.storageengine.api.ReadableChannel;
 import org.neo4j.storageengine.api.WritableChannel;
 
@@ -110,30 +111,16 @@ else if ( content instanceof DummyRequest )
         }
     }
 
-    public ContentBuilder<ReplicatedContent> decode( ChunkHandler.ComposedChunks composedChunks ) throws IOException, EndOfStreamException
+    public ContentBuilder<ReplicatedContent> decode( byte contentType, ByteBuf buffer ) throws IOException, EndOfStreamException
     {
-        switch ( composedChunks.contentType() )
+        switch ( contentType )
         {
         case TX_CONTENT_TYPE:
         {
-            try
-            {
-                return ContentBuilder.finished( ReplicatedTransactionSerializer.decode( composedChunks.content() ) );
-            }
-            finally
-            {
-                composedChunks.release();
-            }
+            return ContentBuilder.finished( ReplicatedTransactionSerializer.decode( buffer ) );
         }
         default:
-            try
-            {
-                return read( composedChunks.contentType(), new NetworkReadableClosableChannelNetty4( composedChunks.content() ) );
-            }
-            finally
-            {
-                composedChunks.release();
-            }
+            return read( contentType, new NetworkReadableClosableChannelNetty4( buffer ) );
         }
     }
 

enterprise/causal-clustering/src/main/java/org/neo4j/causalclustering/messaging/marshalling/v2/decoding/ContentTypeDispatcher.java

@@ -24,14 +24,13 @@
 
 import io.netty.buffer.ByteBuf;
 import io.netty.channel.ChannelHandlerContext;
-import io.netty.handler.codec.ByteToMessageDecoder;
-
-import java.util.List;
+import io.netty.channel.ChannelInboundHandlerAdapter;
+import io.netty.util.ReferenceCountUtil;
 
 import org.neo4j.causalclustering.catchup.Protocol;
 import org.neo4j.causalclustering.messaging.marshalling.v2.ContentType;
 
-public class ContentTypeDispatcher extends ByteToMessageDecoder
+public class ContentTypeDispatcher extends ChannelInboundHandlerAdapter
 {
     private final Protocol<ContentType> contentTypeProtocol;
 
@@ -41,17 +40,18 @@ public ContentTypeDispatcher( Protocol<ContentType> contentTypeProtocol )
     }
 
     @Override
-    protected void decode( ChannelHandlerContext ctx, ByteBuf in, List<Object> out )
+    public void channelRead( ChannelHandlerContext ctx, Object msg )
     {
         if ( contentTypeProtocol.isExpecting( ContentType.ContentType ) )
         {
-            byte messageCode = in.readByte();
+            byte messageCode = ((ByteBuf) msg).readByte();
             ContentType contentType = getContentType( messageCode );
             contentTypeProtocol.expect( contentType );
+            ReferenceCountUtil.release( msg );
         }
         else
         {
-            ctx.fireChannelRead( in.retain() );
+            ctx.fireChannelRead( msg );
         }
     }
 

enterprise/causal-clustering/src/main/java/org/neo4j/causalclustering/messaging/marshalling/v2/decoding/DecodingDispatcher.java

@@ -35,13 +35,9 @@
 
 public class DecodingDispatcher extends RequestDecoderDispatcher<ContentType>
 {
-    private final ChunkHandler chunkHandler;
-
-    public DecodingDispatcher( Protocol<ContentType> protocol, LogProvider logProvider, ChunkHandler chunkHandler )
+    public DecodingDispatcher( Protocol<ContentType> protocol, LogProvider logProvider )
     {
         super( protocol, logProvider );
-        this.chunkHandler = chunkHandler;
-
         register( ContentType.ContentType, new ByteToMessageDecoder()
         {
             @Override
@@ -54,14 +50,7 @@ protected void decode( ChannelHandlerContext ctx, ByteBuf in, List<Object> out )
             }
         } );
         register( ContentType.RaftLogEntryTerms, new RaftLogEntryTermsDecoder( protocol ) );
-        register( ContentType.ReplicatedContent, new ReplicatedContentChunkDecoder( chunkHandler ) );
+        register( ContentType.ReplicatedContent, new ReplicatedContentChunkDecoder() );
         register( ContentType.Message, new RaftMessageDecoder( protocol ) );
     }
-
-    @Override
-    public void channelInactive( ChannelHandlerContext ctx )
-    {
-        chunkHandler.close();
-        ctx.fireChannelInactive();
-    }
 }

enterprise/causal-clustering/src/main/java/org/neo4j/causalclustering/messaging/marshalling/v2/decoding/ReplicatedContentChunkDecoder.java

@@ -23,6 +23,7 @@
 package org.neo4j.causalclustering.messaging.marshalling.v2.decoding;
 
 import io.netty.buffer.ByteBuf;
+import io.netty.buffer.ByteBufAllocator;
 import io.netty.channel.ChannelHandlerContext;
 import io.netty.handler.codec.ByteToMessageDecoder;
 
@@ -33,35 +34,35 @@
 public class ReplicatedContentChunkDecoder extends ByteToMessageDecoder
 {
     private final CoreReplicatedContentMarshal contentMarshal = new CoreReplicatedContentMarshal();
-    private final ChunkHandler chunkHandler;
 
-    ReplicatedContentChunkDecoder( ChunkHandler chunkHandler )
+    ReplicatedContentChunkDecoder()
     {
-        this.chunkHandler = chunkHandler;
+        setCumulator( new ContentChunkCumulator() );
     }
 
     @Override
     protected void decode( ChannelHandlerContext ctx, ByteBuf in, List<Object> out ) throws Exception
     {
-        try
+        in.markReaderIndex();
+        boolean isLast = in.readBoolean();
+        if ( isLast )
         {
-            ChunkHandler.ComposedChunks complete = chunkHandler.handle( in );
-            if ( complete != null )
-            {
-                out.add( contentMarshal.decode( complete ) );
-            }
+            out.add( contentMarshal.decode( in.readByte(), in ) );
         }
-        catch ( Throwable e )
+        else
         {
-            try
-            {
-                chunkHandler.close();
-            }
-            catch ( Throwable t1 )
-            {
-                e.addSuppressed( t1 );
-            }
-            throw e;
+            in.resetReaderIndex();
+        }
+    }
+
+    private static class ContentChunkCumulator implements Cumulator
+    {
+        @Override
+        public ByteBuf cumulate( ByteBufAllocator alloc, ByteBuf cumulation, ByteBuf in )
+        {
+            boolean isLast = in.readBoolean();
+            cumulation.setBoolean( 0, isLast );
+            return COMPOSITE_CUMULATOR.cumulate( alloc, cumulation, in.slice() );
         }
     }
 }

enterprise/causal-clustering/src/test/java/org/neo4j/causalclustering/messaging/marshalling/ChunkedReplicatedContentTest.java

@@ -50,24 +50,18 @@ public void shouldProvideExpectedMetaData() throws IOException
 
         // is not last
         assertFalse( byteBuf.readBoolean() );
-        // first chunk unknown length
-        assertEquals( 8, byteBuf.readInt() );
         // first chunk has content
         assertEquals( (byte) 1, byteBuf.readByte() );
         byteBuf.release();
 
         byteBuf = replicatedContent.readChunk( allocator );
         // is not last
         assertFalse( byteBuf.readBoolean() );
-        // second chunk has length -1
-        assertEquals( 8, byteBuf.readInt() );
         byteBuf.release();
 
         byteBuf = replicatedContent.readChunk( allocator );
         // is last
         assertTrue( byteBuf.readBoolean() );
-        // third chunk has a length
-        assertEquals( 8, byteBuf.readInt() );
         byteBuf.release();
 
         assertNull( replicatedContent.readChunk( allocator ) );