github-code-scanning[bot] commented on code in PR #14063:
URL: https://github.com/apache/druid/pull/14063#discussion_r1164721389


##########
indexing-service/src/main/java/org/apache/druid/indexing/common/TaskStorageDirTracker.java:
##########
@@ -19,81 +19,96 @@
 
 package org.apache.druid.indexing.common;
 
-import com.google.common.annotations.VisibleForTesting;
+import com.google.common.collect.ImmutableList;
 import org.apache.druid.indexing.common.config.TaskConfig;
+import org.apache.druid.indexing.worker.config.WorkerConfig;
+import org.apache.druid.java.util.common.FileUtils;
 import org.apache.druid.java.util.common.ISE;
+import org.apache.druid.java.util.common.lifecycle.LifecycleStart;
 
-import javax.inject.Inject;
+import javax.annotation.Nullable;
 import java.io.File;
-import java.util.ArrayList;
-import java.util.HashMap;
+import java.io.IOException;
 import java.util.List;
-import java.util.Map;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.stream.Collectors;
 
 public class TaskStorageDirTracker
 {
-  private int taskDirIndex = 0;
-
-  private final List<File> baseTaskDirs = new ArrayList<>();
-
-  private final Map<String, File> taskToTempDirMap = new HashMap<>();
-
-  @Inject
-  public TaskStorageDirTracker(final TaskConfig taskConfig)
-  {
-    this(taskConfig.getBaseTaskDirPaths());
-  }
-
-  @VisibleForTesting
-  public TaskStorageDirTracker(final List<String> baseTaskDirPaths)
+  public static TaskStorageDirTracker fromConfigs(WorkerConfig workerConfig, 
TaskConfig taskConfig)
   {
-    for (String baseTaskDirPath : baseTaskDirPaths) {
-      baseTaskDirs.add(new File(baseTaskDirPath));
+    if (workerConfig == null) {
+      return new 
TaskStorageDirTracker(ImmutableList.of(taskConfig.getBaseTaskDir()));
+    } else {
+      final List<String> basePaths = workerConfig.getBaseTaskDirs();
+      if (basePaths == null) {
+        return new 
TaskStorageDirTracker(ImmutableList.of(taskConfig.getBaseTaskDir()));
+      }
+      return new TaskStorageDirTracker(
+          basePaths.stream().map(File::new).collect(Collectors.toList())
+      );
     }
   }
 
-  public File getTaskDir(String taskId)
-  {
-    return new File(getBaseTaskDir(taskId), taskId);
-  }
+  private final File[] baseTaskDirs;
+  // Initialize to a negative number because it ensures that we can handle the 
overflow-rollover case
+  private final AtomicInteger iterationCounter = new 
AtomicInteger(Integer.MIN_VALUE);
 
-  public File getTaskWorkDir(String taskId)
+  public TaskStorageDirTracker(List<File> baseTaskDirs)
   {
-    return new File(getTaskDir(taskId), "work");
+    this.baseTaskDirs = baseTaskDirs.toArray(new File[]{});
   }
 
-  public File getTaskTempDir(String taskId)
+  @LifecycleStart
+  public void ensureDirectories()
   {
-    return new File(getTaskDir(taskId), "temp");
+    for (File baseTaskDir : baseTaskDirs) {
+      if (!baseTaskDir.exists()) {
+        try {
+          FileUtils.mkdirp(baseTaskDir);
+        }
+        catch (IOException e) {
+          throw new ISE(
+              e,
+              "base task directory [%s] likely does not exist, please ensure 
it exists and the user has permissions.",
+              baseTaskDir
+          );
+        }
+      }
+    }
   }
 
-  public List<File> getBaseTaskDirs()
+  public File pickBaseDir(String taskId) throws IOException
   {
-    return baseTaskDirs;
-  }
+    if (baseTaskDirs.length == 1) {
+      return baseTaskDirs[0];
+    }
 
-  public synchronized File getBaseTaskDir(final String taskId)
-  {
-    if (!taskToTempDirMap.containsKey(taskId)) {
-      addTask(taskId, baseTaskDirs.get(taskDirIndex));
-      taskDirIndex = (taskDirIndex + 1) % baseTaskDirs.size();
+    // if the task directory already exists, we want to give it precedence, so 
check.
+    for (File baseTaskDir : baseTaskDirs) {
+      if (new File(baseTaskDir, taskId).exists()) {

Review Comment:
   ## Uncontrolled data used in path expression
   
   This path depends on a [user-provided value](1).
   
   [Show more 
details](https://github.com/apache/druid/security/code-scanning/4790)



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: [email protected]

For queries about this service, please contact Infrastructure at:
[email protected]


---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]

Reply via email to