From b5b177e7be754e2b1c578a2edd5614a9bd20bf46 Mon Sep 17 00:00:00 2001 From: Anil Dasari Date: Fri, 20 Feb 2026 12:51:30 -0800 Subject: [PATCH 1/3] Skip lock when channel is active --- .../pinot/core/transport/ServerChannels.java | 14 ++++++++++++-- 1 file changed, 12 insertions(+), 2 deletions(-) diff --git a/pinot-core/src/main/java/org/apache/pinot/core/transport/ServerChannels.java b/pinot-core/src/main/java/org/apache/pinot/core/transport/ServerChannels.java index fae61d0567..28aef6e3b9 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/transport/ServerChannels.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/transport/ServerChannels.java @@ -155,7 +155,7 @@ class ServerChannel { final Bootstrap _bootstrap; // lock to protect channel as requests must be written into channel sequentially final ReentrantLock _channelLock = new ReentrantLock(); - Channel _channel; + volatile Channel _channel; ServerChannel(ServerRoutingInstance serverRoutingInstance) { _serverRoutingInstance = serverRoutingInstance; @@ -216,10 +216,20 @@ void setSilentShutdown() { void sendRequest(String rawTableName, AsyncQueryResponse asyncQueryResponse, ServerRoutingInstance serverRoutingInstance, byte[] requestBytes, long timeoutMs) throws InterruptedException, TimeoutException { + ensureConnected(timeoutMs); + sendRequestWithoutLocking(rawTableName, asyncQueryResponse, serverRoutingInstance, requestBytes); + } + + private void ensureConnected(long timeoutMs) throws InterruptedException, TimeoutException { + // channel is active, no lock needed + if (_channel != null && _channel.isActive()) { + return; + } + // Lock only for reconnect, not for every send + // reconnect needed — serialize with lock if (_channelLock.tryLock(timeoutMs, TimeUnit.MILLISECONDS)) { try { connectWithoutLocking(); - sendRequestWithoutLocking(rawTableName, asyncQueryResponse, serverRoutingInstance, requestBytes); } finally { _channelLock.unlock(); } From ed4f956af42096fcd3ff401a95ad5721a14fa726 Mon Sep 17 00:00:00 2001 From: Anil Dasari Date: Mon, 2 Mar 2026 09:27:30 -0800 Subject: [PATCH 2/3] Use TOCTOU channel --- .../pinot/core/transport/ServerChannels.java | 19 ++++++++++--------- 1 file changed, 10 insertions(+), 9 deletions(-) diff --git a/pinot-core/src/main/java/org/apache/pinot/core/transport/ServerChannels.java b/pinot-core/src/main/java/org/apache/pinot/core/transport/ServerChannels.java index 28aef6e3b9..b57888f547 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/transport/ServerChannels.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/transport/ServerChannels.java @@ -216,20 +216,21 @@ void setSilentShutdown() { void sendRequest(String rawTableName, AsyncQueryResponse asyncQueryResponse, ServerRoutingInstance serverRoutingInstance, byte[] requestBytes, long timeoutMs) throws InterruptedException, TimeoutException { - ensureConnected(timeoutMs); - sendRequestWithoutLocking(rawTableName, asyncQueryResponse, serverRoutingInstance, requestBytes); + Channel channel = ensureConnected(timeoutMs); + sendRequestWithoutLocking(channel, rawTableName, asyncQueryResponse, serverRoutingInstance, requestBytes); } - private void ensureConnected(long timeoutMs) throws InterruptedException, TimeoutException { - // channel is active, no lock needed - if (_channel != null && _channel.isActive()) { - return; + // 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 - // reconnect needed — serialize with lock if (_channelLock.tryLock(timeoutMs, TimeUnit.MILLISECONDS)) { try { connectWithoutLocking(); + return _channel; } finally { _channelLock.unlock(); } @@ -248,10 +249,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); From f7a2f64c57918ef51286fc246c9fe927e083b102 Mon Sep 17 00:00:00 2001 From: Anil Dasari Date: Thu, 12 Mar 2026 15:56:56 -0700 Subject: [PATCH 3/3] Added benchmark & unit test --- .../pinot/core/transport/ServerChannels.java | 3 +- .../core/transport/ServerChannelsTest.java | 176 +++++++++++++++ pinot-perf/pom.xml | 4 + .../perf/BenchmarkServerChannelLocking.java | 203 ++++++++++++++++++ 4 files changed, 385 insertions(+), 1 deletion(-) create mode 100644 pinot-perf/src/main/java/org/apache/pinot/perf/BenchmarkServerChannelLocking.java diff --git a/pinot-core/src/main/java/org/apache/pinot/core/transport/ServerChannels.java b/pinot-core/src/main/java/org/apache/pinot/core/transport/ServerChannels.java index b57888f547..1e5f3a50cd 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/transport/ServerChannels.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/transport/ServerChannels.java @@ -153,7 +153,8 @@ 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(); volatile Channel _channel; diff --git a/pinot-core/src/test/java/org/apache/pinot/core/transport/ServerChannelsTest.java b/pinot-core/src/test/java/org/apache/pinot/core/transport/ServerChannelsTest.java index 23e45c1273..dc89e7ee40 100644 --- a/pinot-core/src/test/java/org/apache/pinot/core/transport/ServerChannelsTest.java +++ b/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) () -> { + // 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 map = + (ConcurrentHashMap) 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; + } } diff --git a/pinot-perf/pom.xml b/pinot-perf/pom.xml index 61c6bec92e..6f6cce993c 100644 --- a/pinot-perf/pom.xml +++ b/pinot-perf/pom.xml @@ -198,6 +198,10 @@ org.apache.pinot.perf.StringDictionaryPerfTest pinot-StringDictionaryPerfTest + + org.apache.pinot.perf.BenchmarkServerChannelLocking + pinot-BenchmarkServerChannelLocking + flat lib diff --git a/pinot-perf/src/main/java/org/apache/pinot/perf/BenchmarkServerChannelLocking.java b/pinot-perf/src/main/java/org/apache/pinot/perf/BenchmarkServerChannelLocking.java new file mode 100644 index 0000000000..9ba9d66191 --- /dev/null +++ b/pinot-perf/src/main/java/org/apache/pinot/perf/BenchmarkServerChannelLocking.java @@ -0,0 +1,203 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.pinot.perf; + +import java.util.concurrent.TimeUnit; +import java.util.concurrent.locks.ReentrantLock; +import org.openjdk.jmh.annotations.Benchmark; +import org.openjdk.jmh.annotations.BenchmarkMode; +import org.openjdk.jmh.annotations.Fork; +import org.openjdk.jmh.annotations.Group; +import org.openjdk.jmh.annotations.GroupThreads; +import org.openjdk.jmh.annotations.Measurement; +import org.openjdk.jmh.annotations.Mode; +import org.openjdk.jmh.annotations.OutputTimeUnit; +import org.openjdk.jmh.annotations.Scope; +import org.openjdk.jmh.annotations.State; +import org.openjdk.jmh.annotations.Warmup; +import org.openjdk.jmh.infra.Blackhole; +import org.openjdk.jmh.runner.Runner; +import org.openjdk.jmh.runner.RunnerException; +import org.openjdk.jmh.runner.options.Options; +import org.openjdk.jmh.runner.options.OptionsBuilder; + + +/** + * Benchmarks the ServerChannels lock-skip optimization. + * + * The optimization avoids acquiring the ReentrantLock on the hot path when the channel is already + * active, using a volatile read + isActive() check instead. The lock is only acquired for reconnect. + * + * This benchmark compares: + * - alwaysLock: Original behavior — acquires the lock for every sendRequest call + * - volatileCheckThenLock: New behavior — volatile read fast-path, lock only when channel is inactive + */ +@BenchmarkMode(Mode.Throughput) +@OutputTimeUnit(TimeUnit.SECONDS) +@Warmup(iterations = 2, time = 5, timeUnit = TimeUnit.SECONDS) +@Measurement(iterations = 3, time = 10, timeUnit = TimeUnit.SECONDS) +@State(Scope.Benchmark) +@Fork(3) +public class BenchmarkServerChannelLocking { + + /** + * Simulates a Netty Channel with configurable active state. + * Avoids the need for a real Netty bootstrap/connection in the benchmark. + */ + static class FakeChannel { + private volatile boolean _active; + + FakeChannel(boolean active) { + _active = active; + } + + boolean isActive() { + return _active; + } + + void setActive(boolean active) { + _active = active; + } + + void writeAndFlush(byte[] data) { + // simulate minimal work + Blackhole.consumeCPU(5); + } + } + + // --- State: original approach (always lock) --- + @State(Scope.Group) + public static class AlwaysLockState { + final ReentrantLock _channelLock = new ReentrantLock(true); + FakeChannel _channel = new FakeChannel(true); + + void sendRequest(byte[] requestBytes) { + _channelLock.lock(); + try { + if (_channel == null || !_channel.isActive()) { + // simulate reconnect + _channel = new FakeChannel(true); + Blackhole.consumeCPU(50); + } + _channel.writeAndFlush(requestBytes); + } finally { + _channelLock.unlock(); + } + } + } + + // --- State: optimized approach (volatile check, lock only for reconnect) --- + @State(Scope.Group) + public static class VolatileCheckState { + final ReentrantLock _channelLock = new ReentrantLock(); + volatile FakeChannel _channel = new FakeChannel(true); + + void sendRequest(byte[] requestBytes) { + FakeChannel channel = _channel; + if (channel != null && channel.isActive()) { + // fast path: no lock needed + channel.writeAndFlush(requestBytes); + return; + } + // slow path: acquire lock for reconnect + _channelLock.lock(); + try { + if (_channel == null || !_channel.isActive()) { + _channel = new FakeChannel(true); + Blackhole.consumeCPU(50); + } + _channel.writeAndFlush(requestBytes); + } finally { + _channelLock.unlock(); + } + } + } + + private static final byte[] REQUEST_BYTES = new byte[64]; + + // --- Benchmarks: active channel (hot path, no reconnect needed) --- + + @Benchmark + @Group("alwaysLock") + @GroupThreads(8) + public void alwaysLock(AlwaysLockState state) { + state.sendRequest(REQUEST_BYTES); + } + + @Benchmark + @Group("volatileCheck") + @GroupThreads(8) + public void volatileCheck(VolatileCheckState state) { + state.sendRequest(REQUEST_BYTES); + } + + // --- Benchmarks: mixed scenario (occasional reconnect under contention) --- + + @Benchmark + @Group("alwaysLockMixed") + @GroupThreads(7) + public void alwaysLockMixedSender(AlwaysLockState state) { + state.sendRequest(REQUEST_BYTES); + } + + @Benchmark + @Group("alwaysLockMixed") + @GroupThreads(1) + public void alwaysLockMixedDisconnector(AlwaysLockState state) { + // Simulate occasional channel disconnect/reconnect + state._channelLock.lock(); + try { + state._channel.setActive(false); + Blackhole.consumeCPU(20); + state._channel = new FakeChannel(true); + } finally { + state._channelLock.unlock(); + } + } + + @Benchmark + @Group("volatileCheckMixed") + @GroupThreads(7) + public void volatileCheckMixedSender(VolatileCheckState state) { + state.sendRequest(REQUEST_BYTES); + } + + @Benchmark + @Group("volatileCheckMixed") + @GroupThreads(1) + public void volatileCheckMixedDisconnector(VolatileCheckState state) { + // Simulate occasional channel disconnect/reconnect + state._channelLock.lock(); + try { + state._channel.setActive(false); + Blackhole.consumeCPU(20); + state._channel = new FakeChannel(true); + } finally { + state._channelLock.unlock(); + } + } + + public static void main(String[] args) + throws RunnerException { + Options opt = new OptionsBuilder() + .include(BenchmarkServerChannelLocking.class.getSimpleName()) + .build(); + new Runner(opt).run(); + } +}