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..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 @@ -18,11 +18,14 @@ */ 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; @@ -87,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, @@ -170,17 +178,51 @@ 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()); + try { + entry.getValue().close(); + } catch (Exception e) { + LOGGER.warn("Failed to close gRPC client for server: {}", entry.getKey(), e); + } + iterator.remove(); + } + } + } + public void shutdown() { for (ServerGrpcQueryClient client : _grpcQueryClientMap.values()) { client.close(); } } + + @VisibleForTesting + Map getGrpcQueryClientMap() { + return _grpcQueryClientMap; + } } /** * 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) { 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..7d70a3451e --- /dev/null +++ b/pinot-broker/src/test/java/org/apache/pinot/broker/requesthandler/GrpcBrokerRequestHandlerTest.java @@ -0,0 +1,424 @@ +/** + * 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 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; +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 { + + @BeforeClass + public void setUp() { + BrokerMetrics.register(mock(BrokerMetrics.class)); + BrokerQueryEventListenerFactory.init(new PinotConfiguration()); + } + + // ======================== + // PinotServerStreamingQueryClient.cleanupStaleClients tests + // ======================== + + @Test + public void testCleanupRemovesDecommissionedServer() { + GrpcBrokerRequestHandler.PinotServerStreamingQueryClient queryClient = createQueryClient(); + + 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 + Map enabledServers = new HashMap<>(); + enabledServers.put("Server_host1_8099", mockServerInstance("host1", 8090)); + enabledServers.put("Server_host3_8099", mockServerInstance("host3", 8090)); + + queryClient.cleanupStaleClients(enabledServers); + + verify(mockClient2).close(); + 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 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); + + Map clientMap = queryClient.getGrpcQueryClientMap(); + clientMap.put("host1_8090", mockClient1); + clientMap.put("host2_8090", mockClient2); + + queryClient.cleanupStaleClients(Collections.emptyMap()); + + verify(mockClient1).close(); + verify(mockClient2).close(); + assertTrue(clientMap.isEmpty()); + } + + @Test + public void testCleanupNoOpWhenAllServersActive() { + GrpcBrokerRequestHandler.PinotServerStreamingQueryClient queryClient = createQueryClient(); + + ServerGrpcQueryClient mockClient1 = mock(ServerGrpcQueryClient.class); + + Map clientMap = queryClient.getGrpcQueryClientMap(); + clientMap.put("host1_8090", mockClient1); + + Map enabledServers = new HashMap<>(); + enabledServers.put("Server_host1_8099", mockServerInstance("host1", 8090)); + + queryClient.cleanupStaleClients(enabledServers); + + verify(mockClient1, never()).close(); + assertEquals(clientMap.size(), 1); + } + + @Test + public void testCleanupNoOpWhenClientMapEmpty() { + GrpcBrokerRequestHandler.PinotServerStreamingQueryClient queryClient = createQueryClient(); + + Map enabledServers = new HashMap<>(); + 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 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); + @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()); + } +} 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(); + } }