fix: ignore TrafficCop calls once the connection is back in the pool

processing/src/main/java/org/apache/druid/java/util/http/client/NettyHttpClient.java

@@ -240,6 +240,7 @@ public <Intermediate, Final> ListenableFuture<Final> go(
           private final Object watermarkLock = new Object();
           private long suspendWatermark = -1;
           private long resumeWatermark = -1;
+          private boolean returnedToPool = false;
 
           @Override
           protected void channelRead0(ChannelHandlerContext ctx, HttpObject msg)
@@ -274,6 +275,9 @@ protected void channelRead0(ChannelHandlerContext ctx, HttpObject msg)
                   public long resume(long resumeChunkNum)
                   {
                     synchronized (watermarkLock) {
+                      if (returnedToPool) {
+                        return 0;
+                      }
                       resumeWatermark = Math.max(resumeWatermark, resumeChunkNum);
 
                       if (suspendWatermark >= 0 && resumeWatermark >= suspendWatermark) {
@@ -292,8 +296,12 @@ public long resume(long resumeChunkNum)
                   @Override
                   public void abort()
                   {
-                    log.debug("[%s] Aborted connection at caller's request.", requestDesc);
-                    channel.close();
+                    // On the event loop, so that it is ordered against finishRequest handing the channel back.
+                    if (channel.eventLoop().inEventLoop()) {
+                      closeUnlessReturnedToPool();
+                    } else {
+                      channel.eventLoop().execute(() -> closeUnlessReturnedToPool());
+                    }
                   }
                 };
 
@@ -388,9 +396,22 @@ private void finishRequest()
             }
             removeHandlers();
             channel.config().setAutoRead(true);
+            synchronized (watermarkLock) {
+              returnedToPool = true;
+            }
             channelResourceContainer.returnResource();
           }
 
+          private void closeUnlessReturnedToPool()
+          {
+            synchronized (watermarkLock) {
+              if (!returnedToPool) {
+                log.debug("[%s] Aborted connection at caller's request.", requestDesc);
+                channel.close();
+              }
+            }
+          }
+
           @Override
           public void exceptionCaught(ChannelHandlerContext context, Throwable cause)
           {

processing/src/main/java/org/apache/druid/java/util/http/client/response/HttpResponseHandler.java

@@ -87,6 +87,10 @@ ClientResponse<IntermediateType> handleChunk(
 
   void exceptionCaught(ClientResponse<IntermediateType> clientResponse, Throwable e);
 
+  /**
+   * Flow control over the connection carrying one response. Once that response has completed, calls have no effect:
+   * the connection may already be carrying the response of another request.
+   */
   interface TrafficCop
   {
     /**

processing/src/test/java/org/apache/druid/java/util/http/client/FriendlyServersTest.java

@@ -32,6 +32,8 @@
 import org.apache.druid.java.util.http.client.response.HttpResponseHandler;
 import org.apache.druid.java.util.http.client.response.StatusResponseHandler;
 import org.apache.druid.java.util.http.client.response.StatusResponseHolder;
+import org.apache.druid.java.util.http.client.response.StringFullResponseHandler;
+import org.apache.druid.java.util.http.client.response.StringFullResponseHolder;
 import org.eclipse.jetty.server.Connector;
 import org.eclipse.jetty.server.HttpConfiguration;
 import org.eclipse.jetty.server.HttpConnectionFactory;
@@ -58,11 +60,13 @@
 import java.nio.charset.StandardCharsets;
 import java.util.Locale;
 import java.util.concurrent.BlockingQueue;
+import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
 import java.util.concurrent.LinkedBlockingQueue;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicReference;
 
@@ -71,6 +75,9 @@
  */
 public class FriendlyServersTest
 {
+  private static final String CHUNKED_OK = "HTTP/1.1 200 OK\r\nTransfer-Encoding: chunked\r\n\r\n";
+  private static final String LAST_CHUNK = "0\r\n\r\n";
+
   @Test
   public void testFriendlyHttpServer() throws Exception
   {
@@ -753,4 +760,198 @@
       lifecycle.stop();
     }
   }
+
+  /**
+   * An abort that arrives once the response is complete must not close the channel: it is back in the pool, and the
+   * next request may already be on it.
+   */
+  @Test
+  public void testAbortAfterCompletionLeavesThePooledChannelOpen() throws Exception
+  {
+    final ExecutorService exec = Executors.newSingleThreadExecutor();
+    final ServerSocket serverSocket = new ServerSocket(0);
+    final BlockingQueue<String> wire = new LinkedBlockingQueue<>();
+    serveChunkedFromWire(exec, serverSocket, wire);
+
+    final Lifecycle lifecycle = new Lifecycle();
+    try {
+      final HttpClient client = HttpClientInit.createClient(
+          HttpClientConfig.builder().withNumConnections(1).build(),
+          lifecycle
+      );
+      final Request request = new Request(
+          HttpMethod.GET,
+          new URL(StringUtils.format("http://localhost:%d/", serverSocket.getLocalPort()))
+      );
+
+      final CapturingTrafficCop<StringFullResponseHolder, StringFullResponseHolder> first =
+          new CapturingTrafficCop<>(new StringFullResponseHandler(StandardCharsets.UTF_8));
+      final ListenableFuture<StringFullResponseHolder> firstResponse = client.go(request, first);
+      wire.put(CHUNKED_OK + chunk("a") + LAST_CHUNK);
+      Assertions.assertEquals("a", firstResponse.get(30, TimeUnit.SECONDS).getContent());
+
+      first.trafficCop.get().abort();
+
+      // The server answers on its only connection, so the second request succeeds only if it reuses the first's.
+      final ListenableFuture<StringFullResponseHolder> secondResponse =
+          client.go(request, new StringFullResponseHandler(StandardCharsets.UTF_8));
+      wire.put(CHUNKED_OK + chunk("b") + LAST_CHUNK);
+      Assertions.assertEquals("b", secondResponse.get(30, TimeUnit.SECONDS).getContent());
+    }
+    finally {
+      exec.shutdownNow();
+      serverSocket.close();
+      lifecycle.stop();
+    }
+  }
+
+  /**
+   * A resume that arrives once the response is complete must not touch the channel: it is back in the pool, and the
+   * next request on it may have suspended reads.
+   */
+  @Test
+  public void testResumeAfterCompletionLeavesTheNextRequestSuspended() throws Exception
+  {
+    final ExecutorService exec = Executors.newSingleThreadExecutor();
+    final ServerSocket serverSocket = new ServerSocket(0);
+    final BlockingQueue<String> wire = new LinkedBlockingQueue<>();
+    serveChunkedFromWire(exec, serverSocket, wire);
+
+    final Lifecycle lifecycle = new Lifecycle();
+    try {
+      final HttpClient client = HttpClientInit.createClient(
+          HttpClientConfig.builder().withNumConnections(1).build(),
+          lifecycle
+      );
+      final Request request = new Request(
+          HttpMethod.GET,
+          new URL(StringUtils.format("http://localhost:%d/", serverSocket.getLocalPort()))
+      );
+
+      // The first request suspends too: only then does its late resume pass its own watermark and reach the channel.
+      final SuspendAfterFirstChunk<StringFullResponseHolder, StringFullResponseHolder> first =
+          new SuspendAfterFirstChunk<>(new StringFullResponseHandler(StandardCharsets.UTF_8));
+      final ListenableFuture<StringFullResponseHolder> firstResponse = client.go(request, first);
+      wire.put(CHUNKED_OK + chunk("a") + LAST_CHUNK);
+      Assertions.assertEquals("a", firstResponse.get(30, TimeUnit.SECONDS).getContent());
+
+      final SuspendAfterFirstChunk<StringFullResponseHolder, StringFullResponseHolder> second =
+          new SuspendAfterFirstChunk<>(new StringFullResponseHandler(StandardCharsets.UTF_8));
+      final ListenableFuture<StringFullResponseHolder> secondResponse = client.go(request, second);
+      wire.put(CHUNKED_OK + chunk("b"));
+      Assertions.assertTrue(second.suspended.await(30, TimeUnit.SECONDS));
+
+      first.trafficCop.get().resume(1);
+      wire.put(chunk("c") + LAST_CHUNK);
+      Assertions.assertThrows(TimeoutException.class, () -> secondResponse.get(1, TimeUnit.SECONDS));
+
+      second.trafficCop.get().resume(1);
+      Assertions.assertEquals("bc", secondResponse.get(30, TimeUnit.SECONDS).getContent());
+    }
+    finally {
+      exec.shutdownNow();
+      serverSocket.close();
+      lifecycle.stop();
+    }
+  }
+
+  private static String chunk(String data)
+  {
+    return StringUtils.format("%x\r\n%s\r\n", data.length(), data);
+  }
+
+  /**
+   * Serves one keep-alive connection. Each request is answered with the segments taken from {@code wire}, up to the
+   * one that ends the chunked body, so the test decides when each part of a response goes out.
+   */
+  private static void serveChunkedFromWire(ExecutorService exec, ServerSocket serverSocket, BlockingQueue<String> wire)
+  {
+    exec.submit(
+        () -> {
+          try (
+              Socket s = serverSocket.accept();
+              BufferedReader in = new BufferedReader(new InputStreamReader(s.getInputStream(), StandardCharsets.UTF_8));
+              OutputStream out = s.getOutputStream()
+          ) {
+            String line;
+            while ((line = in.readLine()) != null) {
+              if (line.isEmpty()) {
+                String segment;
+                do {
+                  segment = wire.take();
+                  out.write(StringUtils.toUtf8(segment));
+                  out.flush();
+                } while (!segment.endsWith(LAST_CHUNK));
+              }
+            }
+          }
+          catch (Exception ignored) {
+            // suppress
+          }
+        }
+    );
+  }
+
+  /**
+   * Exposes the {@link HttpResponseHandler.TrafficCop} of the response it handles.
+   */
+  private static class CapturingTrafficCop<I, F> implements HttpResponseHandler<I, F>
+  {
+    final HttpResponseHandler<I, F> delegate;
+    final AtomicReference<TrafficCop> trafficCop = new AtomicReference<>();
+
+    CapturingTrafficCop(HttpResponseHandler<I, F> delegate)
+    {
+      this.delegate = delegate;
+    }
+
+    @Override
+    public ClientResponse<I> handleResponse(HttpResponse response, TrafficCop trafficCop)
+    {
+      this.trafficCop.set(trafficCop);
+      return delegate.handleResponse(response, trafficCop);
+    }
+
+    @Override
+    public ClientResponse<I> handleChunk(ClientResponse<I> clientResponse, HttpContent chunk, long chunkNum)
+    {
+      return delegate.handleChunk(clientResponse, chunk, chunkNum);
+    }
+
+    @Override
+    public ClientResponse<F> done(ClientResponse<I> clientResponse)
+    {
+      return delegate.done(clientResponse);
+    }
+
+    @Override
+    public void exceptionCaught(ClientResponse<I> clientResponse, Throwable e)
+    {
+      delegate.exceptionCaught(clientResponse, e);
+    }
+  }
+
+  /**
+   * Suspends reads after the first chunk; resume them through {@link #trafficCop}.
+   */
+  private static class SuspendAfterFirstChunk<I, F> extends CapturingTrafficCop<I, F>
+  {
+    final CountDownLatch suspended = new CountDownLatch(1);
+
+    SuspendAfterFirstChunk(HttpResponseHandler<I, F> delegate)
+    {
+      super(delegate);
+    }
+
+    @Override
+    public ClientResponse<I> handleChunk(ClientResponse<I> clientResponse, HttpContent chunk, long chunkNum)
+    {
+      final ClientResponse<I> next = super.handleChunk(clientResponse, chunk, chunkNum);
+      if (chunkNum != 1) {
+        return next;
+      }
+      suspended.countDown();
+      return new ClientResponse<>(next.isFinished(), false, next.getObj());
+    }
+  }
 }

server/src/main/java/org/apache/druid/client/DirectDruidClient.java

@@ -24,7 +24,6 @@
 import com.fasterxml.jackson.dataformat.smile.SmileConstants;
 import com.fasterxml.jackson.dataformat.smile.SmileFactory;
 import com.fasterxml.jackson.jaxrs.smile.SmileMediaTypes;
-import com.google.common.base.Preconditions;
 import com.google.common.util.concurrent.FutureCallback;
 import com.google.common.util.concurrent.Futures;
 import com.google.common.util.concurrent.ListenableFuture;
@@ -222,8 +221,6 @@ private Failure(String message, Throwable cause)
           }
         }
 
-        private final AtomicReference<TrafficCop> trafficCopRef = new AtomicReference<>();
-
         private QueryMetrics<? super Query<T>> queryMetrics;
         private long responseStartTimeNs;
 
@@ -257,7 +254,7 @@ private boolean enqueue(ByteBuf buffer, long chunkNum) throws InterruptedExcepti
           return !usingBackpressure || currentQueuedByteCount < maxQueuedBytes;
         }
 
-        private InputStream dequeue() throws InterruptedException
+        private InputStream dequeue(TrafficCop trafficCop) throws InterruptedException
         {
           final InputStreamHolder holder = queue.poll(checkQueryTimeout(), TimeUnit.MILLISECONDS);
           if (holder == null) {
@@ -266,8 +263,7 @@ private InputStream dequeue() throws InterruptedException
 
           final long currentQueuedByteCount = queuedByteCount.addAndGet(-holder.getLength());
           if (usingBackpressure && currentQueuedByteCount < maxQueuedBytes) {
-            long backPressureTime = Preconditions.checkNotNull(trafficCopRef.get(), "No TrafficCop, how can this be?")
-                                                 .resume(holder.getChunkNum());
+            long backPressureTime = trafficCop.resume(holder.getChunkNum());
             channelSuspendedTime.addAndGet(backPressureTime);
           }
 
@@ -424,7 +420,6 @@ private void throwForNonJsonBody(
         @Override
         public ClientResponse<InputStream> handleResponse(HttpResponse response, TrafficCop trafficCop)
         {
-          trafficCopRef.set(trafficCop);
           checkQueryTimeout();
           // Netty 4: the initial HttpResponse carries no body, so the status and Content-Type are recorded
           // here and the body itself is inspected on the first HttpContent chunk. The goal is to detect a
@@ -511,7 +506,7 @@ public InputStream nextElement()
                       }
 
                       try {
-                        return dequeue();
+                        return dequeue(trafficCop);
                       }
                       catch (InterruptedException e) {
                         Thread.currentThread().interrupt();
@@ -529,20 +524,8 @@ public InputStream nextElement()
                 @Override
                 public void close()
                 {
-                  final TrafficCop trafficCop;
-                  synchronized (done) {
-                    if (done.get()) {
-                      return;
-                    }
-                    // Stop buffering further chunks (see enqueue()) and drop anything already buffered so the
-                    // underlying Netty ChannelBuffers can be released.
-                    discard.set(true);
-                    queue.clear();
-                    trafficCop = trafficCopRef.get();
-                  }
-                  if (trafficCop == null) {
-                    return;
-                  }
+                  discard.set(true);
+                  queue.clear();
                   trafficCop.abort();
                 }
               },