diff --git a/gateway-ha/src/main/java/io/trino/gateway/ha/config/HaGatewayConfiguration.java b/gateway-ha/src/main/java/io/trino/gateway/ha/config/HaGatewayConfiguration.java index 96313b36f..8d717ae4f 100644 --- a/gateway-ha/src/main/java/io/trino/gateway/ha/config/HaGatewayConfiguration.java +++ b/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; diff --git a/gateway-ha/src/main/java/io/trino/gateway/ha/config/RoutingConfiguration.java b/gateway-ha/src/main/java/io/trino/gateway/ha/config/RoutingConfiguration.java new file mode 100644 index 000000000..93136f2c6 --- /dev/null +++ b/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; + } +} diff --git a/gateway-ha/src/main/java/io/trino/gateway/proxyserver/ProxyRequestHandler.java b/gateway-ha/src/main/java/io/trino/gateway/proxyserver/ProxyRequestHandler.java index 2aada9fa6..d65629b59 100644 --- a/gateway-ha/src/main/java/io/trino/gateway/proxyserver/ProxyRequestHandler.java +++ b/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 private void setupAsyncResponse(AsyncResponse asyncResponse, ListenableFuture 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()); }