Fetch active task payloads from memory (#15377)

The TaskQueue maintains a map of active task ids to tasks, which can be utilized to get active task payloads, before falling back to the metadata store.

indexing-service/src/main/java/org/apache/druid/indexing/overlord/TaskQueue.java

@@ -947,6 +947,17 @@ public class TaskQueue
     return stats;
   }
 
+  public Optional<Task> getActiveTask(String id)
+  {
+    giant.lock();
+    try {
+      return Optional.fromNullable(tasks.get(id));
+    }
+    finally {
+      giant.unlock();
+    }
+  }
+
   @VisibleForTesting
   List<Task> getTasks()
   {

indexing-service/src/main/java/org/apache/druid/indexing/overlord/TaskStorageQueryAdapter.java

@@ -48,12 +48,14 @@ public class TaskStorageQueryAdapter
 {
   private final TaskStorage storage;
   private final TaskLockbox taskLockbox;
+  private final Optional<TaskQueue> taskQueue;
 
   @Inject
-  public TaskStorageQueryAdapter(TaskStorage storage, TaskLockbox taskLockbox)
+  public TaskStorageQueryAdapter(TaskStorage storage, TaskLockbox taskLockbox, TaskMaster taskMaster)
   {
     this.storage = storage;
     this.taskLockbox = taskLockbox;
+    this.taskQueue = taskMaster.getTaskQueue();
   }
 
   public List<Task> getActiveTasks()
@@ -104,6 +106,12 @@ public class TaskStorageQueryAdapter
 
   public Optional<Task> getTask(final String taskid)
   {
+    if (taskQueue.isPresent()) {
+      Optional<Task> activeTask = taskQueue.get().getActiveTask(taskid);
+      if (activeTask.isPresent()) {
+        return activeTask;
+      }
+    }
     return storage.getTask(taskid);
   }
 

indexing-service/src/test/java/org/apache/druid/indexing/overlord/TaskLifecycleTest.java

@@ -482,7 +482,10 @@ public class TaskLifecycleTest extends InitializedNullHandlingTest
       default:
         throw new RE("Unknown task storage type [%s]", taskStorageType);
     }
-    tsqa = new TaskStorageQueryAdapter(taskStorage, taskLockbox);
+    TaskMaster taskMaster = EasyMock.createMock(TaskMaster.class);
+    EasyMock.expect(taskMaster.getTaskQueue()).andReturn(Optional.absent()).anyTimes();
+    EasyMock.replay(taskMaster);
+    tsqa = new TaskStorageQueryAdapter(taskStorage, taskLockbox, taskMaster);
     return taskStorage;
   }
 

indexing-service/src/test/java/org/apache/druid/indexing/overlord/http/OverlordTest.java

@@ -257,7 +257,7 @@ public class OverlordTest
     Assert.assertEquals(taskMaster.getCurrentLeader(), druidNode.getHostAndPort());
     Assert.assertEquals(Optional.absent(), taskMaster.getRedirectLocation());
 
-    final TaskStorageQueryAdapter taskStorageQueryAdapter = new TaskStorageQueryAdapter(taskStorage, taskLockbox);
+    final TaskStorageQueryAdapter taskStorageQueryAdapter = new TaskStorageQueryAdapter(taskStorage, taskLockbox, taskMaster);
     final WorkerTaskRunnerQueryAdapter workerTaskRunnerQueryAdapter = new WorkerTaskRunnerQueryAdapter(taskMaster, null);
     // Test Overlord resource stuff
     overlordResource = new OverlordResource(