From d34207814367173ff8aa5837fcfeae7906972f50 Mon Sep 17 00:00:00 2001 From: Siddharth Teotia Date: Fri, 10 Apr 2026 22:48:38 -0700 Subject: [PATCH 1/5] Fix gRPC client resource leak when servers are decommissioned MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit PinotServerStreamingQueryClient cached gRPC clients in a ConcurrentHashMap that grew monotonically — clients were never closed or removed when servers left the cluster. This caused stale ManagedChannel connections to accumulate over the lifetime of the broker. The fix adds a cleanupStaleClients() method that compares cached clients against the current enabled server instance map and closes/removes any that are no longer needed. GrpcBrokerRequestHandler now implements ClusterChangeHandler and is registered for INSTANCE_CONFIG changes, so cleanup runs automatically whenever servers are added or removed. Co-Authored-By: Claude Opus 4.6 (1M context) --- .../broker/helix/BaseBrokerStarter.java | 5 + .../GrpcBrokerRequestHandler.java | 46 +++++- .../GrpcBrokerRequestHandlerTest.java | 144 ++++++++++++++++++ 3 files changed, 194 insertions(+), 1 deletion(-) create mode 100644 pinot-broker/src/test/java/org/apache/pinot/broker/requesthandler/GrpcBrokerRequestHandlerTest.java diff --git a/pinot-broker/src/main/java/org/apache/pinot/broker/broker/helix/BaseBrokerStarter.java b/pinot-broker/src/main/java/org/apache/pinot/broker/broker/helix/BaseBrokerStarter.java index d712ddd030..b2e3f3b373 100644 --- a/pinot-broker/src/main/java/org/apache/pinot/broker/broker/helix/BaseBrokerStarter.java +++ b/pinot-broker/src/main/java/org/apache/pinot/broker/broker/helix/BaseBrokerStarter.java @@ -480,6 +480,11 @@ public void start() } _instanceConfigChangeHandlers.add(_routingManager); _instanceConfigChangeHandlers.add(_queryQuotaManager); + // Register the single-stage broker request handler for instance config changes if it supports it (e.g., + // GrpcBrokerRequestHandler uses this to clean up stale gRPC clients when servers are decommissioned) + if (singleStageBrokerRequestHandler instanceof ClusterChangeHandler) { + _instanceConfigChangeHandlers.add((ClusterChangeHandler) singleStageBrokerRequestHandler); + } for (ClusterChangeHandler liveInstanceChangeHandler : _liveInstanceChangeHandlers) { liveInstanceChangeHandler.init(_spectatorHelixManager); } diff --git a/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/GrpcBrokerRequestHandler.java b/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/GrpcBrokerRequestHandler.java index e5e81be6ac..b79dc75662 100644 --- a/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/GrpcBrokerRequestHandler.java +++ b/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/GrpcBrokerRequestHandler.java @@ -18,15 +18,21 @@ */ package org.apache.pinot.broker.requesthandler; +import com.google.common.annotations.VisibleForTesting; import io.grpc.ConnectivityState; import java.util.HashMap; +import java.util.HashSet; import java.util.Iterator; import java.util.List; import java.util.Map; +import java.util.Set; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.TimeUnit; import javax.annotation.concurrent.ThreadSafe; +import org.apache.helix.HelixConstants; +import org.apache.helix.HelixManager; import org.apache.pinot.broker.broker.AccessControlFactory; +import org.apache.pinot.broker.broker.helix.ClusterChangeHandler; import org.apache.pinot.broker.queryquota.QueryQuotaManager; import org.apache.pinot.common.config.GrpcConfig; import org.apache.pinot.common.config.provider.TableCache; @@ -55,7 +61,7 @@ * The GrpcBrokerRequestHandler class communicates query request via GRPC. */ @ThreadSafe -public class GrpcBrokerRequestHandler extends BaseSingleStageBrokerRequestHandler { +public class GrpcBrokerRequestHandler extends BaseSingleStageBrokerRequestHandler implements ClusterChangeHandler { private static final Logger LOGGER = LoggerFactory.getLogger(GrpcBrokerRequestHandler.class); private final StreamingReduceService _streamingReduceService; @@ -87,6 +93,18 @@ public void shutDown() { _streamingReduceService.shutDown(); } + @Override + public void init(HelixManager helixManager) { + // No-op: routing manager is already available via constructor injection + } + + @Override + public void processClusterChange(HelixConstants.ChangeType changeType) { + if (changeType == HelixConstants.ChangeType.INSTANCE_CONFIG) { + _streamingQueryClient.cleanupStaleClients(_routingManager.getEnabledServerInstanceMap()); + } + } + @Override protected BrokerResponseNative processBrokerRequest(long requestId, BrokerRequest originalBrokerRequest, BrokerRequest serverBrokerRequest, TableRouteInfo route, @@ -170,11 +188,37 @@ private ServerGrpcQueryClient getOrCreateGrpcQueryClient(ServerInstance serverIn k -> new ServerGrpcQueryClient(serverInstance.getHostname(), serverInstance.getGrpcPort(), _config)); } + /** + * Closes and removes gRPC clients for servers that are no longer in the enabled server instance map. + * This prevents resource leaks when servers are decommissioned from the cluster. + */ + public void cleanupStaleClients(Map enabledServerInstances) { + Set activeHostnamePorts = new HashSet<>(); + for (ServerInstance serverInstance : enabledServerInstances.values()) { + activeHostnamePorts.add( + String.format("%s_%d", serverInstance.getHostname(), serverInstance.getGrpcPort())); + } + Iterator> iterator = _grpcQueryClientMap.entrySet().iterator(); + while (iterator.hasNext()) { + Map.Entry entry = iterator.next(); + if (!activeHostnamePorts.contains(entry.getKey())) { + LOGGER.info("Closing stale gRPC client for decommissioned server: {}", entry.getKey()); + entry.getValue().close(); + iterator.remove(); + } + } + } + public void shutdown() { for (ServerGrpcQueryClient client : _grpcQueryClientMap.values()) { client.close(); } } + + @VisibleForTesting + Map getGrpcQueryClientMap() { + return _grpcQueryClientMap; + } } /** diff --git a/pinot-broker/src/test/java/org/apache/pinot/broker/requesthandler/GrpcBrokerRequestHandlerTest.java b/pinot-broker/src/test/java/org/apache/pinot/broker/requesthandler/GrpcBrokerRequestHandlerTest.java new file mode 100644 index 0000000000..9e30199a69 --- /dev/null +++ b/pinot-broker/src/test/java/org/apache/pinot/broker/requesthandler/GrpcBrokerRequestHandlerTest.java @@ -0,0 +1,144 @@ +/** + * 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.broker.requesthandler; + +import java.util.Collections; +import java.util.HashMap; +import java.util.Map; +import org.apache.pinot.common.config.GrpcConfig; +import org.apache.pinot.common.utils.grpc.ServerGrpcQueryClient; +import org.apache.pinot.core.transport.ServerInstance; +import org.testng.annotations.Test; + +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertFalse; +import static org.testng.Assert.assertTrue; + + +public class GrpcBrokerRequestHandlerTest { + + @Test + public void testCleanupStaleClientsRemovesDecommissionedServers() { + GrpcBrokerRequestHandler.PinotServerStreamingQueryClient queryClient = + new GrpcBrokerRequestHandler.PinotServerStreamingQueryClient( + new GrpcConfig(Collections.emptyMap())); + + ServerGrpcQueryClient mockClient1 = mock(ServerGrpcQueryClient.class); + ServerGrpcQueryClient mockClient2 = mock(ServerGrpcQueryClient.class); + ServerGrpcQueryClient mockClient3 = mock(ServerGrpcQueryClient.class); + + Map clientMap = queryClient.getGrpcQueryClientMap(); + clientMap.put("host1_8090", mockClient1); + clientMap.put("host2_8090", mockClient2); + clientMap.put("host3_8090", mockClient3); + + // Only host1 and host3 are still enabled; host2 has been decommissioned + Map enabledServers = new HashMap<>(); + ServerInstance server1 = mock(ServerInstance.class); + when(server1.getHostname()).thenReturn("host1"); + when(server1.getGrpcPort()).thenReturn(8090); + enabledServers.put("Server_host1_8099", server1); + + ServerInstance server3 = mock(ServerInstance.class); + when(server3.getHostname()).thenReturn("host3"); + when(server3.getGrpcPort()).thenReturn(8090); + enabledServers.put("Server_host3_8099", server3); + + queryClient.cleanupStaleClients(enabledServers); + + // host2's client should be closed and removed + verify(mockClient2).close(); + // host1 and host3 clients should not be closed + verify(mockClient1, never()).close(); + verify(mockClient3, never()).close(); + + assertEquals(clientMap.size(), 2); + assertTrue(clientMap.containsKey("host1_8090")); + assertFalse(clientMap.containsKey("host2_8090")); + assertTrue(clientMap.containsKey("host3_8090")); + } + + @Test + public void testCleanupStaleClientsWithEmptyEnabledServers() { + GrpcBrokerRequestHandler.PinotServerStreamingQueryClient queryClient = + new GrpcBrokerRequestHandler.PinotServerStreamingQueryClient( + new GrpcConfig(Collections.emptyMap())); + + ServerGrpcQueryClient mockClient1 = mock(ServerGrpcQueryClient.class); + ServerGrpcQueryClient mockClient2 = mock(ServerGrpcQueryClient.class); + + Map clientMap = queryClient.getGrpcQueryClientMap(); + clientMap.put("host1_8090", mockClient1); + clientMap.put("host2_8090", mockClient2); + + // No enabled servers — all clients should be cleaned up + queryClient.cleanupStaleClients(Collections.emptyMap()); + + verify(mockClient1).close(); + verify(mockClient2).close(); + assertTrue(clientMap.isEmpty()); + } + + @Test + public void testCleanupStaleClientsWithNoStaleClients() { + GrpcBrokerRequestHandler.PinotServerStreamingQueryClient queryClient = + new GrpcBrokerRequestHandler.PinotServerStreamingQueryClient( + new GrpcConfig(Collections.emptyMap())); + + ServerGrpcQueryClient mockClient1 = mock(ServerGrpcQueryClient.class); + + Map clientMap = queryClient.getGrpcQueryClientMap(); + clientMap.put("host1_8090", mockClient1); + + // host1 is still enabled — no cleanup needed + Map enabledServers = new HashMap<>(); + ServerInstance server1 = mock(ServerInstance.class); + when(server1.getHostname()).thenReturn("host1"); + when(server1.getGrpcPort()).thenReturn(8090); + enabledServers.put("Server_host1_8099", server1); + + queryClient.cleanupStaleClients(enabledServers); + + verify(mockClient1, never()).close(); + assertEquals(clientMap.size(), 1); + assertTrue(clientMap.containsKey("host1_8090")); + } + + @Test + public void testCleanupStaleClientsWithEmptyClientMap() { + GrpcBrokerRequestHandler.PinotServerStreamingQueryClient queryClient = + new GrpcBrokerRequestHandler.PinotServerStreamingQueryClient( + new GrpcConfig(Collections.emptyMap())); + + // No clients exist — cleanup should be a no-op + Map enabledServers = new HashMap<>(); + ServerInstance server1 = mock(ServerInstance.class); + when(server1.getHostname()).thenReturn("host1"); + when(server1.getGrpcPort()).thenReturn(8090); + enabledServers.put("Server_host1_8099", server1); + + queryClient.cleanupStaleClients(enabledServers); + + assertTrue(queryClient.getGrpcQueryClientMap().isEmpty()); + } +} From 46a22d375032fd4b801010101e0cb560fb722992 Mon Sep 17 00:00:00 2001 From: Siddharth Teotia Date: Fri, 10 Apr 2026 22:57:04 -0700 Subject: [PATCH 2/5] Simplify: clean up stale gRPC clients from retryUnhealthyServer MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Remove ClusterChangeHandler implementation — request handlers don't implement that interface in this codebase. Instead, call cleanupStaleClients() opportunistically from retryUnhealthyServer(), which already runs periodically via the failure detector's retry thread. This keeps the fix self-contained within GrpcBrokerRequestHandler. Co-Authored-By: Claude Opus 4.6 (1M context) --- .../broker/helix/BaseBrokerStarter.java | 5 ----- .../GrpcBrokerRequestHandler.java | 21 +++++-------------- 2 files changed, 5 insertions(+), 21 deletions(-) diff --git a/pinot-broker/src/main/java/org/apache/pinot/broker/broker/helix/BaseBrokerStarter.java b/pinot-broker/src/main/java/org/apache/pinot/broker/broker/helix/BaseBrokerStarter.java index b2e3f3b373..d712ddd030 100644 --- a/pinot-broker/src/main/java/org/apache/pinot/broker/broker/helix/BaseBrokerStarter.java +++ b/pinot-broker/src/main/java/org/apache/pinot/broker/broker/helix/BaseBrokerStarter.java @@ -480,11 +480,6 @@ public void start() } _instanceConfigChangeHandlers.add(_routingManager); _instanceConfigChangeHandlers.add(_queryQuotaManager); - // Register the single-stage broker request handler for instance config changes if it supports it (e.g., - // GrpcBrokerRequestHandler uses this to clean up stale gRPC clients when servers are decommissioned) - if (singleStageBrokerRequestHandler instanceof ClusterChangeHandler) { - _instanceConfigChangeHandlers.add((ClusterChangeHandler) singleStageBrokerRequestHandler); - } for (ClusterChangeHandler liveInstanceChangeHandler : _liveInstanceChangeHandlers) { liveInstanceChangeHandler.init(_spectatorHelixManager); } diff --git a/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/GrpcBrokerRequestHandler.java b/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/GrpcBrokerRequestHandler.java index b79dc75662..067b474386 100644 --- a/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/GrpcBrokerRequestHandler.java +++ b/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/GrpcBrokerRequestHandler.java @@ -29,10 +29,7 @@ import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.TimeUnit; import javax.annotation.concurrent.ThreadSafe; -import org.apache.helix.HelixConstants; -import org.apache.helix.HelixManager; import org.apache.pinot.broker.broker.AccessControlFactory; -import org.apache.pinot.broker.broker.helix.ClusterChangeHandler; import org.apache.pinot.broker.queryquota.QueryQuotaManager; import org.apache.pinot.common.config.GrpcConfig; import org.apache.pinot.common.config.provider.TableCache; @@ -61,7 +58,7 @@ * The GrpcBrokerRequestHandler class communicates query request via GRPC. */ @ThreadSafe -public class GrpcBrokerRequestHandler extends BaseSingleStageBrokerRequestHandler implements ClusterChangeHandler { +public class GrpcBrokerRequestHandler extends BaseSingleStageBrokerRequestHandler { private static final Logger LOGGER = LoggerFactory.getLogger(GrpcBrokerRequestHandler.class); private final StreamingReduceService _streamingReduceService; @@ -93,18 +90,6 @@ public void shutDown() { _streamingReduceService.shutDown(); } - @Override - public void init(HelixManager helixManager) { - // No-op: routing manager is already available via constructor injection - } - - @Override - public void processClusterChange(HelixConstants.ChangeType changeType) { - if (changeType == HelixConstants.ChangeType.INSTANCE_CONFIG) { - _streamingQueryClient.cleanupStaleClients(_routingManager.getEnabledServerInstanceMap()); - } - } - @Override protected BrokerResponseNative processBrokerRequest(long requestId, BrokerRequest originalBrokerRequest, BrokerRequest serverBrokerRequest, TableRouteInfo route, @@ -223,8 +208,12 @@ Map getGrpcQueryClientMap() { /** * Check if a server that was previously detected as unhealthy is now healthy. + * Also opportunistically cleans up stale gRPC clients for servers that have been decommissioned. */ private FailureDetector.ServerState retryUnhealthyServer(String instanceId) { + // Opportunistically clean up gRPC clients for servers no longer in the cluster + _streamingQueryClient.cleanupStaleClients(_routingManager.getEnabledServerInstanceMap()); + LOGGER.info("Checking gRPC connection to unhealthy server: {}", instanceId); ServerInstance serverInstance = _routingManager.getEnabledServerInstanceMap().get(instanceId); if (serverInstance == null) { From 95fc765f19a716811533b31497565b938f6132c6 Mon Sep 17 00:00:00 2001 From: Siddharth Teotia Date: Fri, 10 Apr 2026 23:06:54 -0700 Subject: [PATCH 3/5] Add comprehensive tests and fix close() resilience in cleanup - Add try-catch around close() in cleanupStaleClients so one failing client doesn't prevent cleanup of remaining stale clients - Add @VisibleForTesting getter for streamingQueryClient - Add 13 tests covering: - cleanupStaleClients: single/multiple removals, port specificity, empty maps, idempotency, close() exception resilience - retryUnhealthyServer integration: stale client cleanup on retry, return states (UNHEALTHY/UNKNOWN/HEALTHY) for each code path Co-Authored-By: Claude Opus 4.6 (1M context) --- .../GrpcBrokerRequestHandler.java | 11 +- .../GrpcBrokerRequestHandlerTest.java | 321 +++++++++++++++--- 2 files changed, 290 insertions(+), 42 deletions(-) diff --git a/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/GrpcBrokerRequestHandler.java b/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/GrpcBrokerRequestHandler.java index 067b474386..19aba3f0ea 100644 --- a/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/GrpcBrokerRequestHandler.java +++ b/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/GrpcBrokerRequestHandler.java @@ -90,6 +90,11 @@ public void shutDown() { _streamingReduceService.shutDown(); } + @VisibleForTesting + PinotServerStreamingQueryClient getStreamingQueryClient() { + return _streamingQueryClient; + } + @Override protected BrokerResponseNative processBrokerRequest(long requestId, BrokerRequest originalBrokerRequest, BrokerRequest serverBrokerRequest, TableRouteInfo route, @@ -188,7 +193,11 @@ public void cleanupStaleClients(Map enabledServerInstanc Map.Entry entry = iterator.next(); if (!activeHostnamePorts.contains(entry.getKey())) { LOGGER.info("Closing stale gRPC client for decommissioned server: {}", entry.getKey()); - entry.getValue().close(); + try { + entry.getValue().close(); + } catch (Exception e) { + LOGGER.warn("Failed to close gRPC client for server: {}", entry.getKey(), e); + } iterator.remove(); } } diff --git a/pinot-broker/src/test/java/org/apache/pinot/broker/requesthandler/GrpcBrokerRequestHandlerTest.java b/pinot-broker/src/test/java/org/apache/pinot/broker/requesthandler/GrpcBrokerRequestHandlerTest.java index 9e30199a69..f8cdb2e895 100644 --- a/pinot-broker/src/test/java/org/apache/pinot/broker/requesthandler/GrpcBrokerRequestHandlerTest.java +++ b/pinot-broker/src/test/java/org/apache/pinot/broker/requesthandler/GrpcBrokerRequestHandlerTest.java @@ -18,14 +18,28 @@ */ package org.apache.pinot.broker.requesthandler; +import io.grpc.ConnectivityState; +import io.grpc.ManagedChannel; import java.util.Collections; import java.util.HashMap; import java.util.Map; +import java.util.function.Function; +import org.apache.pinot.broker.broker.AllowAllAccessControlFactory; +import org.apache.pinot.broker.queryquota.QueryQuotaManager; import org.apache.pinot.common.config.GrpcConfig; +import org.apache.pinot.common.failuredetector.FailureDetector; +import org.apache.pinot.common.metrics.BrokerMetrics; import org.apache.pinot.common.utils.grpc.ServerGrpcQueryClient; +import org.apache.pinot.core.routing.RoutingManager; import org.apache.pinot.core.transport.ServerInstance; +import org.apache.pinot.spi.accounting.ThreadAccountantUtils; +import org.apache.pinot.spi.env.PinotConfiguration; +import org.apache.pinot.spi.eventlistener.query.BrokerQueryEventListenerFactory; +import org.mockito.ArgumentCaptor; +import org.testng.annotations.BeforeClass; import org.testng.annotations.Test; +import static org.mockito.Mockito.doThrow; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.never; import static org.mockito.Mockito.verify; @@ -37,11 +51,19 @@ public class GrpcBrokerRequestHandlerTest { + @BeforeClass + public void setUp() { + BrokerMetrics.register(mock(BrokerMetrics.class)); + BrokerQueryEventListenerFactory.init(new PinotConfiguration()); + } + + // ======================== + // PinotServerStreamingQueryClient.cleanupStaleClients tests + // ======================== + @Test - public void testCleanupStaleClientsRemovesDecommissionedServers() { - GrpcBrokerRequestHandler.PinotServerStreamingQueryClient queryClient = - new GrpcBrokerRequestHandler.PinotServerStreamingQueryClient( - new GrpcConfig(Collections.emptyMap())); + public void testCleanupRemovesDecommissionedServer() { + GrpcBrokerRequestHandler.PinotServerStreamingQueryClient queryClient = createQueryClient(); ServerGrpcQueryClient mockClient1 = mock(ServerGrpcQueryClient.class); ServerGrpcQueryClient mockClient2 = mock(ServerGrpcQueryClient.class); @@ -52,26 +74,16 @@ public void testCleanupStaleClientsRemovesDecommissionedServers() { clientMap.put("host2_8090", mockClient2); clientMap.put("host3_8090", mockClient3); - // Only host1 and host3 are still enabled; host2 has been decommissioned + // Only host1 and host3 are still enabled Map enabledServers = new HashMap<>(); - ServerInstance server1 = mock(ServerInstance.class); - when(server1.getHostname()).thenReturn("host1"); - when(server1.getGrpcPort()).thenReturn(8090); - enabledServers.put("Server_host1_8099", server1); - - ServerInstance server3 = mock(ServerInstance.class); - when(server3.getHostname()).thenReturn("host3"); - when(server3.getGrpcPort()).thenReturn(8090); - enabledServers.put("Server_host3_8099", server3); + enabledServers.put("Server_host1_8099", mockServerInstance("host1", 8090)); + enabledServers.put("Server_host3_8099", mockServerInstance("host3", 8090)); queryClient.cleanupStaleClients(enabledServers); - // host2's client should be closed and removed verify(mockClient2).close(); - // host1 and host3 clients should not be closed verify(mockClient1, never()).close(); verify(mockClient3, never()).close(); - assertEquals(clientMap.size(), 2); assertTrue(clientMap.containsKey("host1_8090")); assertFalse(clientMap.containsKey("host2_8090")); @@ -79,10 +91,61 @@ public void testCleanupStaleClientsRemovesDecommissionedServers() { } @Test - public void testCleanupStaleClientsWithEmptyEnabledServers() { - GrpcBrokerRequestHandler.PinotServerStreamingQueryClient queryClient = - new GrpcBrokerRequestHandler.PinotServerStreamingQueryClient( - new GrpcConfig(Collections.emptyMap())); + public void testCleanupRemovesMultipleDecommissionedServers() { + GrpcBrokerRequestHandler.PinotServerStreamingQueryClient queryClient = createQueryClient(); + + ServerGrpcQueryClient mockClient1 = mock(ServerGrpcQueryClient.class); + ServerGrpcQueryClient mockClient2 = mock(ServerGrpcQueryClient.class); + ServerGrpcQueryClient mockClient3 = mock(ServerGrpcQueryClient.class); + ServerGrpcQueryClient mockClient4 = mock(ServerGrpcQueryClient.class); + + Map clientMap = queryClient.getGrpcQueryClientMap(); + clientMap.put("host1_8090", mockClient1); + clientMap.put("host2_8090", mockClient2); + clientMap.put("host3_8090", mockClient3); + clientMap.put("host4_8090", mockClient4); + + // Only host1 remains; host2, host3, host4 all decommissioned + Map enabledServers = new HashMap<>(); + enabledServers.put("Server_host1_8099", mockServerInstance("host1", 8090)); + + queryClient.cleanupStaleClients(enabledServers); + + verify(mockClient1, never()).close(); + verify(mockClient2).close(); + verify(mockClient3).close(); + verify(mockClient4).close(); + assertEquals(clientMap.size(), 1); + assertTrue(clientMap.containsKey("host1_8090")); + } + + @Test + public void testCleanupDistinguishesSameHostDifferentPorts() { + GrpcBrokerRequestHandler.PinotServerStreamingQueryClient queryClient = createQueryClient(); + + ServerGrpcQueryClient clientPort8090 = mock(ServerGrpcQueryClient.class); + ServerGrpcQueryClient clientPort9090 = mock(ServerGrpcQueryClient.class); + + Map clientMap = queryClient.getGrpcQueryClientMap(); + clientMap.put("host1_8090", clientPort8090); + clientMap.put("host1_9090", clientPort9090); + + // Only port 8090 is still enabled + Map enabledServers = new HashMap<>(); + enabledServers.put("Server_host1_8099", mockServerInstance("host1", 8090)); + + queryClient.cleanupStaleClients(enabledServers); + + verify(clientPort8090, never()).close(); + verify(clientPort9090).close(); + assertEquals(clientMap.size(), 1); + assertTrue(clientMap.containsKey("host1_8090")); + assertFalse(clientMap.containsKey("host1_9090")); + } + + @Test + public void testCleanupWithEmptyEnabledServersRemovesAll() { + GrpcBrokerRequestHandler.PinotServerStreamingQueryClient queryClient = createQueryClient(); ServerGrpcQueryClient mockClient1 = mock(ServerGrpcQueryClient.class); ServerGrpcQueryClient mockClient2 = mock(ServerGrpcQueryClient.class); @@ -91,7 +154,6 @@ public void testCleanupStaleClientsWithEmptyEnabledServers() { clientMap.put("host1_8090", mockClient1); clientMap.put("host2_8090", mockClient2); - // No enabled servers — all clients should be cleaned up queryClient.cleanupStaleClients(Collections.emptyMap()); verify(mockClient1).close(); @@ -100,45 +162,222 @@ public void testCleanupStaleClientsWithEmptyEnabledServers() { } @Test - public void testCleanupStaleClientsWithNoStaleClients() { - GrpcBrokerRequestHandler.PinotServerStreamingQueryClient queryClient = - new GrpcBrokerRequestHandler.PinotServerStreamingQueryClient( - new GrpcConfig(Collections.emptyMap())); + public void testCleanupNoOpWhenAllServersActive() { + GrpcBrokerRequestHandler.PinotServerStreamingQueryClient queryClient = createQueryClient(); ServerGrpcQueryClient mockClient1 = mock(ServerGrpcQueryClient.class); Map clientMap = queryClient.getGrpcQueryClientMap(); clientMap.put("host1_8090", mockClient1); - // host1 is still enabled — no cleanup needed Map enabledServers = new HashMap<>(); - ServerInstance server1 = mock(ServerInstance.class); - when(server1.getHostname()).thenReturn("host1"); - when(server1.getGrpcPort()).thenReturn(8090); - enabledServers.put("Server_host1_8099", server1); + enabledServers.put("Server_host1_8099", mockServerInstance("host1", 8090)); queryClient.cleanupStaleClients(enabledServers); verify(mockClient1, never()).close(); assertEquals(clientMap.size(), 1); - assertTrue(clientMap.containsKey("host1_8090")); } @Test - public void testCleanupStaleClientsWithEmptyClientMap() { - GrpcBrokerRequestHandler.PinotServerStreamingQueryClient queryClient = - new GrpcBrokerRequestHandler.PinotServerStreamingQueryClient( - new GrpcConfig(Collections.emptyMap())); + public void testCleanupNoOpWhenClientMapEmpty() { + GrpcBrokerRequestHandler.PinotServerStreamingQueryClient queryClient = createQueryClient(); - // No clients exist — cleanup should be a no-op Map enabledServers = new HashMap<>(); - ServerInstance server1 = mock(ServerInstance.class); - when(server1.getHostname()).thenReturn("host1"); - when(server1.getGrpcPort()).thenReturn(8090); - enabledServers.put("Server_host1_8099", server1); + enabledServers.put("Server_host1_8099", mockServerInstance("host1", 8090)); queryClient.cleanupStaleClients(enabledServers); assertTrue(queryClient.getGrpcQueryClientMap().isEmpty()); } + + @Test + public void testCleanupIsIdempotent() { + GrpcBrokerRequestHandler.PinotServerStreamingQueryClient queryClient = createQueryClient(); + + ServerGrpcQueryClient mockClient1 = mock(ServerGrpcQueryClient.class); + ServerGrpcQueryClient mockClient2 = mock(ServerGrpcQueryClient.class); + + Map clientMap = queryClient.getGrpcQueryClientMap(); + clientMap.put("host1_8090", mockClient1); + clientMap.put("host2_8090", mockClient2); + + Map enabledServers = new HashMap<>(); + enabledServers.put("Server_host1_8099", mockServerInstance("host1", 8090)); + + queryClient.cleanupStaleClients(enabledServers); + assertEquals(clientMap.size(), 1); + verify(mockClient2).close(); + + // Second call should be a no-op — host2 is already gone + queryClient.cleanupStaleClients(enabledServers); + assertEquals(clientMap.size(), 1); + assertTrue(clientMap.containsKey("host1_8090")); + } + + @Test + public void testCleanupContinuesWhenCloseThrows() { + GrpcBrokerRequestHandler.PinotServerStreamingQueryClient queryClient = createQueryClient(); + + ServerGrpcQueryClient failingClient = mock(ServerGrpcQueryClient.class); + doThrow(new RuntimeException("channel close failed")).when(failingClient).close(); + ServerGrpcQueryClient normalClient = mock(ServerGrpcQueryClient.class); + + Map clientMap = queryClient.getGrpcQueryClientMap(); + clientMap.put("host1_8090", failingClient); + clientMap.put("host2_8090", normalClient); + + // All servers decommissioned — both should be removed even if one close() throws + queryClient.cleanupStaleClients(Collections.emptyMap()); + + verify(failingClient).close(); + verify(normalClient).close(); + assertTrue(clientMap.isEmpty()); + } + + // ======================== + // retryUnhealthyServer integration tests + // ======================== + + @Test + public void testRetryUnhealthyServerCleansUpStaleClients() { + // Capture the retrier function registered with the failure detector + FailureDetector failureDetector = mock(FailureDetector.class); + @SuppressWarnings("unchecked") + ArgumentCaptor> retrierCaptor = + ArgumentCaptor.forClass(Function.class); + + RoutingManager routingManager = mock(RoutingManager.class); + Map enabledServers = new HashMap<>(); + enabledServers.put("Server_host1_8099", mockServerInstance("host1", 8090)); + when(routingManager.getEnabledServerInstanceMap()).thenReturn(enabledServers); + + GrpcBrokerRequestHandler handler = createHandler(routingManager, failureDetector); + verify(failureDetector).registerUnhealthyServerRetrier(retrierCaptor.capture()); + Function retrier = retrierCaptor.getValue(); + + // Inject a stale client for a decommissioned server + Map clientMap = handler.getStreamingQueryClient().getGrpcQueryClientMap(); + ServerGrpcQueryClient staleClient = mock(ServerGrpcQueryClient.class); + clientMap.put("host2_8090", staleClient); + + // Retrying any server should trigger cleanup of host2's stale client + retrier.apply("Server_host2_8099"); + + verify(staleClient).close(); + assertFalse(clientMap.containsKey("host2_8090")); + } + + @Test + public void testRetryReturnsUnhealthyWhenServerNotInRouting() { + FailureDetector failureDetector = mock(FailureDetector.class); + @SuppressWarnings("unchecked") + ArgumentCaptor> retrierCaptor = + ArgumentCaptor.forClass(Function.class); + + RoutingManager routingManager = mock(RoutingManager.class); + when(routingManager.getEnabledServerInstanceMap()).thenReturn(Collections.emptyMap()); + + createHandler(routingManager, failureDetector); + verify(failureDetector).registerUnhealthyServerRetrier(retrierCaptor.capture()); + Function retrier = retrierCaptor.getValue(); + + assertEquals(retrier.apply("Server_host1_8099"), FailureDetector.ServerState.UNHEALTHY); + } + + @Test + public void testRetryReturnsUnknownWhenNoGrpcClient() { + FailureDetector failureDetector = mock(FailureDetector.class); + @SuppressWarnings("unchecked") + ArgumentCaptor> retrierCaptor = + ArgumentCaptor.forClass(Function.class); + + RoutingManager routingManager = mock(RoutingManager.class); + Map enabledServers = new HashMap<>(); + enabledServers.put("Server_host1_8099", mockServerInstance("host1", 8090)); + when(routingManager.getEnabledServerInstanceMap()).thenReturn(enabledServers); + + createHandler(routingManager, failureDetector); + verify(failureDetector).registerUnhealthyServerRetrier(retrierCaptor.capture()); + Function retrier = retrierCaptor.getValue(); + + // Server is in routing but no gRPC client has been created for it + assertEquals(retrier.apply("Server_host1_8099"), FailureDetector.ServerState.UNKNOWN); + } + + @Test + public void testRetryReturnsHealthyWhenChannelReady() { + FailureDetector failureDetector = mock(FailureDetector.class); + @SuppressWarnings("unchecked") + ArgumentCaptor> retrierCaptor = + ArgumentCaptor.forClass(Function.class); + + RoutingManager routingManager = mock(RoutingManager.class); + Map enabledServers = new HashMap<>(); + enabledServers.put("Server_host1_8099", mockServerInstance("host1", 8090)); + when(routingManager.getEnabledServerInstanceMap()).thenReturn(enabledServers); + + GrpcBrokerRequestHandler handler = createHandler(routingManager, failureDetector); + verify(failureDetector).registerUnhealthyServerRetrier(retrierCaptor.capture()); + Function retrier = retrierCaptor.getValue(); + + // Inject a client with a READY channel + ServerGrpcQueryClient mockClient = mock(ServerGrpcQueryClient.class); + ManagedChannel mockChannel = mock(ManagedChannel.class); + when(mockClient.getChannel()).thenReturn(mockChannel); + when(mockChannel.getState(true)).thenReturn(ConnectivityState.READY); + handler.getStreamingQueryClient().getGrpcQueryClientMap().put("host1_8090", mockClient); + + assertEquals(retrier.apply("Server_host1_8099"), FailureDetector.ServerState.HEALTHY); + } + + @Test + public void testRetryReturnsUnhealthyWhenChannelNotReady() { + FailureDetector failureDetector = mock(FailureDetector.class); + @SuppressWarnings("unchecked") + ArgumentCaptor> retrierCaptor = + ArgumentCaptor.forClass(Function.class); + + RoutingManager routingManager = mock(RoutingManager.class); + Map enabledServers = new HashMap<>(); + enabledServers.put("Server_host1_8099", mockServerInstance("host1", 8090)); + when(routingManager.getEnabledServerInstanceMap()).thenReturn(enabledServers); + + GrpcBrokerRequestHandler handler = createHandler(routingManager, failureDetector); + verify(failureDetector).registerUnhealthyServerRetrier(retrierCaptor.capture()); + Function retrier = retrierCaptor.getValue(); + + // Inject a client with a TRANSIENT_FAILURE channel + ServerGrpcQueryClient mockClient = mock(ServerGrpcQueryClient.class); + ManagedChannel mockChannel = mock(ManagedChannel.class); + when(mockClient.getChannel()).thenReturn(mockChannel); + when(mockChannel.getState(true)).thenReturn(ConnectivityState.TRANSIENT_FAILURE); + handler.getStreamingQueryClient().getGrpcQueryClientMap().put("host1_8090", mockClient); + + assertEquals(retrier.apply("Server_host1_8099"), FailureDetector.ServerState.UNHEALTHY); + } + + // ======================== + // Helpers + // ======================== + + private static GrpcBrokerRequestHandler.PinotServerStreamingQueryClient createQueryClient() { + return new GrpcBrokerRequestHandler.PinotServerStreamingQueryClient(new GrpcConfig(Collections.emptyMap())); + } + + private static ServerInstance mockServerInstance(String hostname, int grpcPort) { + ServerInstance si = mock(ServerInstance.class); + when(si.getHostname()).thenReturn(hostname); + when(si.getGrpcPort()).thenReturn(grpcPort); + return si; + } + + private static GrpcBrokerRequestHandler createHandler(RoutingManager routingManager, + FailureDetector failureDetector) { + PinotConfiguration config = new PinotConfiguration(); + return new GrpcBrokerRequestHandler(config, "testBrokerId", new BrokerRequestIdGenerator(), routingManager, + new AllowAllAccessControlFactory(), mock(QueryQuotaManager.class), mock( + org.apache.pinot.common.config.provider.TableCache.class), failureDetector, + ThreadAccountantUtils.getNoOpAccountant()); + } } From 3440be2c3066d981bebebbc94576a0affa770406 Mon Sep 17 00:00:00 2001 From: Siddharth Teotia Date: Sat, 11 Apr 2026 00:13:07 -0700 Subject: [PATCH 4/5] Add missing test assertions and combined cleanup+retry test - Assert return value in testRetryUnhealthyServerCleansUpStaleClients - Add testRetryCleanupAndHealthyReturnInSameCall: verifies stale client cleanup AND healthy return for the retried server happen in one call Co-Authored-By: Claude Opus 4.6 (1M context) --- .../GrpcBrokerRequestHandlerTest.java | 45 ++++++++++++++++++- 1 file changed, 43 insertions(+), 2 deletions(-) diff --git a/pinot-broker/src/test/java/org/apache/pinot/broker/requesthandler/GrpcBrokerRequestHandlerTest.java b/pinot-broker/src/test/java/org/apache/pinot/broker/requesthandler/GrpcBrokerRequestHandlerTest.java index f8cdb2e895..7d70a3451e 100644 --- a/pinot-broker/src/test/java/org/apache/pinot/broker/requesthandler/GrpcBrokerRequestHandlerTest.java +++ b/pinot-broker/src/test/java/org/apache/pinot/broker/requesthandler/GrpcBrokerRequestHandlerTest.java @@ -261,13 +261,54 @@ public void testRetryUnhealthyServerCleansUpStaleClients() { ServerGrpcQueryClient staleClient = mock(ServerGrpcQueryClient.class); clientMap.put("host2_8090", staleClient); - // Retrying any server should trigger cleanup of host2's stale client - retrier.apply("Server_host2_8099"); + // Retrying the decommissioned server should clean up its client and return UNHEALTHY + FailureDetector.ServerState result = retrier.apply("Server_host2_8099"); + assertEquals(result, FailureDetector.ServerState.UNHEALTHY); verify(staleClient).close(); assertFalse(clientMap.containsKey("host2_8090")); } + @Test + public void testRetryCleanupAndHealthyReturnInSameCall() { + FailureDetector failureDetector = mock(FailureDetector.class); + @SuppressWarnings("unchecked") + ArgumentCaptor> retrierCaptor = + ArgumentCaptor.forClass(Function.class); + + RoutingManager routingManager = mock(RoutingManager.class); + Map enabledServers = new HashMap<>(); + enabledServers.put("Server_host1_8099", mockServerInstance("host1", 8090)); + when(routingManager.getEnabledServerInstanceMap()).thenReturn(enabledServers); + + GrpcBrokerRequestHandler handler = createHandler(routingManager, failureDetector); + verify(failureDetector).registerUnhealthyServerRetrier(retrierCaptor.capture()); + Function retrier = retrierCaptor.getValue(); + + Map clientMap = handler.getStreamingQueryClient().getGrpcQueryClientMap(); + + // Inject a stale client for a decommissioned server + ServerGrpcQueryClient staleClient = mock(ServerGrpcQueryClient.class); + clientMap.put("host2_8090", staleClient); + + // Inject a healthy client for the server being retried + ServerGrpcQueryClient healthyClient = mock(ServerGrpcQueryClient.class); + ManagedChannel mockChannel = mock(ManagedChannel.class); + when(healthyClient.getChannel()).thenReturn(mockChannel); + when(mockChannel.getState(true)).thenReturn(ConnectivityState.READY); + clientMap.put("host1_8090", healthyClient); + + // Retrying host1: should clean up host2's stale client AND return HEALTHY for host1 + FailureDetector.ServerState result = retrier.apply("Server_host1_8099"); + + assertEquals(result, FailureDetector.ServerState.HEALTHY); + verify(staleClient).close(); + assertFalse(clientMap.containsKey("host2_8090")); + // host1's client should still be there, not closed + verify(healthyClient, never()).close(); + assertTrue(clientMap.containsKey("host1_8090")); + } + @Test public void testRetryReturnsUnhealthyWhenServerNotInRouting() { FailureDetector failureDetector = mock(FailureDetector.class); From 702a64ac10d531f46525aeac0c6df31ca03515cb Mon Sep 17 00:00:00 2001 From: Siddharth Teotia Date: Thu, 16 Apr 2026 01:22:26 -0700 Subject: [PATCH 5/5] Fix OA-1090: silent channel failures and stale idle connections in ServerChannels Two bugs in ServerChannels caused queries to hang until timeout instead of failing fast when a server-side Netty connection was unhealthy: 1. writeAndFlush() failure not propagated: The ChannelFuture listener that fires after every write never checked f.isSuccess(). A write to a dead channel would silently complete the listener, mark the request as sent, and leave the query waiting forever. Fixed by calling QueryRouter.markServerDown() on write failure so all in-flight queries on that channel fail immediately. 2. No idle connection eviction: SO_KEEPALIVE relied on OS-level TCP keepalive (~2h default on Linux). A silently hung connection (network partition, frozen process) would never trigger channelInactive(), and isActive() would still return true. Added an IdleStateHandler (default 300s, configurable via NettyConfig.channelIdleTimeoutSeconds) that closes fully idle channels, firing channelInactive() and triggering the existing markServerDown() path. Also added CONNECT_TIMEOUT_MILLIS (default 10s, configurable via NettyConfig.channelConnectTimeoutMs) to the Bootstrap so that bootstrap.connect().sync() cannot block indefinitely when a server is unreachable. Tests added: - testWriteFailureMarksServerDown: mocks a failed ChannelFuture and verifies QueryRouter.markServerDown() is called and markRequestSent() is not. - testIdleConnectionEviction: connects with a 1s idle timeout and verifies the channel is closed and markServerDown() is called within 5s. Co-Authored-By: Claude Sonnet 4.6 (1M context) --- .../pinot/common/config/NettyConfig.java | 28 ++++++ .../pinot/core/transport/ServerChannels.java | 40 +++++++- .../core/transport/ServerChannelsTest.java | 98 +++++++++++++++++++ 3 files changed, 164 insertions(+), 2 deletions(-) diff --git a/pinot-common/src/main/java/org/apache/pinot/common/config/NettyConfig.java b/pinot-common/src/main/java/org/apache/pinot/common/config/NettyConfig.java index 5bb80fd651..02628780d1 100644 --- a/pinot-common/src/main/java/org/apache/pinot/common/config/NettyConfig.java +++ b/pinot-common/src/main/java/org/apache/pinot/common/config/NettyConfig.java @@ -26,7 +26,15 @@ */ public class NettyConfig { private static final String NATIVE_TRANSPORTS_ENABLED = "native.transports.enabled"; + private static final String CHANNEL_IDLE_TIMEOUT_SECONDS = "channel.idle.timeout.seconds"; + private static final String CHANNEL_CONNECT_TIMEOUT_MS = "channel.connect.timeout.ms"; + + public static final int DEFAULT_CHANNEL_IDLE_TIMEOUT_SECONDS = 300; + public static final int DEFAULT_CHANNEL_CONNECT_TIMEOUT_MS = 10_000; + private boolean _nativeTransportsEnabled = false; + private int _channelIdleTimeoutSeconds = DEFAULT_CHANNEL_IDLE_TIMEOUT_SECONDS; + private int _channelConnectTimeoutMs = DEFAULT_CHANNEL_CONNECT_TIMEOUT_MS; private static String key(String namespace, String suffix) { return namespace + "." + suffix; @@ -42,6 +50,10 @@ public static NettyConfig extractNettyConfig(PinotConfiguration pinotConfig, Str NettyConfig nettyConfig = new NettyConfig(); nettyConfig.setNativeTransportsEnabled(pinotConfig.getProperty(key(namespace, NATIVE_TRANSPORTS_ENABLED), defaultConfig.isNativeTransportsEnabled())); + nettyConfig.setChannelIdleTimeoutSeconds(pinotConfig.getProperty(key(namespace, CHANNEL_IDLE_TIMEOUT_SECONDS), + defaultConfig.getChannelIdleTimeoutSeconds())); + nettyConfig.setChannelConnectTimeoutMs(pinotConfig.getProperty(key(namespace, CHANNEL_CONNECT_TIMEOUT_MS), + defaultConfig.getChannelConnectTimeoutMs())); return nettyConfig; } @@ -53,4 +65,20 @@ public boolean isNativeTransportsEnabled() { public void setNativeTransportsEnabled(boolean nativeTransportsEnabled) { _nativeTransportsEnabled = nativeTransportsEnabled; } + + public int getChannelIdleTimeoutSeconds() { + return _channelIdleTimeoutSeconds; + } + + public void setChannelIdleTimeoutSeconds(int channelIdleTimeoutSeconds) { + _channelIdleTimeoutSeconds = channelIdleTimeoutSeconds; + } + + public int getChannelConnectTimeoutMs() { + return _channelConnectTimeoutMs; + } + + public void setChannelConnectTimeoutMs(int channelConnectTimeoutMs) { + _channelConnectTimeoutMs = channelConnectTimeoutMs; + } } 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..dbc2118774 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 @@ -23,6 +23,8 @@ import io.netty.buffer.PooledByteBufAllocatorMetric; import io.netty.buffer.Unpooled; import io.netty.channel.Channel; +import io.netty.channel.ChannelDuplexHandler; +import io.netty.channel.ChannelHandlerContext; import io.netty.channel.ChannelInitializer; import io.netty.channel.ChannelOption; import io.netty.channel.EventLoopGroup; @@ -35,6 +37,8 @@ import io.netty.channel.nio.NioEventLoopGroup; import io.netty.channel.socket.SocketChannel; import io.netty.channel.socket.nio.NioSocketChannel; +import io.netty.handler.timeout.IdleStateEvent; +import io.netty.handler.timeout.IdleStateHandler; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.TimeUnit; import java.util.concurrent.TimeoutException; @@ -82,9 +86,11 @@ public class ServerChannels { private final EventLoopGroup _eventLoopGroup; private final Class _channelClass; private final ThreadAccountant _threadAccountant; + private final int _channelIdleTimeoutSeconds; + private final int _channelConnectTimeoutMs; private final BrokerMetrics _brokerMetrics = BrokerMetrics.get(); - private final ConcurrentHashMap _serverToChannelMap = new ConcurrentHashMap<>(); + final ConcurrentHashMap _serverToChannelMap = new ConcurrentHashMap<>(); /** * Create a server channel with TLS config @@ -125,6 +131,10 @@ public ServerChannels(QueryRouter queryRouter, @Nullable NettyConfig nettyConfig _queryRouter = queryRouter; _tlsConfig = tlsConfig; _threadAccountant = threadAccountant; + _channelIdleTimeoutSeconds = nettyConfig != null + ? nettyConfig.getChannelIdleTimeoutSeconds() : NettyConfig.DEFAULT_CHANNEL_IDLE_TIMEOUT_SECONDS; + _channelConnectTimeoutMs = nettyConfig != null + ? nettyConfig.getChannelConnectTimeoutMs() : NettyConfig.DEFAULT_CHANNEL_CONNECT_TIMEOUT_MS; } public void sendRequest(String rawTableName, AsyncQueryResponse asyncQueryResponse, @@ -175,7 +185,9 @@ class ServerChannel { _bootstrap = new Bootstrap().remoteAddress(serverRoutingInstance.getHostname(), serverRoutingInstance.getPort()) .option(ChannelOption.ALLOCATOR, bufAllocatorWithLimits).group(_eventLoopGroup).channel(_channelClass) - .option(ChannelOption.SO_KEEPALIVE, true).handler(new ChannelInitializer() { + .option(ChannelOption.SO_KEEPALIVE, true) + .option(ChannelOption.CONNECT_TIMEOUT_MILLIS, _channelConnectTimeoutMs) + .handler(new ChannelInitializer() { @Override protected void initChannel(SocketChannel ch) { if (_tlsConfig != null) { @@ -186,6 +198,22 @@ protected void initChannel(SocketChannel ch) { ch.pipeline().addLast(ChannelHandlerFactory.getLengthFieldBasedFrameDecoder()); ch.pipeline().addLast(ChannelHandlerFactory.getLengthFieldPrepender()); + // Close channels that have been completely idle (no read or write) for the configured duration. + // This evicts stale connections to unreachable servers before they are reused for a new query. + // When closed, channelInactive() fires on DataTableHandler which marks any in-flight queries as failed. + ch.pipeline().addLast(new IdleStateHandler(0, 0, _channelIdleTimeoutSeconds)); + ch.pipeline().addLast(new ChannelDuplexHandler() { + @Override + public void userEventTriggered(ChannelHandlerContext ctx, Object evt) throws Exception { + if (evt instanceof IdleStateEvent) { + LOGGER.warn("Channel to server: {} has been idle for {}s, closing connection", + _serverRoutingInstance, _channelIdleTimeoutSeconds); + ctx.close(); + } else { + super.userEventTriggered(ctx, evt); + } + } + }); ch.pipeline().addLast( ChannelHandlerFactory.getDirectOOMHandler(_queryRouter, _serverRoutingInstance, _serverToChannelMap, null, null)); @@ -242,6 +270,14 @@ void sendRequestWithoutLocking(String rawTableName, AsyncQueryResponse asyncQuer ServerRoutingInstance serverRoutingInstance, byte[] requestBytes) { long startTimeMs = System.currentTimeMillis(); _channel.writeAndFlush(Unpooled.wrappedBuffer(requestBytes)).addListener(f -> { + if (!f.isSuccess()) { + // The write failed — the channel is in a bad state. Mark the server down so all in-flight queries + // on this channel fail immediately rather than hanging until query timeout. + LOGGER.error("Failed to send request to server: {}", serverRoutingInstance, f.cause()); + _queryRouter.markServerDown(serverRoutingInstance, + new RuntimeException("Failed to send request to server: " + serverRoutingInstance, f.cause())); + return; + } int requestSentLatencyMs = (int) (System.currentTimeMillis() - startTimeMs); _brokerMetrics.addTimedTableValue(rawTableName, BrokerTimer.NETTY_CONNECTION_SEND_REQUEST_LATENCY, requestSentLatencyMs, TimeUnit.MILLISECONDS); 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..a363b71ed2 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,18 +19,30 @@ package org.apache.pinot.core.transport; import com.sun.net.httpserver.HttpServer; +import io.netty.channel.Channel; +import io.netty.channel.ChannelFuture; +import io.netty.util.concurrent.GenericFutureListener; import java.net.InetSocketAddress; import org.apache.pinot.common.config.NettyConfig; import org.apache.pinot.common.request.BrokerRequest; import org.apache.pinot.common.request.InstanceRequest; import org.apache.pinot.spi.accounting.ThreadAccountantUtils; import org.apache.pinot.spi.config.table.TableType; +import org.apache.pinot.util.TestUtils; import org.testng.annotations.AfterClass; import org.testng.annotations.BeforeClass; import org.testng.annotations.DataProvider; import org.testng.annotations.Test; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.doAnswer; import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; +import static org.testng.Assert.assertNotNull; +import static org.testng.Assert.assertTrue; public class ServerChannelsTest { @@ -82,4 +94,90 @@ public void testConnect(boolean nativeTransportEnabled) serverChannels.sendRequest("dummy_table_name", asyncQueryResponse, serverRoutingInstance, instanceRequest, 1000); serverChannels.shutDown(); } + + /** + * Verifies that a writeAndFlush failure immediately marks the server down on the QueryRouter + * rather than leaving the query hanging until timeout. + * + * Scenario: the channel appears active (isActive() == true) but the actual write to the + * network fails — a race that can happen when a connection dies between the isActive() check + * and the writeAndFlush() call. + */ + @Test + public void testWriteFailureMarksServerDown() + throws Exception { + QueryRouter mockQueryRouter = mock(QueryRouter.class); + ServerRoutingInstance serverRoutingInstance = + new ServerRoutingInstance("localhost", 9999, TableType.REALTIME); + + ServerChannels serverChannels = + new ServerChannels(mockQueryRouter, null, null, ThreadAccountantUtils.getNoOpAccountant()); + + // Create a ServerChannel and inject a mock Channel that returns a failed ChannelFuture + ServerChannels.ServerChannel serverChannel = serverChannels.new ServerChannel(serverRoutingInstance); + + Channel mockChannel = mock(Channel.class); + ChannelFuture failedFuture = mock(ChannelFuture.class); + when(failedFuture.isSuccess()).thenReturn(false); + when(failedFuture.cause()).thenReturn(new RuntimeException("Simulated write failure")); + when(mockChannel.isActive()).thenReturn(true); + when(mockChannel.writeAndFlush(any())).thenReturn(failedFuture); + // Make the listener fire synchronously so we can assert without sleeping + doAnswer(invocation -> { + GenericFutureListener listener = invocation.getArgument(0); + listener.operationComplete(failedFuture); + return failedFuture; + }).when(failedFuture).addListener(any(GenericFutureListener.class)); + + // Inject the mock channel directly (package-private field, same package) + serverChannel._channel = mockChannel; + + AsyncQueryResponse mockAsyncResponse = mock(AsyncQueryResponse.class); + + // Trigger the write + serverChannel.sendRequestWithoutLocking("testTable", mockAsyncResponse, serverRoutingInstance, new byte[]{1, 2, 3}); + + // The router must be told the server is down so all in-flight queries fail fast + verify(mockQueryRouter).markServerDown(eq(serverRoutingInstance), any(RuntimeException.class)); + + // markRequestSent must NOT be called — the request was never delivered + verify(mockAsyncResponse, never()).markRequestSent(any(), any(Integer.class)); + + serverChannels.shutDown(); + } + + /** + * Verifies that a channel idle for longer than the configured timeout is closed automatically. + * This evicts stale connections so that the next query triggers a fresh connect attempt, + * failing fast rather than writing into a silently dead TCP connection. + * + * Uses a 1-second idle timeout so the test runs quickly without long sleeps. + */ + @Test + public void testIdleConnectionEviction() + throws Exception { + NettyConfig nettyConfig = new NettyConfig(); + nettyConfig.setChannelIdleTimeoutSeconds(1); + + QueryRouter mockQueryRouter = mock(QueryRouter.class); + ServerRoutingInstance serverRoutingInstance = + new ServerRoutingInstance("localhost", _dummyServer.getAddress().getPort(), TableType.REALTIME); + + ServerChannels serverChannels = + new ServerChannels(mockQueryRouter, nettyConfig, null, ThreadAccountantUtils.getNoOpAccountant()); + serverChannels.connect(serverRoutingInstance); + + ServerChannels.ServerChannel serverChannel = serverChannels._serverToChannelMap.get(serverRoutingInstance); + assertNotNull(serverChannel); + assertTrue(serverChannel._channel.isActive(), "Channel should be active immediately after connect"); + + // Wait for the idle timeout (1s) to fire and close the channel + TestUtils.waitForCondition(aVoid -> !serverChannel._channel.isActive(), 5_000L, + "Channel should have been closed by idle timeout"); + + // channelInactive() fires on close, which calls QueryRouter.markServerDown() for in-flight queries + verify(mockQueryRouter).markServerDown(eq(serverRoutingInstance), any(RuntimeException.class)); + + serverChannels.shutDown(); + } }