perf: avoid rebuilding the eligible-worker snapshot per pending task in HttpRemoteTaskRunner

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

@@ -46,6 +46,7 @@
 
 import java.net.URL;
 import java.util.ArrayList;
+import java.util.Collections;
 import java.util.List;
 import java.util.Map;
 import java.util.concurrent.ConcurrentHashMap;
@@ -77,6 +78,8 @@ public class WorkerHolder
   // Known list of tasks running/completed on this worker.
   protected final AtomicReference<Map<String, TaskAnnouncement>> tasksSnapshotRef;
 
+  private final AtomicReference<ImmutableWorkerInfo> announcementsViewRef = new AtomicReference<>();
+
   private final AtomicReference<DateTime> lastCompletedTaskTime = new AtomicReference<>(DateTimes.nowUtc());
   private final AtomicReference<DateTime> blacklistedUntil = new AtomicReference<>();
   private final AtomicInteger continuouslyFailedTasksCount = new AtomicInteger(0);
@@ -125,6 +128,7 @@ public WorkerHolder(
       knownAnnouncements.forEach(e -> announcements.put(e.getTaskId(), e));
     }
     tasksSnapshotRef = new AtomicReference<>(announcements);
+    refreshAnnouncementsView();
   }
 
   public Worker getWorker()
@@ -149,12 +153,39 @@ public void setBlacklistedUntil(DateTime blacklistedUntil)
 
   public ImmutableWorkerInfo toImmutable()
   {
-    return ImmutableWorkerInfo.fromWorkerAnnouncements(
+    final ImmutableWorkerInfo announcementsView = announcementsViewRef.get();
+    return new ImmutableWorkerInfo(
         getWorker(),
+        announcementsView.getCurrCapacityUsed(),
+        announcementsView.getCurrParallelIndexCapacityUsed(),
+        announcementsView.getCurrCapacityUsedByTaskType(),
+        announcementsView.getAvailabilityGroups(),
+        announcementsView.getRunningTasks(),
+        lastCompletedTaskTime.get(),
+        blacklistedUntil.get()
+    );
+  }
+
+  protected void refreshAnnouncementsView()
+  {
+    final ImmutableWorkerInfo view = ImmutableWorkerInfo.fromWorkerAnnouncements(
+        enabledWorker,
         tasksSnapshotRef.get(),
         lastCompletedTaskTime.get(),
         blacklistedUntil.get()
     );
+    announcementsViewRef.set(
+        new ImmutableWorkerInfo(
+            view.getWorker(),
+            view.getCurrCapacityUsed(),
+            view.getCurrParallelIndexCapacityUsed(),
+            Collections.unmodifiableMap(view.getCurrCapacityUsedByTaskType()),
+            view.getAvailabilityGroups(),
+            view.getRunningTasks(),
+            view.getLastCompletedTaskTime(),
+            view.getBlacklistedUntil()
+        )
+    );
   }
 
   public int getContinuouslyFailedTasksCount()
@@ -368,6 +399,7 @@ public void fullSync(List<WorkerHistoryItem> changes)
         }
 
         tasksSnapshotRef.set(newSnapshot);
+        refreshAnnouncementsView();
 
         notifyListener(delta, isWorkerDisabled);
       }
@@ -419,6 +451,8 @@ public void deltaSync(List<WorkerHistoryItem> changes)
           }
         }
 
+        refreshAnnouncementsView();
+
         notifyListener(delta, isWorkerDisabled);
       }
 

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

@@ -2165,6 +2165,7 @@ public void start()
                   () -> {
                     for (TaskAnnouncement announcement : preExistingTaskAnnouncements) {
                       tasksSnapshotRef.get().put(announcement.getTaskId(), announcement);
+                      refreshAnnouncementsView();
                       listener.taskAddedOrUpdated(announcement, this);
                     }
                     ticks.incrementAndGet();
@@ -2261,6 +2262,7 @@ public boolean assignTask(Task task)
                 }
 
                 tasksSnapshotRef.get().put(announcement.getTaskId(), announcement);
+                refreshAnnouncementsView();
                 listener.taskAddedOrUpdated(announcement, this);
               }
 

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

@@ -20,6 +20,8 @@
 package org.apache.druid.indexing.overlord.hrtr;
 
 import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.ImmutableSet;
 import org.apache.druid.indexer.TaskLocation;
 import org.apache.druid.indexer.TaskStatus;
 import org.apache.druid.indexing.common.task.NoopTask;
@@ -29,15 +31,19 @@
 import org.apache.druid.indexing.worker.Worker;
 import org.apache.druid.indexing.worker.WorkerHistoryItem;
 import org.apache.druid.indexing.worker.config.WorkerConfig;
+import org.apache.druid.java.util.common.DateTimes;
 import org.apache.druid.java.util.http.client.HttpClient;
 import org.apache.druid.segment.TestHelper;
 import org.apache.druid.server.coordination.ChangeRequestHttpSyncer;
 import org.easymock.EasyMock;
+import org.joda.time.DateTime;
 import org.junit.jupiter.api.Assertions;
 import org.junit.jupiter.api.Test;
 
 import java.util.ArrayList;
 import java.util.List;
+import java.util.Map;
+import java.util.Set;
 import java.util.concurrent.ScheduledExecutorService;
 
 /**
@@ -206,4 +212,134 @@ public void stateChanged(boolean enabled, WorkerHolder workerHolder)
 
     updates.clear();
   }
+
+  @Test
+  public void testToImmutableIsUpdatedBeforeListenerIsNotified()
+  {
+    final Task task0 = NoopTask.create();
+    final Task task1 = NoopTask.create();
+    final RunningTasksRecorder recorder = new RunningTasksRecorder();
+    final WorkerHolder workerHolder = createWorkerHolder(recorder, running(task0));
+    final ChangeRequestHttpSyncer.Listener<WorkerHistoryItem> syncListener = workerHolder.createSyncListener();
+
+    // Known announcements are reflected before any sync
+    Assertions.assertEquals(ImmutableSet.of(task0.getId()), workerHolder.toImmutable().getRunningTasks());
+
+    syncListener.fullSync(
+        ImmutableList.of(new WorkerHistoryItem.Metadata(false), update(running(task0)), update(running(task1)))
+    );
+    final Set<String> bothTasks = ImmutableSet.of(task0.getId(), task1.getId());
+    Assertions.assertEquals(ImmutableList.of(bothTasks, bothTasks), recorder.drain());
+    Assertions.assertFalse(workerHolder.toImmutable().getWorker().isDisabled());
+
+    syncListener.deltaSync(ImmutableList.of(update(succeeded(task0))));
+    Assertions.assertEquals(ImmutableList.of(ImmutableSet.of(task1.getId())), recorder.drain());
+
+    syncListener.deltaSync(ImmutableList.of(new WorkerHistoryItem.TaskRemoval(task1.getId())));
+    Assertions.assertEquals(ImmutableList.of(ImmutableSet.of()), recorder.drain());
+    Assertions.assertEquals(0, workerHolder.toImmutable().getCurrCapacityUsed());
+  }
+
+  @Test
+  public void testToImmutableReflectsFieldsNotDerivedFromAnnouncementsWithoutSync()
+  {
+    final WorkerHolder workerHolder = createWorkerHolder(new RunningTasksRecorder());
+    final DateTime lastCompletedTaskTime = DateTimes.of("2026-01-01");
+    final DateTime blacklistedUntil = DateTimes.of("2026-01-02");
+
+    workerHolder.setLastCompletedTaskTime(lastCompletedTaskTime);
+    workerHolder.setBlacklistedUntil(blacklistedUntil);
+
+    Assertions.assertEquals(lastCompletedTaskTime, workerHolder.toImmutable().getLastCompletedTaskTime());
+    Assertions.assertEquals(blacklistedUntil, workerHolder.toImmutable().getBlacklistedUntil());
+    // Disabled until the first sync enables it
+    Assertions.assertTrue(workerHolder.toImmutable().getWorker().isDisabled());
+  }
+
+  @Test
+  public void testWorkerRejoiningWithPlaceholderAnnouncements()
+  {
+    final String taskId = "task1";
+    final TaskAnnouncement placeholder = TaskAnnouncement.create(
+        taskId,
+        NoopTask.TYPE,
+        null,
+        TaskStatus.running(taskId),
+        TaskLocation.unknown(),
+        null
+    );
+    final WorkerHolder workerHolder = createWorkerHolder(new RunningTasksRecorder(), placeholder);
+
+    Assertions.assertEquals(ImmutableSet.of(taskId), workerHolder.toImmutable().getRunningTasks());
+    Assertions.assertEquals(1, workerHolder.toImmutable().getCurrCapacityUsed());
+  }
+
+  @Test
+  public void testToImmutableDoesNotExposeCachedCapacityByTaskType()
+  {
+    final Task task = NoopTask.create();
+    final WorkerHolder workerHolder = createWorkerHolder(new RunningTasksRecorder(), running(task));
+    final Map<String, Integer> capacityByTaskType = workerHolder.toImmutable().getCurrCapacityUsedByTaskType();
+
+    Assertions.assertThrows(UnsupportedOperationException.class, () -> capacityByTaskType.put(task.getType(), 100));
+    Assertions.assertEquals(
+        ImmutableMap.of(task.getType(), 1),
+        workerHolder.toImmutable().getCurrCapacityUsedByTaskType()
+    );
+  }
+
+  private static WorkerHolder createWorkerHolder(WorkerHolder.Listener listener, TaskAnnouncement... knownAnnouncements)
+  {
+    return new WorkerHolder(
+        TestHelper.makeJsonMapper(),
+        EasyMock.createNiceMock(HttpClient.class),
+        new HttpRemoteTaskRunnerConfig(),
+        EasyMock.createNiceMock(ScheduledExecutorService.class),
+        listener,
+        new Worker("http", "localhost", "127.0.0.1", 5, "v1", WorkerConfig.DEFAULT_CATEGORY),
+        ImmutableList.copyOf(knownAnnouncements)
+    );
+  }
+
+  private static TaskAnnouncement running(Task task)
+  {
+    return TaskAnnouncement.create(task, TaskStatus.running(task.getId()), TaskLocation.unknown());
+  }
+
+  private static TaskAnnouncement succeeded(Task task)
+  {
+    return TaskAnnouncement.create(task, TaskStatus.success(task.getId()), TaskLocation.unknown());
+  }
+
+  private static WorkerHistoryItem update(TaskAnnouncement announcement)
+  {
+    return new WorkerHistoryItem.TaskUpdate(announcement);
+  }
+
+  /**
+   * Records the running tasks seen through {@link WorkerHolder#toImmutable()} from inside each listener callback.
+   */
+  private static class RunningTasksRecorder implements WorkerHolder.Listener
+  {
+    private final List<Set<String>> seen = new ArrayList<>();
+
+    @Override
+    public void taskAddedOrUpdated(TaskAnnouncement announcement, WorkerHolder workerHolder)
+    {
+      seen.add(workerHolder.toImmutable().getRunningTasks());
+    }
+
+    @Override
+    public void stateChanged(boolean enabled, WorkerHolder workerHolder)
+    {
+      // Do nothing
+    }
+
+    List<Set<String>> drain()
+    {
+      final List<Set<String>> drained = new ArrayList<>(seen);
+      seen.clear();
+      return drained;
+    }
+  }
 }