Make asyncTimeout configurable

gateway-ha/src/main/java/io/trino/gateway/ha/config/HaGatewayConfiguration.java

@@ -24,6 +24,7 @@ public class HaGatewayConfiguration
     private NotifierConfiguration notifier;
     private DataStoreConfiguration dataStore;
     private MonitorConfiguration monitor = new MonitorConfiguration();
+    private RoutingConfiguration routing = new RoutingConfiguration();
     private RoutingRulesConfiguration routingRules = new RoutingRulesConfiguration();
     private AuthenticationConfiguration authentication;
     private AuthorizationConfiguration authorization;
@@ -83,6 +84,16 @@ public void setMonitor(MonitorConfiguration monitor)
         this.monitor = monitor;
     }
 
+    public RoutingConfiguration getRouting()
+    {
+        return routing;
+    }
+
+    public void setRouting(RoutingConfiguration routing)
+    {
+        this.routing = routing;
+    }
+
     public RoutingRulesConfiguration getRoutingRules()
     {
         return this.routingRules;

gateway-ha/src/main/java/io/trino/gateway/ha/config/RoutingConfiguration.java

@@ -0,0 +1,33 @@
+/*
+ * 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.config;
+
+import io.airlift.units.Duration;
+
+import static java.util.concurrent.TimeUnit.MINUTES;
+
+public class RoutingConfiguration
+{
+    private Duration asyncTimeout = new Duration(2, MINUTES);
+
+    public Duration getAsyncTimeout()
+    {
+        return asyncTimeout;
+    }
+
+    public void setAsyncTimeout(Duration asyncTimeout)
+    {
+        this.asyncTimeout = asyncTimeout;
+    }
+}

gateway-ha/src/main/java/io/trino/gateway/proxyserver/ProxyRequestHandler.java

@@ -24,6 +24,7 @@
 import io.airlift.log.Logger;
 import io.airlift.units.Duration;
 import io.trino.gateway.ha.config.GatewayCookieConfigurationPropertiesProvider;
+import io.trino.gateway.ha.config.HaGatewayConfiguration;
 import io.trino.gateway.ha.router.GatewayCookie;
 import io.trino.gateway.ha.router.OAuth2GatewayCookie;
 import io.trino.gateway.ha.router.QueryHistoryManager;
@@ -70,14 +71,13 @@
 import static java.util.Collections.list;
 import static java.util.Objects.requireNonNull;
 import static java.util.concurrent.Executors.newCachedThreadPool;
-import static java.util.concurrent.TimeUnit.MINUTES;
 
 public class ProxyRequestHandler
 {
     private static final Logger log = Logger.get(ProxyRequestHandler.class);
-    private static final Duration ASYNC_TIMEOUT = new Duration(2, MINUTES);
     private static final ObjectMapper OBJECT_MAPPER = new ObjectMapper();
 
+    private final Duration asyncTimeout;
     private final ExecutorService executor = newCachedThreadPool(daemonThreadsNamed("proxy-%s"));
     private final HttpClient httpClient;
     private final RoutingManager routingManager;
@@ -88,12 +88,14 @@ public class ProxyRequestHandler
     public ProxyRequestHandler(
             @ForProxy HttpClient httpClient,
             RoutingManager routingManager,
-            QueryHistoryManager queryHistoryManager)
+            QueryHistoryManager queryHistoryManager,
+            HaGatewayConfiguration haGatewayConfiguration)
     {
         this.httpClient = requireNonNull(httpClient, "httpClient is null");
         this.routingManager = requireNonNull(routingManager, "routingManager is null");
         this.queryHistoryManager = requireNonNull(queryHistoryManager, "queryHistoryManager is null");
         cookiesEnabled = GatewayCookieConfigurationPropertiesProvider.getInstance().isEnabled();
+        asyncTimeout = haGatewayConfiguration.getRouting().getAsyncTimeout();
     }
 
     @PreDestroy
@@ -198,10 +200,10 @@ private Response buildResponse(ProxyResponse response, ImmutableList<NewCookie>
     private void setupAsyncResponse(AsyncResponse asyncResponse, ListenableFuture<Response> future)
     {
         bindAsyncResponse(asyncResponse, future, executor)
-                .withTimeout(ASYNC_TIMEOUT, () -> Response
+                .withTimeout(asyncTimeout, () -> Response
                         .status(BAD_GATEWAY)
                         .type(TEXT_PLAIN_TYPE)
-                        .entity("Request to remote Trino server timed out after" + ASYNC_TIMEOUT)
+                        .entity("Request to remote Trino server timed out after" + asyncTimeout)
                         .build());
     }