Reduce excessive off-heap memory consumption caused by occasional large keys.

Author: likun

src/main/java/io/lettuce/core/protocol/CommandHandler.java

@@ -55,6 +55,7 @@
 import io.lettuce.core.tracing.Tracer;
 import io.lettuce.core.tracing.Tracing;
 import io.netty.buffer.ByteBuf;
+import io.netty.buffer.ByteBufAllocator;
 import io.netty.channel.Channel;
 import io.netty.channel.ChannelDuplexHandler;
 import io.netty.channel.ChannelHandler;
@@ -93,6 +94,8 @@ public class CommandHandler extends ChannelDuplexHandler implements HasQueuedCom
 
     private static final AtomicLong COMMAND_HANDLER_COUNTER = new AtomicLong();
 
+    private static final int DEFAULT_READER_BUFFER_CAPACITY = 8192 * 8;
+
     private final ClientOptions clientOptions;
 
     private final ClientResources clientResources;
@@ -139,6 +142,10 @@ public class CommandHandler extends ChannelDuplexHandler implements HasQueuedCom
 
     private Tracing.Endpoint tracedEndpoint;
 
+    private ByteBuf tmpReadBuffer;
+
+    private ByteBufAllocator byteBufAllocator;
+
     /**
      * Initialize a new instance that handles commands from the supplied queue.
      *
@@ -222,7 +229,8 @@ public void channelRegistered(ChannelHandlerContext ctx) throws Exception {
 
         setState(LifecycleState.REGISTERED);
 
-        readBuffer = ctx.alloc().buffer(8192 * 8);
+        byteBufAllocator = ctx.alloc();
+        readBuffer = ctx.alloc().buffer(DEFAULT_READER_BUFFER_CAPACITY);
         rsm = new RedisStateMachine();
         ctx.fireChannelRegistered();
     }
@@ -614,6 +622,15 @@ public void channelRead(ChannelHandlerContext ctx, Object msg) throws Exception
             if (traceEnabled) {
                 logger.trace("{} Buffer: {}", logPrefix(), input.toString(Charset.defaultCharset()).trim());
             }
+            // if buffer capacity larger than default capacity, then create a new buffer with double capacity
+            // in most cases, key size is less than 64k
+            if (readBuffer.capacity() == DEFAULT_READER_BUFFER_CAPACITY && readBuffer.writableBytes() < input.readableBytes()
+                    && byteBufAllocator != null) {
+                ByteBuf byteBuf = byteBufAllocator.directBuffer(readBuffer.capacity() << 1);
+                byteBuf.writeBytes(readBuffer);
+                tmpReadBuffer = readBuffer;
+                readBuffer = byteBuf;
+            }
 
             readBuffer.touch("CommandHandler.read(…)");
             readBuffer.writeBytes(input);
@@ -661,6 +678,8 @@ protected void decode(ChannelHandlerContext ctx, ByteBuf buffer) throws Interrup
 
                     ctx.close();
                     throw e;
+                } finally {
+                    releaseLargeBufferIfNecessary(buffer);
                 }
 
                 hasDecodeProgress = false;
@@ -704,14 +723,26 @@ protected void decode(ChannelHandlerContext ctx, ByteBuf buffer) throws Interrup
                             complete(command);
                         } catch (Exception e) {
                             logger.warn("{} Unexpected exception during request: {}", logPrefix, e.toString(), e);
+                        } finally {
+                            releaseLargeBufferIfNecessary(buffer);
                         }
                     }
                 }
                 afterDecode(ctx, command);
             }
         }
+        if (buffer.refCnt() > 0) {
+            decodeBufferPolicy.afterDecoding(buffer);
+        }
+    }
 
-        decodeBufferPolicy.afterDecoding(buffer);
+    private void releaseLargeBufferIfNecessary(ByteBuf buffer) {
+        if (this.tmpReadBuffer != null) {
+            buffer.release();
+            this.readBuffer = tmpReadBuffer;
+            this.readBuffer.clear();
+            this.tmpReadBuffer = null;
+        }
     }
 
     protected void notifyPushListeners(PushMessage notification) {
@@ -734,7 +765,7 @@ protected void notifyPushListeners(PushMessage notification) {
      * @return
      */
     protected boolean canDecode(ByteBuf buffer) {
-        return buffer.isReadable() && (isMessageDecode() || isPushDecode(buffer));
+        return buffer.refCnt()>0 && buffer.isReadable() && (isMessageDecode() || isPushDecode(buffer));
     }
 
     private boolean isPushMessage(ByteBuf buffer) {

src/test/java/io/lettuce/core/protocol/CommandHandlerUnitTests.java

@@ -30,6 +30,7 @@
 import java.net.Inet4Address;
 import java.net.InetSocketAddress;
 import java.nio.ByteBuffer;
+import java.nio.charset.StandardCharsets;
 import java.time.Duration;
 import java.util.Arrays;
 import java.util.Collection;
@@ -65,6 +66,7 @@
 import io.lettuce.core.output.KeyValueListOutput;
 import io.lettuce.core.output.StatusOutput;
 import io.lettuce.core.output.ValueListOutput;
+import io.lettuce.core.output.ValueOutput;
 import io.lettuce.core.resource.ClientResources;
 import io.lettuce.core.tracing.Tracing;
 import io.lettuce.test.Delay;
@@ -650,4 +652,49 @@ void shouldHandleNullBuffers() throws Exception {
         sut.channelUnregistered(context);
     }
 
+    /**
+     * if large keys are received ,the large buffer will created and then released
+     */
+    @Test
+    void shouldLargeBufferCreatedAndRelease() throws Exception {
+        ChannelPromise channelPromise = new DefaultChannelPromise(channel, ImmediateEventExecutor.INSTANCE);
+        channelPromise.setSuccess();
+        sut.channelRegistered(context);
+        sut.channelActive(context);
+        sut.getStack().add(new Command<>(CommandType.GET, new ValueOutput<String, String>(StringCodec.UTF8)));
+
+        ByteBuf internalBuffer = ReflectionTestUtils.getField(sut, "readBuffer");
+
+        // step1 Receive First TCP Packet
+        ByteBuf msg = context.alloc().buffer(13);
+        // 1+5+2+7 ($+length+\r\n+len(value))
+        msg.writeBytes("$65536\r\nval_abc".getBytes(StandardCharsets.UTF_8));
+        sut.channelRead(context, msg);
+
+        int markedReaderIndex = ReflectionTestUtils.getField(internalBuffer, "markedReaderIndex");
+        assertThat(markedReaderIndex).isEqualTo(8);
+        assertThat(internalBuffer.readerIndex()).isEqualTo(8);
+        assertThat(internalBuffer.writerIndex()).isEqualTo(15);
+
+        // step2 Receive Second TCP Packet
+        ByteBuf msg2 = context.alloc().buffer(64 * 1024);
+        StringBuilder sb = new StringBuilder();
+        // 65536-7
+        for (int i = 0; i < 65529; i++) {
+            sb.append((char) ('a' + i % 26));
+        }
+        sb.append("\r\n");
+        msg2.writeBytes(sb.toString().getBytes(StandardCharsets.UTF_8));
+        sut.channelRead(context, msg2);
+
+        // step3 Got Result: readBuffer.capacity = 64k and tmpBuffer is null
+        ByteBuf readBuffer = ReflectionTestUtils.getField(sut, "readBuffer");
+        assertThat(readBuffer.capacity()).isEqualTo(64 * 1024);
+        assertThat(readBuffer.readerIndex()).isZero();
+        assertThat(readBuffer.writerIndex()).isZero();
+
+        ByteBuf tmpBuffer = ReflectionTestUtils.getField(sut, "tmpReadBuffer");
+        assertThat(tmpBuffer).isNull();
+    }
+
 }