[Do not Merge] Skip channel lock when channel is active

pinot-core/src/main/java/org/apache/pinot/core/transport/ServerChannels.java

@@ -153,9 +153,10 @@ public void shutDown() {
   class ServerChannel {
     final ServerRoutingInstance _serverRoutingInstance;
     final Bootstrap _bootstrap;
-    // lock to protect channel as requests must be written into channel sequentially
+    // Lock to serialize reconnect: prevents two threads from simultaneously seeing _channel == null
+    // and each initiating a TCP reconnect. Normal sends bypass this lock entirely.
     final ReentrantLock _channelLock = new ReentrantLock();
-    Channel _channel;
+    volatile Channel _channel;
 
     ServerChannel(ServerRoutingInstance serverRoutingInstance) {
       _serverRoutingInstance = serverRoutingInstance;
@@ -216,10 +217,21 @@ void setSilentShutdown() {
     void sendRequest(String rawTableName, AsyncQueryResponse asyncQueryResponse,
         ServerRoutingInstance serverRoutingInstance, byte[] requestBytes, long timeoutMs)
         throws InterruptedException, TimeoutException {
+      Channel channel = ensureConnected(timeoutMs);
+      sendRequestWithoutLocking(channel, rawTableName, asyncQueryResponse, serverRoutingInstance, requestBytes);
+    }
+
+    // Returns the channel snapshot that is guaranteed to have been active at check time.
+    private Channel ensureConnected(long timeoutMs) throws InterruptedException, TimeoutException {
+      Channel channel = _channel;
+      if (channel != null && channel.isActive()) {
+        return channel;
+      }
+      // Lock only for reconnect, not for every send
       if (_channelLock.tryLock(timeoutMs, TimeUnit.MILLISECONDS)) {
         try {
           connectWithoutLocking();
-          sendRequestWithoutLocking(rawTableName, asyncQueryResponse, serverRoutingInstance, requestBytes);
+          return _channel;
         } finally {
           _channelLock.unlock();
         }
@@ -238,10 +250,10 @@ void connectWithoutLocking()
       }
     }
 
-    void sendRequestWithoutLocking(String rawTableName, AsyncQueryResponse asyncQueryResponse,
+    void sendRequestWithoutLocking(Channel channel, String rawTableName, AsyncQueryResponse asyncQueryResponse,
         ServerRoutingInstance serverRoutingInstance, byte[] requestBytes) {
       long startTimeMs = System.currentTimeMillis();
-      _channel.writeAndFlush(Unpooled.wrappedBuffer(requestBytes)).addListener(f -> {
+      channel.writeAndFlush(Unpooled.wrappedBuffer(requestBytes)).addListener(f -> {
         int requestSentLatencyMs = (int) (System.currentTimeMillis() - startTimeMs);
         _brokerMetrics.addTimedTableValue(rawTableName, BrokerTimer.NETTY_CONNECTION_SEND_REQUEST_LATENCY,
             requestSentLatencyMs, TimeUnit.MILLISECONDS);

pinot-core/src/test/java/org/apache/pinot/core/transport/ServerChannelsTest.java

@@ -19,7 +19,15 @@
 package org.apache.pinot.core.transport;
 
 import com.sun.net.httpserver.HttpServer;
+import io.netty.channel.Channel;
+import java.lang.reflect.Field;
 import java.net.InetSocketAddress;
+import java.util.concurrent.Callable;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.Executors;
+import java.util.concurrent.Future;
+import java.util.concurrent.TimeoutException;
 import org.apache.pinot.common.config.NettyConfig;
 import org.apache.pinot.common.request.BrokerRequest;
 import org.apache.pinot.common.request.InstanceRequest;
@@ -31,6 +39,9 @@
 import org.testng.annotations.Test;
 
 import static org.mockito.Mockito.mock;
+import static org.testng.Assert.assertFalse;
+import static org.testng.Assert.assertNotSame;
+import static org.testng.Assert.assertTrue;
 
 
 public class ServerChannelsTest {
@@ -82,4 +93,169 @@ public void testConnect(boolean nativeTransportEnabled)
     serverChannels.sendRequest("dummy_table_name", asyncQueryResponse, serverRoutingInstance, instanceRequest, 1000);
     serverChannels.shutDown();
   }
+
+  /**
+   * When the channel is already active, ensureConnected() must return on the fast path without
+   * acquiring _channelLock. Verified by holding the lock externally: if the fast path works the
+   * send completes successfully; if it mistakenly tries to lock it would block/timeout.
+   */
+  @Test(dataProvider = "parameters")
+  public void testActiveChannelSkipsLock(boolean nativeTransportEnabled)
+      throws Exception {
+    NettyConfig nettyConfig = new NettyConfig();
+    nettyConfig.setNativeTransportsEnabled(nativeTransportEnabled);
+    ServerChannels serverChannels =
+        new ServerChannels(mock(QueryRouter.class), nettyConfig, null, ThreadAccountantUtils.getNoOpAccountant());
+    ServerRoutingInstance serverRoutingInstance =
+        new ServerRoutingInstance("localhost", _dummyServer.getAddress().getPort(), TableType.REALTIME);
+    try {
+      serverChannels.connect(serverRoutingInstance);
+      ServerChannels.ServerChannel serverChannel = getServerChannel(serverChannels, serverRoutingInstance);
+      assertTrue(serverChannel._channel.isActive());
+
+      // Hold the lock — fast path must not attempt to acquire it
+      serverChannel._channelLock.lock();
+      try {
+        // Short timeout: would expire if slow path were taken
+        serverChannels.sendRequest("table", mock(AsyncQueryResponse.class), serverRoutingInstance,
+            buildInstanceRequest(), 100);
+      } finally {
+        serverChannel._channelLock.unlock();
+      }
+    } finally {
+      serverChannels.shutDown();
+    }
+  }
+
+  /**
+   * When the current channel is inactive (closed), sendRequest must reconnect before writing.
+   * After the call the channel field must hold a fresh active channel distinct from the closed one.
+   */
+  @Test(dataProvider = "parameters")
+  public void testInactiveChannelTriggersReconnect(boolean nativeTransportEnabled)
+      throws Exception {
+    NettyConfig nettyConfig = new NettyConfig();
+    nettyConfig.setNativeTransportsEnabled(nativeTransportEnabled);
+    ServerChannels serverChannels =
+        new ServerChannels(mock(QueryRouter.class), nettyConfig, null, ThreadAccountantUtils.getNoOpAccountant());
+    ServerRoutingInstance serverRoutingInstance =
+        new ServerRoutingInstance("localhost", _dummyServer.getAddress().getPort(), TableType.REALTIME);
+    try {
+      serverChannels.connect(serverRoutingInstance);
+      ServerChannels.ServerChannel serverChannel = getServerChannel(serverChannels, serverRoutingInstance);
+      Channel originalChannel = serverChannel._channel;
+      assertTrue(originalChannel.isActive());
+
+      // Force the channel inactive
+      originalChannel.close().sync();
+      assertFalse(serverChannel._channel.isActive());
+
+      // sendRequest must reconnect transparently
+      serverChannels.sendRequest("table", mock(AsyncQueryResponse.class), serverRoutingInstance,
+          buildInstanceRequest(), 5000);
+
+      assertTrue(serverChannel._channel.isActive());
+      assertNotSame(serverChannel._channel, originalChannel);
+    } finally {
+      serverChannels.shutDown();
+    }
+  }
+
+  /**
+   * When the channel is inactive and another thread holds _channelLock, sendRequest must throw
+   * TimeoutException once the timeout elapses. The lock is held on the test thread while a
+   * separate thread calls sendRequest, avoiding ReentrantLock re-entrancy.
+   */
+  @Test(dataProvider = "parameters")
+  public void testLockTimeoutWhenChannelInactive(boolean nativeTransportEnabled)
+      throws Exception {
+    NettyConfig nettyConfig = new NettyConfig();
+    nettyConfig.setNativeTransportsEnabled(nativeTransportEnabled);
+    ServerChannels serverChannels =
+        new ServerChannels(mock(QueryRouter.class), nettyConfig, null, ThreadAccountantUtils.getNoOpAccountant());
+    ServerRoutingInstance serverRoutingInstance =
+        new ServerRoutingInstance("localhost", _dummyServer.getAddress().getPort(), TableType.REALTIME);
+    try {
+      serverChannels.connect(serverRoutingInstance);
+      ServerChannels.ServerChannel serverChannel = getServerChannel(serverChannels, serverRoutingInstance);
+
+      // Force inactive so the slow path (lock) is taken
+      serverChannel._channel.close().sync();
+      assertFalse(serverChannel._channel.isActive());
+
+      // acquire lock using main thread.
+      serverChannel._channelLock.lock();
+      try {
+        Future<?> future = Executors.newSingleThreadExecutor().submit((Callable<Void>) () -> {
+          // send request using non main thread, so that it cannot acquire the lock and will timeout instead
+          serverChannels.sendRequest("table", mock(AsyncQueryResponse.class), serverRoutingInstance,
+              buildInstanceRequest(), 50);
+          return null;
+        });
+        try {
+          future.get();
+          throw new AssertionError("Expected TimeoutException was not thrown");
+        } catch (ExecutionException e) {
+          assertTrue(e.getCause() instanceof TimeoutException);
+          assertTrue(e.getCause().getMessage().contains(ServerChannels.CHANNEL_LOCK_TIMEOUT_MSG));
+        }
+      } finally {
+        serverChannel._channelLock.unlock();
+      }
+    } finally {
+      serverChannels.shutDown();
+    }
+  }
+
+  /**
+   * Verifies the TOCTOU fix: sendRequestWithoutLocking uses the channel snapshot returned by
+   * ensureConnected, not the _channel field. Even if _channel is nulled out after ensureConnected
+   * returns, the write targets the originally-active channel.
+   */
+  @Test(dataProvider = "parameters")
+  public void testSendRequestWithoutLockingUsesChannelSnapshot(boolean nativeTransportEnabled)
+      throws Exception {
+    NettyConfig nettyConfig = new NettyConfig();
+    nettyConfig.setNativeTransportsEnabled(nativeTransportEnabled);
+    ServerChannels serverChannels =
+        new ServerChannels(mock(QueryRouter.class), nettyConfig, null, ThreadAccountantUtils.getNoOpAccountant());
+    ServerRoutingInstance serverRoutingInstance =
+        new ServerRoutingInstance("localhost", _dummyServer.getAddress().getPort(), TableType.REALTIME);
+    try {
+      serverChannels.connect(serverRoutingInstance);
+      ServerChannels.ServerChannel serverChannel = getServerChannel(serverChannels, serverRoutingInstance);
+      Channel snapshot = serverChannel._channel;
+      assertTrue(snapshot.isActive());
+
+      // Null out _channel to simulate a race where the field changes after ensureConnected returned
+      serverChannel._channel = null;
+
+      // Must complete without NPE because sendRequestWithoutLocking uses the snapshot, not _channel
+      serverChannel.sendRequestWithoutLocking(snapshot, "table", mock(AsyncQueryResponse.class),
+          serverRoutingInstance, new byte[64]);
+
+      // Restore so shutDown is clean
+      serverChannel._channel = snapshot;
+    } finally {
+      serverChannels.shutDown();
+    }
+  }
+
+  @SuppressWarnings("unchecked")
+  private static ServerChannels.ServerChannel getServerChannel(ServerChannels serverChannels,
+      ServerRoutingInstance instance)
+      throws Exception {
+    Field mapField = ServerChannels.class.getDeclaredField("_serverToChannelMap");
+    mapField.setAccessible(true);
+    ConcurrentHashMap<ServerRoutingInstance, ServerChannels.ServerChannel> map =
+        (ConcurrentHashMap<ServerRoutingInstance, ServerChannels.ServerChannel>) mapField.get(serverChannels);
+    return map.get(instance);
+  }
+
+  private static InstanceRequest buildInstanceRequest() {
+    InstanceRequest instanceRequest = new InstanceRequest();
+    instanceRequest.setRequestId(System.currentTimeMillis());
+    instanceRequest.setQuery(new BrokerRequest());
+    return instanceRequest;
+  }
 }

pinot-perf/pom.xml

@@ -198,6 +198,10 @@
               <mainClass>org.apache.pinot.perf.StringDictionaryPerfTest</mainClass>
               <name>pinot-StringDictionaryPerfTest</name>
             </program>
+            <program>
+              <mainClass>org.apache.pinot.perf.BenchmarkServerChannelLocking</mainClass>
+              <name>pinot-BenchmarkServerChannelLocking</name>
+            </program>
           </programs>
           <repositoryLayout>flat</repositoryLayout>
           <repositoryName>lib</repositoryName>