add gated auto read

Author: mhl-b

modules/transport-netty4/src/main/java/org/elasticsearch/http/netty4/GatedAutoRead.java

@@ -0,0 +1,98 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the "Elastic License
+ * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side
+ * Public License v 1"; you may not use this file except in compliance with, at
+ * your election, the "Elastic License 2.0", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+
+package org.elasticsearch.http.netty4;
+
+import io.netty.channel.Channel;
+import io.netty.channel.ChannelConfig;
+import io.netty.util.AttributeKey;
+
+import org.elasticsearch.core.Releasable;
+
+import java.util.BitSet;
+
+/**
+ * This class provides gated toggle to {@link io.netty.channel.ChannelConfig#setAutoRead(boolean)}.
+ * Channel handlers should get own instance of {@link Gate} to change autoRead config. Setting any
+ * Gate to false will turn off auto-read. The auto-read will turn-on once all gates are set to true,
+ * or closed ({@link Releasable#close()}).
+ * <br><br>
+ * This class is not thread-safe and should be used inside channels' {@link io.netty.channel.EventLoop}.
+ * <br><br>
+ * When channel pipeline has multiple handlers that toggle autoRead it might become problematic.
+ * For example, we use HTTP pipelining feature, that means we dont block incoming requests and
+ * several ones can be flow in the pipeline. Some request might wait for authentication and some hit by
+ * backpressure, abd both controlled by channel autoRead. The backpressure handler might allow unauthenticated
+ * content keep flowing to the Authentication handler.
+ */
+public class GatedAutoRead {
+
+    private static final AttributeKey<GatedAutoRead> GATED_AUTO_READ = AttributeKey.valueOf("GatedAutoRead");
+    private final Channel channel;
+    private final ChannelConfig config;
+    private final BitSet gateIds = new BitSet(1);
+    private final BitSet gateStates = new BitSet(1);
+
+    GatedAutoRead(Channel channel) {
+        this.channel = channel;
+        this.config = channel.config();
+    }
+
+    /**
+     * Returns a new instance to toggle channel's autoRead.
+     * AutoRead will be set to true when ALL gates toggled to true. Toggle-off applies immediately.
+     * Must be closed when goes out of scope, otherwise might leave channel permanently in disabled state.
+     */
+    public static Gate newGate(Channel channel) {
+        assert channel.eventLoop().inEventLoop();
+        var gates = channel.attr(GATED_AUTO_READ).get();
+        if (gates == null) {
+            gates = new GatedAutoRead(channel);
+            channel.attr(GATED_AUTO_READ).set(gates);
+        }
+        return gates.newGate();
+    }
+
+    Gate newGate() {
+        return new Gate();
+    }
+
+    public class Gate implements Releasable {
+        final int id;
+
+        Gate() {
+            this.id = gateIds.nextClearBit(0); // next unused id
+            gateIds.set(id, true);
+        }
+
+        /**
+         * Set {@link ChannelConfig#setAutoRead(boolean)}. Setting to false will take immediate effect.
+         * Setting to true will apply only when ALL gates are set to true.
+         */
+        public void set(boolean flag) {
+            assert channel.eventLoop().inEventLoop();
+            gateStates.set(id, flag == false); // gate=true means closed, autoRead=false
+            config.setAutoRead(gateStates.isEmpty());
+        }
+
+        /**
+         * @return {@link ChannelConfig#isAutoRead()}.
+         */
+        public boolean get() {
+            return config.isAutoRead();
+        }
+
+        @Override
+        public void close() {
+            set(true);
+            gateIds.set(id, false); // release id
+        }
+    }
+
+}

modules/transport-netty4/src/main/java/org/elasticsearch/http/netty4/Netty4HttpHeaderValidator.java

@@ -37,6 +37,7 @@ public class Netty4HttpHeaderValidator extends ChannelInboundHandlerAdapter {
 
     private final HttpValidator validator;
     private final ThreadContext threadContext;
+    private GatedAutoRead.Gate autoRead;
     private ArrayDeque<HttpObject> pending = new ArrayDeque<>(4);
     private State state = WAITING_TO_START;
 
@@ -49,6 +50,12 @@ State getState() {
         return state;
     }
 
+    @Override
+    public void channelActive(ChannelHandlerContext ctx) throws Exception {
+        autoRead = GatedAutoRead.newGate(ctx.channel());
+        super.channelActive(ctx);
+    }
+
     @SuppressWarnings("fallthrough")
     @Override
     public void channelRead(ChannelHandlerContext ctx, Object msg) throws Exception {
@@ -61,7 +68,7 @@ public void channelRead(ChannelHandlerContext ctx, Object msg) throws Exception
                 pending.add(ReferenceCountUtil.retain(httpObject));
                 requestStart(ctx);
                 assert state == QUEUEING_DATA;
-                assert ctx.channel().config().isAutoRead() == false;
+                assert autoRead.get() == false;
                 break;
             case QUEUEING_DATA:
                 pending.add(ReferenceCountUtil.retain(httpObject));
@@ -83,7 +90,7 @@ public void channelRead(ChannelHandlerContext ctx, Object msg) throws Exception
             case DROPPING_DATA_PERMANENTLY:
                 assert pending.isEmpty();
                 ReferenceCountUtil.release(httpObject); // consume without enqueuing
-                ctx.channel().config().setAutoRead(false);
+                autoRead.set(false);
                 break;
         }
     }
@@ -106,7 +113,7 @@ private void requestStart(ChannelHandlerContext ctx) {
         }
 
         state = QUEUEING_DATA;
-        ctx.channel().config().setAutoRead(false);
+        autoRead.set(false);
 
         if (httpRequest == null) {
             // this looks like a malformed request and will forward without validation
@@ -149,10 +156,10 @@ public void onFailure(Exception e) {
     private void forwardFullRequest(ChannelHandlerContext ctx) {
         Transports.assertDefaultThreadContext(threadContext);
         assert ctx.channel().eventLoop().inEventLoop();
-        assert ctx.channel().config().isAutoRead() == false;
+        assert autoRead.get() == false;
         assert state == QUEUEING_DATA;
 
-        ctx.channel().config().setAutoRead(true);
+        autoRead.set(true);
         boolean fullRequestForwarded = forwardData(ctx, pending);
 
         assert fullRequestForwarded || pending.isEmpty();
@@ -169,7 +176,7 @@ private void forwardFullRequest(ChannelHandlerContext ctx) {
     private void forwardRequestWithDecoderExceptionAndNoContent(ChannelHandlerContext ctx, Exception e) {
         Transports.assertDefaultThreadContext(threadContext);
         assert ctx.channel().eventLoop().inEventLoop();
-        assert ctx.channel().config().isAutoRead() == false;
+        assert autoRead.get() == false;
         assert state == QUEUEING_DATA;
 
         HttpObject messageToForward = pending.getFirst();
@@ -180,7 +187,7 @@ private void forwardRequestWithDecoderExceptionAndNoContent(ChannelHandlerContex
         }
         messageToForward.setDecoderResult(DecoderResult.failure(e));
 
-        ctx.channel().config().setAutoRead(true);
+        autoRead.set(true);
         ctx.fireChannelRead(messageToForward);
 
         assert fullRequestDropped || pending.isEmpty();

modules/transport-netty4/src/main/java/org/elasticsearch/http/netty4/Netty4HttpRequestBodyStream.java

@@ -33,6 +33,7 @@
 public class Netty4HttpRequestBodyStream implements HttpBody.Stream {
 
     private final Channel channel;
+    private final GatedAutoRead.Gate autoRead;
     private final ChannelFutureListener closeListener = future -> doClose();
     private final List<ChunkHandler> tracingHandlers = new ArrayList<>(4);
     private final ThreadContext threadContext;
@@ -48,10 +49,11 @@ public class Netty4HttpRequestBodyStream implements HttpBody.Stream {
 
     public Netty4HttpRequestBodyStream(Channel channel, ThreadContext threadContext) {
         this.channel = channel;
+        this.autoRead = GatedAutoRead.newGate(channel);
         this.threadContext = threadContext;
         this.requestContext = threadContext.newStoredContext();
         Netty4Utils.addListener(channel.closeFuture(), closeListener);
-        channel.config().setAutoRead(false);
+        autoRead.set(false);
     }
 
     @Override
@@ -142,7 +144,7 @@ private void send() {
             handler.onNext(bytesRef, hasLast);
         }
         if (hasLast) {
-            channel.config().setAutoRead(true);
+            autoRead.close();
             channel.closeFuture().removeListener(closeListener);
         }
     }
@@ -171,6 +173,6 @@ private void doClose() {
             buf = null;
             bufSize = 0;
         }
-        channel.config().setAutoRead(true);
+        autoRead.close();
     }
 }

modules/transport-netty4/src/test/java/org/elasticsearch/http/netty4/GatedAutoReadTests.java

@@ -0,0 +1,58 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the "Elastic License
+ * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side
+ * Public License v 1"; you may not use this file except in compliance with, at
+ * your election, the "Elastic License 2.0", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+
+package org.elasticsearch.http.netty4;
+
+import io.netty.channel.embedded.EmbeddedChannel;
+
+import org.elasticsearch.core.Releasables;
+import org.elasticsearch.test.ESTestCase;
+
+import java.util.stream.IntStream;
+
+public class GatedAutoReadTests extends ESTestCase {
+
+    public void testToggleGates() {
+        var chan = new EmbeddedChannel();
+        var config = chan.config();
+        config.setAutoRead(true);
+
+        var n = between(10, 20);
+        var gates = IntStream.range(0, n).mapToObj(i -> GatedAutoRead.newGate(chan)).toList();
+        assertTrue("new gate should preserve current value", config.isAutoRead());
+
+        var gateId1 = between(0, n);
+        var gateId2 = (gateId1 + 1) % n;
+
+        var autoRead1 = gates.get(gateId1);
+        var autoRead2 = gates.get(gateId2);
+
+        autoRead1.set(false);
+        assertFalse("any gate toggle-off should disable reads", config.isAutoRead());
+
+        autoRead2.set(false);
+        autoRead1.set(true);
+        assertFalse("toggle-on should not enable reads while other gate is off", config.isAutoRead());
+
+        autoRead2.close();
+        assertTrue("removing toggled-off gate should enable reads", config.isAutoRead());
+    }
+
+    public void testGateIdReuse() {
+        var chan = new EmbeddedChannel();
+        for (int i = 0; i < between(1, 100); i++) {
+            var n = between(1, 10);
+            var gates = IntStream.range(0, n).mapToObj(j -> GatedAutoRead.newGate(chan)).toList();
+            for (int j = 0; j < gates.size(); j++) {
+                assertEquals(j, gates.get(j).id);
+            }
+            Releasables.close(gates);
+        }
+    }
+}

modules/transport-netty4/src/test/java/org/elasticsearch/http/netty4/Netty4HttpHeaderValidatorTests.java

@@ -76,6 +76,7 @@ private void reset() {
         };
         netty4HttpHeaderValidator = new Netty4HttpHeaderValidator(validator, new ThreadContext(Settings.EMPTY));
         channel.pipeline().addLast(netty4HttpHeaderValidator);
+        channel.pipeline().fireChannelActive();
     }
 
     public void testValidationPausesAndResumesData() {