Support configurable ExchangeManager

Author: yangshangqing95

core/trino-main/src/main/java/io/trino/exchange/ExchangeManagerConfig.java

@@ -0,0 +1,35 @@
+/*
+ * 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.exchange;
+
+import io.airlift.configuration.Config;
+
+import java.io.File;
+
+public class ExchangeManagerConfig
+{
+    private File exchangeManagerConfigFile = new File("etc/exchange-manager.properties");
+
+    public File getExchangeManagerConfigFile()
+    {
+        return exchangeManagerConfigFile;
+    }
+
+    @Config("exchange-manager.config-file")
+    public ExchangeManagerConfig setExchangeManagerConfigFile(String exchangeManagerConfigFile)
+    {
+        this.exchangeManagerConfigFile = new File(exchangeManagerConfigFile);
+        return this;
+    }
+}

core/trino-main/src/main/java/io/trino/exchange/ExchangeManagerModule.java

@@ -17,12 +17,15 @@
 import com.google.inject.Module;
 import com.google.inject.Scopes;
 
+import static io.airlift.configuration.ConfigBinder.configBinder;
+
 public class ExchangeManagerModule
         implements Module
 {
     @Override
     public void configure(Binder binder)
     {
+        configBinder(binder).bindConfig(ExchangeManagerConfig.class);
         binder.bind(ExchangeManagerRegistry.class).in(Scopes.SINGLETON);
     }
 }

core/trino-main/src/main/java/io/trino/exchange/ExchangeManagerRegistry.java

@@ -43,7 +43,6 @@ public class ExchangeManagerRegistry
 {
     private static final Logger log = Logger.get(ExchangeManagerRegistry.class);
 
-    private static final File CONFIG_FILE = new File("etc/exchange-manager.properties");
     private static final String EXCHANGE_MANAGER_NAME_PROPERTY = "exchange-manager.name";
 
     private final OpenTelemetry openTelemetry;
@@ -52,16 +51,19 @@ public class ExchangeManagerRegistry
 
     private volatile ExchangeManager exchangeManager;
     private final SecretsResolver secretsResolver;
+    private final ExchangeManagerConfig config;
 
     @Inject
     public ExchangeManagerRegistry(
             OpenTelemetry openTelemetry,
             Tracer tracer,
-            SecretsResolver secretsResolver)
+            SecretsResolver secretsResolver,
+            ExchangeManagerConfig config)
     {
         this.openTelemetry = requireNonNull(openTelemetry, "openTelemetry is null");
         this.tracer = requireNonNull(tracer, "tracer is null");
         this.secretsResolver = requireNonNull(secretsResolver, "secretsResolver is null");
+        this.config = requireNonNull(config, "config is null");
     }
 
     public void addExchangeManagerFactory(ExchangeManagerFactory factory)
@@ -74,13 +76,14 @@ public void addExchangeManagerFactory(ExchangeManagerFactory factory)
 
     public void loadExchangeManager()
     {
-        if (!CONFIG_FILE.exists()) {
+        File configFile = config.getExchangeManagerConfigFile();
+        if (!configFile.exists()) {
             return;
         }
 
-        Map<String, String> properties = loadProperties(CONFIG_FILE);
+        Map<String, String> properties = loadProperties(configFile);
         String name = properties.remove(EXCHANGE_MANAGER_NAME_PROPERTY);
-        checkArgument(!isNullOrEmpty(name), "Exchange manager configuration %s does not contain %s", CONFIG_FILE, EXCHANGE_MANAGER_NAME_PROPERTY);
+        checkArgument(!isNullOrEmpty(name), "Exchange manager configuration %s does not contain %s", configFile, EXCHANGE_MANAGER_NAME_PROPERTY);
 
         loadExchangeManager(name, properties);
     }

core/trino-main/src/main/java/io/trino/server/testing/TestingTrinoServer.java

@@ -50,6 +50,7 @@
 import io.trino.dispatcher.DispatchManager;
 import io.trino.eventlistener.EventListenerConfig;
 import io.trino.eventlistener.EventListenerManager;
+import io.trino.exchange.ExchangeManagerConfig;
 import io.trino.exchange.ExchangeManagerRegistry;
 import io.trino.execution.FailureInjector;
 import io.trino.execution.FailureInjector.InjectedFailureType;
@@ -312,6 +313,7 @@ private TestingTrinoServer(
                             .addBinding()
                             .to(TracingServletFilter.class);
                     binder.bind(EventListenerConfig.class).in(Scopes.SINGLETON);
+                    binder.bind(ExchangeManagerConfig.class).in(Scopes.SINGLETON);
                     binder.bind(AccessControlConfig.class).in(Scopes.SINGLETON);
                     binder.bind(TestingAccessControlManager.class).in(Scopes.SINGLETON);
                     binder.bind(TestingGroupProvider.class).in(Scopes.SINGLETON);

core/trino-main/src/main/java/io/trino/testing/PlanTester.java

@@ -60,6 +60,7 @@
 import io.trino.cost.TaskCountEstimator;
 import io.trino.eventlistener.EventListenerConfig;
 import io.trino.eventlistener.EventListenerManager;
+import io.trino.exchange.ExchangeManagerConfig;
 import io.trino.exchange.ExchangeManagerRegistry;
 import io.trino.execution.DynamicFilterConfig;
 import io.trino.execution.NodeTaskMap;
@@ -473,7 +474,7 @@ private PlanTester(Session defaultSession, int nodeCountForStats)
                 ImmutableSet.of(),
                 ImmutableSet.of(new ExcludeColumnsFunction()));
 
-        exchangeManagerRegistry = new ExchangeManagerRegistry(noop(), noopTracer(), secretsResolver);
+        exchangeManagerRegistry = new ExchangeManagerRegistry(noop(), noopTracer(), secretsResolver, new ExchangeManagerConfig());
         spoolingManagerRegistry = new SpoolingManagerRegistry(
                 new InternalNode("nodeId", URI.create("http://localhost:8080"), NodeVersion.UNKNOWN, false),
                 new ServerConfig(),

core/trino-main/src/test/java/io/trino/exchange/TestExchangeConfig.java

@@ -0,0 +1,50 @@
+/*
+ * 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.exchange;
+
+import com.google.common.collect.ImmutableMap;
+import org.junit.jupiter.api.Test;
+
+import java.io.IOException;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.util.Map;
+
+import static io.airlift.configuration.testing.ConfigAssertions.assertFullMapping;
+import static io.airlift.configuration.testing.ConfigAssertions.assertRecordedDefaults;
+import static io.airlift.configuration.testing.ConfigAssertions.recordDefaults;
+
+public class TestExchangeConfig
+{
+    @Test
+    public void testDefaults()
+    {
+        assertRecordedDefaults(recordDefaults(ExchangeManagerConfig.class)
+                .setExchangeManagerConfigFile("etc/exchange-manager.properties"));
+    }
+
+    @Test
+    public void testExplicitPropertyMappings()
+            throws IOException
+    {
+        Path exchangeConfig = Files.createTempFile(null, null);
+
+        Map<String, String> properties = ImmutableMap.of("exchange-manager.config-file", exchangeConfig.toString());
+
+        ExchangeManagerConfig expected = new ExchangeManagerConfig()
+                .setExchangeManagerConfigFile(exchangeConfig.toFile().getPath());
+
+        assertFullMapping(properties, expected);
+    }
+}

core/trino-main/src/test/java/io/trino/exchange/TestLazyExchangeDataSource.java

@@ -45,7 +45,7 @@ public void testIsBlockedCancellationIsolationInInitializationPhase()
                     throw new UnsupportedOperationException();
                 },
                 RetryPolicy.NONE,
-                new ExchangeManagerRegistry(OpenTelemetry.noop(), Tracing.noopTracer(), new SecretsResolver(ImmutableMap.of())))) {
+                new ExchangeManagerRegistry(OpenTelemetry.noop(), Tracing.noopTracer(), new SecretsResolver(ImmutableMap.of()), new ExchangeManagerConfig()))) {
             ListenableFuture<Void> first = source.isBlocked();
             ListenableFuture<Void> second = source.isBlocked();
             assertThat(first)

core/trino-main/src/test/java/io/trino/execution/BaseTestSqlTaskManager.java

@@ -29,6 +29,7 @@
 import io.trino.connector.CatalogHandle;
 import io.trino.connector.ConnectorServices;
 import io.trino.connector.ConnectorServicesProvider;
+import io.trino.exchange.ExchangeManagerConfig;
 import io.trino.exchange.ExchangeManagerRegistry;
 import io.trino.execution.buffer.BufferResult;
 import io.trino.execution.buffer.BufferState;
@@ -337,7 +338,7 @@ private SqlTaskManager createSqlTaskManager(TaskManagerConfig taskManagerConfig,
                 new NodeSpillConfig(),
                 new TestingGcMonitor(),
                 noopTracer(),
-                new ExchangeManagerRegistry(OpenTelemetry.noop(), Tracing.noopTracer(), new SecretsResolver(ImmutableMap.of())));
+                new ExchangeManagerRegistry(OpenTelemetry.noop(), Tracing.noopTracer(), new SecretsResolver(ImmutableMap.of()), new ExchangeManagerConfig()));
     }
 
     private TaskInfo createTask(SqlTaskManager sqlTaskManager, TaskId taskId, Set<ScheduledSplit> splits, OutputBuffers outputBuffers)

core/trino-main/src/test/java/io/trino/execution/MockRemoteTaskFactory.java

@@ -31,6 +31,7 @@
 import io.opentelemetry.api.trace.Span;
 import io.trino.Session;
 import io.trino.cost.StatsAndCosts;
+import io.trino.exchange.ExchangeManagerConfig;
 import io.trino.exchange.ExchangeManagerRegistry;
 import io.trino.execution.NodeTaskMap.PartitionedSplitCountTracker;
 import io.trino.execution.buffer.LazyOutputBuffer;
@@ -235,7 +236,7 @@ public MockRemoteTask(
                     DataSize.ofBytes(1),
                     () -> new SimpleLocalMemoryContext(newSimpleAggregatedMemoryContext(), "test"),
                     () -> {},
-                    new ExchangeManagerRegistry(OpenTelemetry.noop(), Tracing.noopTracer(), new SecretsResolver(ImmutableMap.of())));
+                    new ExchangeManagerRegistry(OpenTelemetry.noop(), Tracing.noopTracer(), new SecretsResolver(ImmutableMap.of()), new ExchangeManagerConfig()));
 
             this.fragment = requireNonNull(fragment, "fragment is null");
             this.nodeId = requireNonNull(nodeId, "nodeId is null");

core/trino-main/src/test/java/io/trino/execution/TaskTestUtils.java

@@ -22,6 +22,7 @@
 import io.trino.connector.CatalogHandle;
 import io.trino.connector.CatalogServiceProvider;
 import io.trino.cost.StatsAndCosts;
+import io.trino.exchange.ExchangeManagerConfig;
 import io.trino.exchange.ExchangeManagerRegistry;
 import io.trino.execution.BaseTestSqlTaskManager.MockDirectExchangeClientSupplier;
 import io.trino.execution.buffer.OutputBuffers;
@@ -180,7 +181,7 @@ public static LocalExecutionPlanner createTestingPlanner()
                 blockTypeOperators,
                 PLANNER_CONTEXT.getTypeOperators(),
                 new TableExecuteContextManager(),
-                new ExchangeManagerRegistry(noop(), noopTracer(), new SecretsResolver(ImmutableMap.of())),
+                new ExchangeManagerRegistry(noop(), noopTracer(), new SecretsResolver(ImmutableMap.of()), new ExchangeManagerConfig()),
                 new NodeVersion("test"),
                 new CompilerConfig());
     }