AmatyaAvadhanula commented on code in PR #14063:
URL: https://github.com/apache/druid/pull/14063#discussion_r1164072588
##########
indexing-service/src/main/java/org/apache/druid/indexing/common/TaskStorageDirTracker.java:
##########
@@ -19,81 +19,94 @@
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 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;
+ private final AtomicInteger iterationCounter = new AtomicInteger(0);
- 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()) {
+ return baseTaskDir;
+ }
}
- return taskToTempDirMap.get(taskId);
+ // if it doesn't exist, pick one round-robin and return. This will roll
negative, but that's okay because we
+ // are always modding it.
+ final int currIncrement = iterationCounter.getAndIncrement() %
baseTaskDirs.length;
+ return baseTaskDirs[currIncrement % baseTaskDirs.length];
Review Comment:
negative number modulo positive number is negative in java.
Even though overflow is unlikely to happen, it might be nice to add a check
here
--
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]
