Support HttpClient options for external routing group selector

docs/routing-rules.md

@@ -50,6 +50,21 @@ If there is error parsing the routing rules configuration file, an error is
 logged, and requests are routed using the routing group header
 `X-Trino-Routing-Group` as default.
 
+### Configuring API Requests with HttpClient Config
+
+Trino Gateway uses the Airlift HttpClient library when making HTTP requests.
+You can configure the HttpClient by adding the following configuration to
+the `serverConfig:` section with the `router` prefix.
+
+```yaml
+serverConfig:
+    router.http-client.request-timeout: 1s
+```
+
+Please refer to the [Trino HTTP Client properties](
+https://trino.io/docs/current/admin/properties-http-client.html)
+documentation for more.
+
 ### Use an external service for routing rules
 
 You can use an external service for processing your routing by setting the

gateway-ha/src/main/java/io/trino/gateway/baseapp/BaseApp.java

@@ -31,6 +31,7 @@
 import io.trino.gateway.ha.resource.LoginResource;
 import io.trino.gateway.ha.resource.PublicResource;
 import io.trino.gateway.ha.resource.TrinoResource;
+import io.trino.gateway.ha.router.ForRouter;
 import io.trino.gateway.ha.security.AuthorizedExceptionMapper;
 import io.trino.gateway.proxyserver.ForProxy;
 import io.trino.gateway.proxyserver.ProxyRequestHandler;
@@ -170,5 +171,6 @@ private static void registerProxyResources(Binder binder)
         jaxrsBinder(binder).bind(ProxyRequestHandler.class);
         httpClientBinder(binder).bindHttpClient("proxy", ForProxy.class);
         httpClientBinder(binder).bindHttpClient("monitor", ForMonitor.class);
+        httpClientBinder(binder).bindHttpClient("router", ForRouter.class);
     }
 }

gateway-ha/src/main/java/io/trino/gateway/ha/module/HaGatewayProviderModule.java

@@ -17,6 +17,7 @@
 import com.google.inject.AbstractModule;
 import com.google.inject.Provides;
 import com.google.inject.Singleton;
+import io.airlift.http.client.HttpClient;
 import io.trino.gateway.ha.config.AuthenticationConfiguration;
 import io.trino.gateway.ha.config.AuthorizationConfiguration;
 import io.trino.gateway.ha.config.GatewayCookieConfigurationPropertiesProvider;
@@ -26,6 +27,7 @@
 import io.trino.gateway.ha.config.RulesExternalConfiguration;
 import io.trino.gateway.ha.config.UserConfiguration;
 import io.trino.gateway.ha.router.BackendStateManager;
+import io.trino.gateway.ha.router.ForRouter;
 import io.trino.gateway.ha.router.RoutingGroupSelector;
 import io.trino.gateway.ha.security.ApiAuthenticator;
 import io.trino.gateway.ha.security.AuthorizationManager;
@@ -176,7 +178,7 @@ public BackendStateManager getBackendStateConnectionManager()
 
     @Provides
     @Singleton
-    public RoutingGroupSelector getRoutingGroupSelector()
+    public RoutingGroupSelector getRoutingGroupSelector(@ForRouter HttpClient httpClient)
     {
         RoutingRulesConfiguration routingRulesConfig = configuration.getRoutingRules();
         if (routingRulesConfig.isRulesEngineEnabled()) {
@@ -188,7 +190,7 @@ public RoutingGroupSelector getRoutingGroupSelector()
                     }
                     case EXTERNAL -> {
                         RulesExternalConfiguration rulesExternalConfiguration = routingRulesConfig.getRulesExternalConfiguration();
-                        yield RoutingGroupSelector.byRoutingExternal(rulesExternalConfiguration, configuration.getRequestAnalyzerConfig());
+                        yield RoutingGroupSelector.byRoutingExternal(httpClient, rulesExternalConfiguration, configuration.getRequestAnalyzerConfig());
                     }
                 };
             }

gateway-ha/src/main/java/io/trino/gateway/ha/router/ExternalRoutingGroupSelector.java

@@ -18,11 +18,9 @@
 import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.Multimap;
 import io.airlift.http.client.HttpClient;
-import io.airlift.http.client.HttpClientConfig;
 import io.airlift.http.client.JsonBodyGenerator;
 import io.airlift.http.client.JsonResponseHandler;
 import io.airlift.http.client.Request;
-import io.airlift.http.client.jetty.JettyHttpClient;
 import io.airlift.json.JsonCodec;
 import io.airlift.log.Logger;
 import io.trino.gateway.ha.config.RequestAnalyzerConfig;
@@ -59,7 +57,7 @@ public class ExternalRoutingGroupSelector
             createJsonResponseHandler(jsonCodec(RoutingGroupExternalResponse.class));
 
     @VisibleForTesting
-    ExternalRoutingGroupSelector(RulesExternalConfiguration rulesExternalConfiguration, RequestAnalyzerConfig requestAnalyzerConfig)
+    ExternalRoutingGroupSelector(RulesExternalConfiguration rulesExternalConfiguration, RequestAnalyzerConfig requestAnalyzerConfig, HttpClient httpClient)
     {
         this.excludeHeaders = ImmutableSet.<String>builder()
                 .add("Content-Length")
@@ -76,7 +74,7 @@ public class ExternalRoutingGroupSelector
             throw new RuntimeException("Invalid URL provided, using "
                     + "routing group header as default.", e);
         }
-        httpClient = new JettyHttpClient(new HttpClientConfig());
+        this.httpClient = requireNonNull(httpClient, "httpClient is null");
     }
 
     @Override

gateway-ha/src/main/java/io/trino/gateway/ha/router/ForRouter.java

@@ -0,0 +1,29 @@
+/*
+ * Licensed 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 io.trino.gateway.ha.router;
+
+import com.google.inject.BindingAnnotation;
+
+import java.lang.annotation.Retention;
+import java.lang.annotation.Target;
+
+import static java.lang.annotation.ElementType.FIELD;
+import static java.lang.annotation.ElementType.METHOD;
+import static java.lang.annotation.ElementType.PARAMETER;
+import static java.lang.annotation.RetentionPolicy.RUNTIME;
+
+@Retention(RUNTIME)
+@Target({FIELD, PARAMETER, METHOD})
+@BindingAnnotation
+public @interface ForRouter{}

gateway-ha/src/main/java/io/trino/gateway/ha/router/RoutingGroupSelector.java

@@ -13,6 +13,7 @@
  */
 package io.trino.gateway.ha.router;
 
+import io.airlift.http.client.HttpClient;
 import io.trino.gateway.ha.config.RequestAnalyzerConfig;
 import io.trino.gateway.ha.config.RulesExternalConfiguration;
 import jakarta.servlet.http.HttpServletRequest;
@@ -46,9 +47,12 @@ static RoutingGroupSelector byRoutingRulesEngine(String rulesConfigPath, Request
      * Routing group selector that uses RESTful API
      * to determine the right routing group.
      */
-    static RoutingGroupSelector byRoutingExternal(RulesExternalConfiguration rulesExternalConfiguration, RequestAnalyzerConfig requestAnalyzerConfig)
+    static RoutingGroupSelector byRoutingExternal(
+            HttpClient httpClient,
+            RulesExternalConfiguration rulesExternalConfiguration,
+            RequestAnalyzerConfig requestAnalyzerConfig)
     {
-        return new ExternalRoutingGroupSelector(rulesExternalConfiguration, requestAnalyzerConfig);
+        return new ExternalRoutingGroupSelector(rulesExternalConfiguration, requestAnalyzerConfig, httpClient);
     }
 
     /**

gateway-ha/src/test/java/io/trino/gateway/ha/router/TestRoutingGroupSelectorExternal.java

@@ -133,7 +133,7 @@ void testApiFailure()
     {
         RulesExternalConfiguration rulesExternalConfiguration = provideRoutingRuleExternalConfig();
         RoutingGroupSelector routingGroupSelector =
-                RoutingGroupSelector.byRoutingExternal(rulesExternalConfiguration, requestAnalyzerConfig);
+                RoutingGroupSelector.byRoutingExternal(httpClient, rulesExternalConfiguration, requestAnalyzerConfig);
 
         HttpServletRequest mockRequest = prepareMockRequest();
         setMockHeaders(mockRequest);
@@ -162,7 +162,7 @@ void testNullUri()
         rulesExternalConfiguration.setUrlPath(null);
 
         // Assert that a RuntimeException is thrown with message
-        assertThatThrownBy(() -> RoutingGroupSelector.byRoutingExternal(rulesExternalConfiguration, requestAnalyzerConfig))
+        assertThatThrownBy(() -> RoutingGroupSelector.byRoutingExternal(httpClient, rulesExternalConfiguration, requestAnalyzerConfig))
                 .isInstanceOf(RuntimeException.class)
                 .hasMessage("Invalid URL provided, using routing group header as default.");
     }
@@ -175,7 +175,7 @@ void testExcludeHeader()
         rulesExternalConfiguration.setExcludeHeaders(List.of("test-exclude-header"));
 
         RoutingGroupSelector routingGroupSelector =
-                RoutingGroupSelector.byRoutingExternal(rulesExternalConfiguration, requestAnalyzerConfig);
+                RoutingGroupSelector.byRoutingExternal(httpClient, rulesExternalConfiguration, requestAnalyzerConfig);
 
         // Mock headers to be read by mockRequest
         HttpServletRequest mockRequest = mock(HttpServletRequest.class);