HDDS-15997. Fix QuotaRepairTask scan hang and partial counts on worker failure

hadoop-ozone/ozone-manager/src/main/java/org/apache/hadoop/ozone/om/service/QuotaRepairTask.java

@@ -23,6 +23,7 @@
 import static org.apache.hadoop.ozone.OzoneConsts.OM_KEY_PREFIX;
 import static org.apache.hadoop.ozone.om.helpers.SnapshotInfo.SnapshotStatus.SNAPSHOT_ACTIVE;
 
+import com.google.common.annotations.VisibleForTesting;
 import com.google.common.util.concurrent.UncheckedExecutionException;
 import com.google.protobuf.ServiceException;
 import java.io.File;
@@ -49,6 +50,7 @@
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicLong;
+import java.util.function.Consumer;
 import org.apache.commons.io.FileUtils;
 import org.apache.hadoop.hdds.conf.OzoneConfiguration;
 import org.apache.hadoop.hdds.server.JsonUtils;
@@ -82,7 +84,8 @@
 public class QuotaRepairTask {
   private static final Logger LOG = LoggerFactory.getLogger(
       QuotaRepairTask.class);
-  private static final int BATCH_SIZE = 5000;
+  @VisibleForTesting
+  static final int BATCH_SIZE = 5000;
   private static final int TASK_THREAD_CNT = 3;
   /**
    * Parallel full-table scans: OBS keys, FSO files, dirs, active deleted keys/dirs,
@@ -363,8 +366,10 @@ private void repairCount(
         }
       }));
 
-      for (Future<?> f : tasks) {
-        f.get();
+      // await every scan before propagating a failure, so no scan outlives the checkpoint
+      Exception scanFailure = awaitAll(tasks, null);
+      if (scanFailure != null) {
+        throw scanFailure;
       }
     } catch (UncheckedIOException ex) {
       LOG.error("quota repair failure", ex.getCause());
@@ -576,6 +581,21 @@ private <VALUE> void recalculateUsages(
       Table<String, VALUE> table, Map<String, CountPair> prefixUsageMap,
       String strType, boolean haveValue) throws UncheckedIOException,
       UncheckedExecutionException {
+    try (Table.KeyValueIterator<String, VALUE> keyIter
+        = table.iterator(null, haveValue ? KEY_AND_VALUE : KEY_ONLY)) {
+      scanTableInBatches(executor, keyIter, strType,
+          kv -> extractCount(kv, prefixUsageMap, haveValue));
+    } catch (IOException ex) {
+      throw new UncheckedIOException(ex);
+    }
+  }
+
+  @VisibleForTesting
+  static <VALUE> void scanTableInBatches(
+      ExecutorService executor,
+      Table.KeyValueIterator<String, VALUE> keyIter, String strType,
+      Consumer<Table.KeyValue<String, VALUE>> kvConsumer)
+      throws UncheckedIOException, UncheckedExecutionException {
     LOG.info("Starting recalculate {}", strType);
 
     List<Table.KeyValue<String, VALUE>> kvList = new ArrayList<>(BATCH_SIZE);
@@ -584,53 +604,110 @@ private <VALUE> void recalculateUsages(
     List<Future<?>> tasks = new ArrayList<>();
     AtomicBoolean isRunning = new AtomicBoolean(true);
     for (int i = 0; i < TASK_THREAD_CNT; ++i) {
-      tasks.add(executor.submit(() -> captureCount(
-          prefixUsageMap, q, isRunning, haveValue)));
+      tasks.add(executor.submit(() -> captureCount(q, isRunning, kvConsumer)));
     }
     int count = 0;
     long startTime = Time.monotonicNow();
-    try (Table.KeyValueIterator<String, VALUE> keyIter
-        = table.iterator(null, haveValue ? KEY_AND_VALUE : KEY_ONLY)) {
+    Exception failure = null;
+    try {
       while (keyIter.hasNext()) {
         count++;
         kvList.add(keyIter.next());
         if (kvList.size() == BATCH_SIZE) {
-          q.put(kvList);
+          putBatch(q, kvList, tasks);
           kvList = new ArrayList<>(BATCH_SIZE);
         }
       }
-      q.put(kvList);
+      putBatch(q, kvList, tasks);
+    } catch (InterruptedException ex) {
+      Thread.currentThread().interrupt();
+      failure = ex;
+      q.clear();
+    } catch (ExecutionException | RuntimeException ex) {
+      failure = ex;
+      q.clear();
+    } finally {
       isRunning.set(false);
-      for (Future<?> f : tasks) {
-        f.get();
+    }
+    // always await workers so none outlives this scan and touches a closed table
+    failure = awaitAll(tasks, failure);
+    if (failure != null) {
+      throw new UncheckedExecutionException(failure);
+    }
+    LOG.info("Recalculate {} completed, count {} time {}ms", strType,
+        count, (Time.monotonicNow() - startTime));
+  }
+
+  /**
+   * Awaits every task, retrying an interrupted wait so the interrupted future is not
+   * abandoned mid-run; an interrupt seen while waiting is restored before returning.
+   * Returns the passed-in failure or the first failure seen, with later ones suppressed.
+   */
+  private static Exception awaitAll(List<Future<?>> tasks, Exception failure) {
+    boolean interrupted = false;
+    for (Future<?> f : tasks) {
+      boolean done = false;
+      while (!done) {
+        try {
+          f.get();
+          done = true;
+        } catch (InterruptedException ex) {
+          interrupted = true;
+          if (failure == null) {
+            failure = ex;
+          }
+        } catch (ExecutionException ex) {
+          done = true;
+          if (failure == null) {
+            failure = ex;
+          } else {
+            failure.addSuppressed(ex);
+          }
+        }
       }
-      LOG.info("Recalculate {} completed, count {} time {}ms", strType,
-          count, (Time.monotonicNow() - startTime));
-    } catch (IOException ex) {
-      throw new UncheckedIOException(ex);
-    } catch (InterruptedException ex) {
+    }
+    if (interrupted) {
       Thread.currentThread().interrupt();
-    } catch (ExecutionException ex) {
-      throw new UncheckedExecutionException(ex);
     }
+    return failure;
   }
-
+
+  /**
+   * Blocks until the batch is queued. Fails fast if a worker has already exited,
+   * otherwise a failed worker set could leave the producer blocked forever on a full queue.
+   */
+  private static <VALUE> void putBatch(
+      BlockingQueue<List<Table.KeyValue<String, VALUE>>> q,
+      List<Table.KeyValue<String, VALUE>> kvList,
+      List<Future<?>> tasks) throws InterruptedException, ExecutionException {
+    while (!q.offer(kvList, 100, TimeUnit.MILLISECONDS)) {
+      for (Future<?> f : tasks) {
+        if (f.isDone()) {
+          f.get();
+          throw new IllegalStateException("quota repair scan worker exited prematurely");
+        }
+      }
+    }
+  }
+
   private static <VALUE> void captureCount(
-      Map<String, CountPair> prefixUsageMap,
       BlockingQueue<List<Table.KeyValue<String, VALUE>>> q,
-      AtomicBoolean isRunning, boolean haveValue) throws UncheckedIOException {
+      AtomicBoolean isRunning,
+      Consumer<Table.KeyValue<String, VALUE>> kvConsumer) throws UncheckedIOException {
     try {
       while (isRunning.get() || !q.isEmpty()) {
         List<Table.KeyValue<String, VALUE>> kvList
             = q.poll(100, TimeUnit.MILLISECONDS);
         if (null != kvList) {
           for (Table.KeyValue<String, VALUE> kv : kvList) {
-            extractCount(kv, prefixUsageMap, haveValue);
+            kvConsumer.accept(kv);
           }
         }
       }
     } catch (InterruptedException ex) {
       Thread.currentThread().interrupt();
+      // fail the scan instead of returning partial counts as success
+      throw new IllegalStateException("quota repair scan worker interrupted", ex);
     }
   }
 

hadoop-ozone/ozone-manager/src/test/java/org/apache/hadoop/ozone/om/service/TestQuotaRepairTask.java

@@ -20,20 +20,30 @@
 import static org.apache.hadoop.hdds.protocol.proto.HddsProtos.ReplicationFactor.ONE;
 import static org.apache.hadoop.hdds.protocol.proto.HddsProtos.ReplicationFactor.THREE;
 import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertInstanceOf;
 import static org.junit.jupiter.api.Assertions.assertNull;
+import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.anyLong;
 import static org.mockito.Mockito.doAnswer;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.when;
 
+import com.google.common.util.concurrent.UncheckedExecutionException;
 import java.io.IOException;
+import java.io.UncheckedIOException;
 import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.atomic.AtomicReference;
 import org.apache.hadoop.hdds.client.RatisReplicationConfig;
 import org.apache.hadoop.hdds.utils.db.BatchOperation;
+import org.apache.hadoop.hdds.utils.db.Table;
 import org.apache.hadoop.hdds.utils.db.cache.CacheKey;
 import org.apache.hadoop.hdds.utils.db.cache.CacheValue;
 import org.apache.hadoop.ozone.om.helpers.BucketLayout;
@@ -336,4 +346,107 @@ private void zeroOutBucketUsedBytes(String volumeName, String bucketName,
             CacheValue.get(trxnLogIndex, bucketInfo));
     omMetadataManager.getBucketTable().put(dbKey, bucketInfo);
   }
+
+  @Test
+  public void testScanTableInBatchesFailsFastOnWorkerFailure() throws Exception {
+    int totalEntries = QuotaRepairTask.BATCH_SIZE * 8;
+    AtomicInteger remaining = new AtomicInteger(totalEntries);
+    @SuppressWarnings("unchecked")
+    Table.KeyValueIterator<String, OmKeyInfo> keyIter = mock(Table.KeyValueIterator.class);
+    when(keyIter.hasNext()).thenAnswer(inv -> remaining.get() > 0);
+    when(keyIter.next()).thenAnswer(inv -> {
+      remaining.decrementAndGet();
+      return Table.newKeyValue("/vol/bucket/key", null);
+    });
+    ExecutorService executor = Executors.newFixedThreadPool(3);
+    try {
+      UncheckedExecutionException ex = assertThrows(UncheckedExecutionException.class,
+          () -> QuotaRepairTask.scanTableInBatches(executor, keyIter, "worker failure test", kv -> {
+            throw new UncheckedIOException(new IOException("injected worker failure"));
+          }));
+      assertInstanceOf(UncheckedIOException.class, ex.getCause().getCause());
+      // no worker may outlive the scan
+      executor.shutdown();
+      assertTrue(executor.awaitTermination(10, TimeUnit.SECONDS));
+    } finally {
+      executor.shutdownNow();
+    }
+  }
+
+  @Test
+  public void testScanTableInBatchesFailsOnProducerInterrupt() throws Exception {
+    @SuppressWarnings("unchecked")
+    Table.KeyValueIterator<String, OmKeyInfo> keyIter = mock(Table.KeyValueIterator.class);
+    when(keyIter.hasNext()).thenReturn(true);
+    when(keyIter.next()).thenAnswer(inv -> Table.newKeyValue("/vol/bucket/key", null));
+    ExecutorService executor = Executors.newFixedThreadPool(3);
+    AtomicReference<Throwable> thrown = new AtomicReference<>();
+    Thread producer = new Thread(() -> {
+      try {
+        QuotaRepairTask.scanTableInBatches(executor, keyIter, "interrupt test", kv -> { });
+      } catch (Throwable t) {
+        thrown.set(t);
+      }
+    });
+    try {
+      producer.start();
+      producer.interrupt();
+      producer.join(TimeUnit.SECONDS.toMillis(60));
+      assertFalse(producer.isAlive());
+      UncheckedExecutionException ex = assertInstanceOf(UncheckedExecutionException.class, thrown.get());
+      assertInstanceOf(InterruptedException.class, ex.getCause());
+      executor.shutdown();
+      assertTrue(executor.awaitTermination(10, TimeUnit.SECONDS));
+    } finally {
+      executor.shutdownNow();
+    }
+  }
+
+  @Test
+  public void testInterruptedScanStillAwaitsWorkers() throws Exception {
+    AtomicInteger remaining = new AtomicInteger(QuotaRepairTask.BATCH_SIZE);
+    @SuppressWarnings("unchecked")
+    Table.KeyValueIterator<String, OmKeyInfo> keyIter = mock(Table.KeyValueIterator.class);
+    when(keyIter.hasNext()).thenAnswer(inv -> remaining.get() > 0);
+    when(keyIter.next()).thenAnswer(inv -> {
+      remaining.decrementAndGet();
+      return Table.newKeyValue("/vol/bucket/key", null);
+    });
+    CountDownLatch workerStarted = new CountDownLatch(1);
+    CountDownLatch releaseWorker = new CountDownLatch(1);
+    ExecutorService executor = Executors.newFixedThreadPool(3);
+    AtomicReference<Throwable> thrown = new AtomicReference<>();
+    Thread producer = new Thread(() -> {
+      try {
+        QuotaRepairTask.scanTableInBatches(executor, keyIter, "await workers test", kv -> {
+          workerStarted.countDown();
+          try {
+            releaseWorker.await();
+          } catch (InterruptedException ex) {
+            throw new IllegalStateException(ex);
+          }
+        });
+      } catch (Throwable t) {
+        thrown.set(t);
+      }
+    });
+    try {
+      producer.start();
+      assertTrue(workerStarted.await(30, TimeUnit.SECONDS));
+      producer.interrupt();
+      // the interrupted producer must keep waiting for the blocked worker instead of exiting
+      producer.join(TimeUnit.SECONDS.toMillis(1));
+      assertTrue(producer.isAlive());
+      releaseWorker.countDown();
+      producer.join(TimeUnit.SECONDS.toMillis(60));
+      assertFalse(producer.isAlive());
+      UncheckedExecutionException ex = assertInstanceOf(UncheckedExecutionException.class, thrown.get());
+      assertInstanceOf(InterruptedException.class, ex.getCause());
+      executor.shutdown();
+      assertTrue(executor.awaitTermination(10, TimeUnit.SECONDS));
+    } finally {
+      releaseWorker.countDown();
+      executor.shutdownNow();
+    }
+  }
 }