Fix Processing Partition picked up without active workers

Author: mohityadav766

openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/distributed/DistributedSearchIndexExecutor.java

@@ -23,6 +23,7 @@
 import java.util.Optional;
 import java.util.Set;
 import java.util.UUID;
+import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
@@ -100,6 +101,7 @@ public static boolean isCoordinatingJob(UUID jobId) {
   @Getter private SearchIndexJob currentJob;
   private DistributedJobStatsAggregator statsAggregator;
   private ExecutorService workerExecutor;
+  private final Set<UUID> activePartitions = ConcurrentHashMap.newKeySet();
   private final List<PartitionWorker> activeWorkers = new ArrayList<>();
   private Thread lockRefreshThread;
   private Thread partitionHeartbeatThread;
@@ -620,29 +622,41 @@ private void runWorkerLoop(
             partition.getId(),
             partition.getEntityType());
 
-        PartitionWorker.PartitionResult result = worker.processPartition(partition);
-        totalSuccess.addAndGet(result.successCount());
-        totalFailed.addAndGet(result.failedCount());
+        activePartitions.add(partition.getId());
+        try {
+          PartitionWorker.PartitionResult result = worker.processPartition(partition);
+          totalSuccess.addAndGet(result.successCount());
+          totalFailed.addAndGet(result.failedCount());
 
-        // Accumulate into coordinator stats for persistence
-        // readerSuccess = entities successfully processed through reader
-        // readerFailed = specifically reader failures (not sink failures)
-        // readerWarnings = warnings from reader (e.g., stale references)
-        coordinatorReaderSuccess.addAndGet(result.successCount());
-        coordinatorReaderFailed.addAndGet(result.readerFailed());
-        coordinatorReaderWarnings.addAndGet(result.readerWarnings());
-        coordinatorPartitionsCompleted.incrementAndGet();
+          coordinatorReaderSuccess.addAndGet(result.successCount());
+          coordinatorReaderFailed.addAndGet(result.readerFailed());
+          coordinatorReaderWarnings.addAndGet(result.readerWarnings());
+          coordinatorPartitionsCompleted.incrementAndGet();
 
-        LOG.info(
-            "Worker {} completed partition {} (success: {}, failed: {}, readerFailed: {}, readerWarnings: {})",
-            workerId,
-            partition.getId(),
-            result.successCount(),
-            result.failedCount(),
-            result.readerFailed(),
-            result.readerWarnings());
-
-        // Stats are tracked per-entityType by StageStatsTracker in PartitionWorker
+          LOG.info(
+              "Worker {} completed partition {} (success: {}, failed: {}, readerFailed: {}, readerWarnings: {})",
+              workerId,
+              partition.getId(),
+              result.successCount(),
+              result.failedCount(),
+              result.readerFailed(),
+              result.readerWarnings());
+        } catch (Exception e) {
+          LOG.error(
+              "Worker {} failed partition {} for {}: {}",
+              workerId,
+              partition.getId(),
+              partition.getEntityType(),
+              e.getMessage(),
+              e);
+          try {
+            coordinator.failPartition(partition.getId(), e.getMessage());
+          } catch (Exception fe) {
+            LOG.error("Failed to mark partition {} as failed", partition.getId(), fe);
+          }
+        } finally {
+          activePartitions.remove(partition.getId());
+        }
       }
     } finally {
       synchronized (activeWorkers) {
@@ -738,8 +752,8 @@ private void runPartitionHeartbeatLoop() {
         int updated = 0;
         long now = System.currentTimeMillis();
         for (SearchIndexPartition partition : processing) {
-          if (serverId.equals(partition.getAssignedServer())) {
-            // Update the heartbeat for this partition
+          if (serverId.equals(partition.getAssignedServer())
+              && activePartitions.contains(partition.getId())) {
             collectionDAO
                 .searchIndexPartitionDAO()
                 .updateHeartbeat(partition.getId().toString(), now);