Add TmpStorageProvider

apache · Dec 18, 2024 · 2864ed2 · 2864ed2
1 parent d5eb94d
commit 2864ed2
Show file tree

Hide file tree

Showing 9 changed files with 101 additions and 0 deletions.
diff --git a/processing/src/main/java/org/apache/druid/storage/local/LocalTmpStorage.java b/processing/src/main/java/org/apache/druid/storage/local/LocalTmpStorage.java
@@ -0,0 +1,52 @@
+/*
+ * 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.druid.storage.local;
+
+import com.google.inject.Provider;
+import org.apache.druid.java.util.common.FileUtils;
+
+import java.io.File;
+
+public interface LocalTmpStorage
+{
+  /**
+   * Get a temporary directory.
+   *
+   * @return a temporary directory
+  */
+  File getTmpDir();
+
+  class DefaultLocalTmpStorageProvider implements Provider<LocalTmpStorage>
+  {
+    private final String prefix;
+
+    public DefaultLocalTmpStorageProvider(String prefix)
+    {
+      this.prefix = prefix;
+    }
+
+    @Override
+    public LocalTmpStorage get()
+    {
+      File tmpStorage = FileUtils.createTempDir(prefix);
+      return () -> tmpStorage;
+    }
+  }
+}
diff --git a/services/src/main/java/org/apache/druid/cli/CliBroker.java b/services/src/main/java/org/apache/druid/cli/CliBroker.java
@@ -76,6 +76,7 @@
 import org.apache.druid.server.router.TieredBrokerConfig;
 import org.apache.druid.sql.calcite.schema.MetadataSegmentView;
 import org.apache.druid.sql.guice.SqlModule;
+import org.apache.druid.storage.local.LocalTmpStorage;
 import org.apache.druid.timeline.PruneLoadSpec;
 import org.eclipse.jetty.server.Server;
 
@@ -182,6 +183,10 @@ protected List<? extends Module> getModules()
 
           Jerseys.addResource(binder, SelfDiscoveryResource.class);
           LifecycleModule.registerKey(binder, Key.get(SelfDiscoveryResource.class));
+
+          binder.bind(LocalTmpStorage.class)
+                .toProvider(new LocalTmpStorage.DefaultLocalTmpStorageProvider("broker"))
+                .in(LazySingleton.class);
         },
         new LookupModule(),
         new SqlModule()

diff --git a/services/src/main/java/org/apache/druid/cli/CliCoordinator.java b/services/src/main/java/org/apache/druid/cli/CliCoordinator.java
@@ -130,6 +130,7 @@
 import org.apache.druid.server.lookup.cache.LookupCoordinatorManagerConfig;
 import org.apache.druid.server.metrics.ServiceStatusMonitor;
 import org.apache.druid.server.router.TieredBrokerConfig;
+import org.apache.druid.storage.local.LocalTmpStorage;
 import org.eclipse.jetty.server.Server;
 import org.joda.time.Duration;
 
@@ -293,6 +294,10 @@ public void configure(Binder binder)
             binder.bind(CoordinatorCustomDutyGroups.class)
                   .toProvider(new CoordinatorCustomDutyGroupsProvider())
                   .in(LazySingleton.class);
+
+            binder.bind(LocalTmpStorage.class)
+                  .toProvider(new LocalTmpStorage.DefaultLocalTmpStorageProvider("coordinator"))
+                  .in(LazySingleton.class);
           }
 
           @Provides

diff --git a/services/src/main/java/org/apache/druid/cli/CliHistorical.java b/services/src/main/java/org/apache/druid/cli/CliHistorical.java
@@ -59,6 +59,7 @@
 import org.apache.druid.server.http.SelfDiscoveryResource;
 import org.apache.druid.server.initialization.jetty.JettyServerInitializer;
 import org.apache.druid.server.metrics.QueryCountStatsProvider;
+import org.apache.druid.storage.local.LocalTmpStorage;
 import org.apache.druid.timeline.PruneLastCompactionState;
 import org.eclipse.jetty.server.Server;
 
@@ -140,6 +141,10 @@ protected List<? extends Module> getModules()
 
           Jerseys.addResource(binder, SelfDiscoveryResource.class);
           LifecycleModule.registerKey(binder, Key.get(SelfDiscoveryResource.class));
+
+          binder.bind(LocalTmpStorage.class)
+                .toProvider(new LocalTmpStorage.DefaultLocalTmpStorageProvider("historical"))
+                .in(LazySingleton.class);
         },
         new LookupModule()
     );

diff --git a/services/src/main/java/org/apache/druid/cli/CliIndexer.java b/services/src/main/java/org/apache/druid/cli/CliIndexer.java
@@ -82,6 +82,7 @@
 import org.apache.druid.server.initialization.jetty.CliIndexerServerModule;
 import org.apache.druid.server.initialization.jetty.JettyServerInitializer;
 import org.apache.druid.server.metrics.IndexerTaskCountStatsProvider;
+import org.apache.druid.storage.local.LocalTmpStorage;
 import org.eclipse.jetty.server.Server;
 
 import java.util.List;
@@ -196,6 +197,9 @@ public void configure(Binder binder)
 
             Jerseys.addResource(binder, SelfDiscoveryResource.class);
             LifecycleModule.registerKey(binder, Key.get(SelfDiscoveryResource.class));
+            binder.bind(LocalTmpStorage.class)
+                  .toProvider(new LocalTmpStorage.DefaultLocalTmpStorageProvider("indexer"))
+                  .in(LazySingleton.class);
           }
 
           @Provides

diff --git a/services/src/main/java/org/apache/druid/cli/CliMiddleManager.java b/services/src/main/java/org/apache/druid/cli/CliMiddleManager.java
@@ -81,6 +81,7 @@
 import org.apache.druid.server.initialization.jetty.JettyServerInitializer;
 import org.apache.druid.server.metrics.ServiceStatusMonitor;
 import org.apache.druid.server.metrics.WorkerTaskCountStatsProvider;
+import org.apache.druid.storage.local.LocalTmpStorage;
 import org.apache.druid.timeline.PruneLastCompactionState;
 import org.eclipse.jetty.server.Server;
 
@@ -184,6 +185,10 @@ public void configure(Binder binder)
             LifecycleModule.registerKey(binder, Key.get(SelfDiscoveryResource.class));
 
             configureIntermediaryData(binder);
+
+            binder.bind(LocalTmpStorage.class)
+                  .toProvider(new LocalTmpStorage.DefaultLocalTmpStorageProvider("middle-manager"))
+                  .in(LazySingleton.class);
           }
 
           private void configureIntermediaryData(Binder binder)

diff --git a/services/src/main/java/org/apache/druid/cli/CliOverlord.java b/services/src/main/java/org/apache/druid/cli/CliOverlord.java
@@ -132,6 +132,7 @@
 import org.apache.druid.server.security.AuthenticationUtils;
 import org.apache.druid.server.security.Authenticator;
 import org.apache.druid.server.security.AuthenticatorMapper;
+import org.apache.druid.storage.local.LocalTmpStorage;
 import org.apache.druid.tasklogs.TaskLogStreamer;
 import org.apache.druid.tasklogs.TaskLogs;
 import org.eclipse.jetty.rewrite.handler.RewriteHandler;
@@ -317,6 +318,10 @@ public void configure(Binder binder)
 
             Jerseys.addResource(binder, SelfDiscoveryResource.class);
             LifecycleModule.registerKey(binder, Key.get(SelfDiscoveryResource.class));
+
+            binder.bind(LocalTmpStorage.class)
+                  .toProvider(new LocalTmpStorage.DefaultLocalTmpStorageProvider("overlord"))
+                  .in(LazySingleton.class);
           }
 
           private void configureTaskStorage(Binder binder)

diff --git a/services/src/main/java/org/apache/druid/cli/CliPeon.java b/services/src/main/java/org/apache/druid/cli/CliPeon.java
@@ -134,6 +134,7 @@
 import org.apache.druid.server.lookup.cache.LookupLoadingSpec;
 import org.apache.druid.server.metrics.DataSourceTaskIdHolder;
 import org.apache.druid.server.metrics.ServiceStatusMonitor;
+import org.apache.druid.storage.local.LocalTmpStorage;
 import org.apache.druid.tasklogs.TaskPayloadManager;
 import org.eclipse.jetty.server.Server;
 
@@ -355,6 +356,20 @@ public BroadcastDatasourceLoadingSpec getBroadcastDatasourcesToLoad(final Task t
           {
             return task.getBroadcastDatasourceLoadingSpec();
           }
+
+          @Provides
+          @LazySingleton
+          public LocalTmpStorage getTmpStorageProvider()
+          {
+            Supplier<File> tmpDirSupplier = () -> {
+              File tmpDir = new File(taskDirPath, "tmp");
+              if (!tmpDir.mkdirs()) {
+                log.warn("Failed to create tmp directory [%s]", tmpDir);
+              }
+              return tmpDir;
+            };
+            return tmpDirSupplier::get;
+          }
         },
         new QueryablePeonModule(),
         new IndexingServiceInputSourceModule(),

diff --git a/services/src/main/java/org/apache/druid/cli/CliRouter.java b/services/src/main/java/org/apache/druid/cli/CliRouter.java
@@ -56,6 +56,7 @@
 import org.apache.druid.server.router.TieredBrokerHostSelector;
 import org.apache.druid.server.router.TieredBrokerSelectorStrategiesProvider;
 import org.apache.druid.server.router.TieredBrokerSelectorStrategy;
+import org.apache.druid.storage.local.LocalTmpStorage;
 import org.eclipse.jetty.server.Server;
 
 import java.util.List;
@@ -126,6 +127,10 @@ protected List<? extends Module> getModules()
 
           Jerseys.addResource(binder, SelfDiscoveryResource.class);
           LifecycleModule.registerKey(binder, Key.get(SelfDiscoveryResource.class));
+
+          binder.bind(LocalTmpStorage.class)
+                .toProvider(new LocalTmpStorage.DefaultLocalTmpStorageProvider("router"))
+                .in(LazySingleton.class);
         },
         new LookupSerdeModule()
     );