HDDS-13639. Optimize container iterator for frequent operation

hadoop-hdds/container-service/src/main/java/org/apache/hadoop/ozone/container/common/impl/ContainerSet.java

@@ -205,6 +205,10 @@ private boolean addContainer(Container<?> container, boolean overwrite) throws
         recoveringContainerMap.put(
             clock.millis() + recoveringTimeout, containerId);
       }
+      HddsVolume volume = container.getContainerData().getVolume();
+      if (volume != null) {
+        volume.addContainer(containerId);
+      }
       return true;
     } else {
       LOG.warn("Container already exists with container Id {}", containerId);
@@ -299,6 +303,10 @@ private boolean removeContainer(long containerId, boolean markMissing, boolean r
           "containerMap", containerId);
       return false;
     } else {
+      HddsVolume volume = removed.getContainerData().getVolume();
+      if (volume != null) {
+        volume.removeContainer(containerId);
+      }
       LOG.debug("Container with containerId {} is removed from containerMap",
           containerId);
       return true;
@@ -409,13 +417,19 @@ public Iterator<Map.Entry<Long, Long>> getRecoveringContainerIterator() {
    */
   public Iterator<Container<?>> getContainerIterator(HddsVolume volume) {
     Preconditions.checkNotNull(volume);
-    Preconditions.checkNotNull(volume.getStorageID());
-    String volumeUuid = volume.getStorageID();
-    return containerMap.values().stream()
-        .filter(x -> volumeUuid.equals(x.getContainerData().getVolume()
-            .getStorageID()))
-        .sorted(ContainerDataScanOrder.INSTANCE)
-        .iterator();
+    Iterator<Long> containerIdIterator = volume.getContainerIterator();
+
+    List<Container<?>> containers = new ArrayList<>();
+    while (containerIdIterator.hasNext()) {
+      Long containerId = containerIdIterator.next();
+      Container<?> container = containerMap.get(containerId);
+      if (container != null) {
+        containers.add(container);
+      }
+    }
+    containers.sort(ContainerDataScanOrder.INSTANCE);
+
+    return containers.iterator();
   }
 
   /**
@@ -426,11 +440,7 @@ public Iterator<Container<?>> getContainerIterator(HddsVolume volume) {
    */
   public long containerCount(HddsVolume volume) {
     Preconditions.checkNotNull(volume);
-    Preconditions.checkNotNull(volume.getStorageID());
-    String volumeUuid = volume.getStorageID();
-    return containerMap.values().stream()
-        .filter(x -> volumeUuid.equals(x.getContainerData().getVolume()
-        .getStorageID())).count();
+    return volume.getContainerCount();
   }
 
   /**

hadoop-hdds/container-service/src/main/java/org/apache/hadoop/ozone/container/common/volume/HddsVolume.java

@@ -25,9 +25,11 @@
 import jakarta.annotation.Nullable;
 import java.io.File;
 import java.io.IOException;
+import java.util.Iterator;
 import java.util.LinkedList;
 import java.util.List;
 import java.util.Queue;
+import java.util.concurrent.ConcurrentSkipListSet;
 import java.util.concurrent.ThreadLocalRandom;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicInteger;
@@ -93,6 +95,8 @@ public class HddsVolume extends StorageVolume {
   private final AtomicLong committedBytes = new AtomicLong(); // till Open containers become full
   private Function<HddsVolume, Long> gatherContainerUsages = (K) -> 0L;
 
+  private final ConcurrentSkipListSet<Long> containerIds = new ConcurrentSkipListSet<>();
+
   // Mentions the type of volume
   private final VolumeType type = VolumeType.DATA_VOLUME;
   // The dedicated DbVolume that the db instance of this HddsVolume resides.
@@ -529,6 +533,22 @@ public long getContainers() {
     return 0;
   }
 
+  public void addContainer(long containerId) {
+    containerIds.add(containerId);
+  }
+
+  public void removeContainer(long containerId) {
+    containerIds.remove(containerId);
+  }
+
+  public Iterator<Long> getContainerIterator() {
+    return containerIds.iterator();
+  }
+
+  public long getContainerCount() {
+    return containerIds.size();
+  }
+
   /**
    * Pick a DbVolume for HddsVolume and init db instance.
    * Use the HddsVolume directly if no DbVolume found.

hadoop-hdds/container-service/src/main/java/org/apache/hadoop/ozone/container/ozoneimpl/OzoneContainer.java

@@ -591,11 +591,13 @@ public ContainerSet getContainerSet() {
 
   public Long gatherContainerUsages(HddsVolume storageVolume) {
     AtomicLong usages = new AtomicLong();
-    containerSet.getContainerMapIterator().forEachRemaining(e -> {
-      if (e.getValue().getContainerData().getVolume().getStorageID().equals(storageVolume.getStorageID())) {
-        usages.addAndGet(e.getValue().getContainerData().getBytesUsed());
+    Iterator<Long> containerIdIterator = storageVolume.getContainerIterator();
+    while (containerIdIterator.hasNext()) {
+      Container<?> container = containerSet.getContainer(containerIdIterator.next());
+      if (container != null) {
+        usages.addAndGet(container.getContainerData().getBytesUsed());
       }
-    });
+    }
     return usages.get();
   }
   /**

hadoop-hdds/container-service/src/test/java/org/apache/hadoop/ozone/container/common/impl/TestContainerSet.java

@@ -28,6 +28,7 @@
 import static org.junit.jupiter.api.Assertions.fail;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.anyString;
+import static org.mockito.Mockito.doAnswer;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.when;
 
@@ -40,6 +41,7 @@
 import java.util.Optional;
 import java.util.Random;
 import java.util.UUID;
+import java.util.concurrent.ConcurrentSkipListSet;
 import java.util.concurrent.atomic.AtomicLong;
 import java.util.stream.LongStream;
 import org.apache.hadoop.conf.StorageUnit;
@@ -69,6 +71,33 @@ private void setLayoutVersion(ContainerLayoutVersion layoutVersion) {
     this.layoutVersion = layoutVersion;
   }
 
+  /**
+   * Create a mock {@link HddsVolume} to track container IDs.
+   */
+  private HddsVolume mockHddsVolume(String storageId) {
+    HddsVolume volume = mock(HddsVolume.class);
+    when(volume.getStorageID()).thenReturn(storageId);
+
+    ConcurrentSkipListSet<Long> containerIds = new ConcurrentSkipListSet<>();
+
+    doAnswer(inv -> {
+      Long containerId = inv.getArgument(0);
+      containerIds.add(containerId);
+      return null;
+    }).when(volume).addContainer(any(Long.class));
+
+    doAnswer(inv -> {
+      Long containerId = inv.getArgument(0);
+      containerIds.remove(containerId);
+      return null;
+    }).when(volume).removeContainer(any(Long.class));
+
+    when(volume.getContainerIterator()).thenAnswer(inv -> containerIds.iterator());
+    when(volume.getContainerCount()).thenAnswer(inv -> (long) containerIds.size());
+
+    return volume;
+  }
+
   @ContainerLayoutTestInfo.ContainerTest
   public void testAddGetRemoveContainer(ContainerLayoutVersion layout)
       throws StorageContainerException {
@@ -157,10 +186,8 @@ public void testIteratorsAndCount(ContainerLayoutVersion layout)
   public void testIteratorPerVolume(ContainerLayoutVersion layout)
       throws StorageContainerException {
     setLayoutVersion(layout);
-    HddsVolume vol1 = mock(HddsVolume.class);
-    when(vol1.getStorageID()).thenReturn("uuid-1");
-    HddsVolume vol2 = mock(HddsVolume.class);
-    when(vol2.getStorageID()).thenReturn("uuid-2");
+    HddsVolume vol1 = mockHddsVolume("uuid-1");
+    HddsVolume vol2 = mockHddsVolume("uuid-2");
 
     ContainerSet containerSet = newContainerSet();
     for (int i = 0; i < 10; i++) {
@@ -202,8 +229,7 @@ public void testIteratorPerVolume(ContainerLayoutVersion layout)
   public void iteratorIsOrderedByScanTime(ContainerLayoutVersion layout)
       throws StorageContainerException {
     setLayoutVersion(layout);
-    HddsVolume vol = mock(HddsVolume.class);
-    when(vol.getStorageID()).thenReturn("uuid-1");
+    HddsVolume vol = mockHddsVolume("uuid-1");
     Random random = new Random();
     ContainerSet containerSet = newContainerSet();
     int containerCount = 50;
@@ -375,4 +401,102 @@ private ContainerSet createContainerSet() throws StorageContainerException {
     return containerSet;
   }
 
+  /**
+   * Test that containerCount per volume returns correct count.
+   */
+  @ContainerLayoutTestInfo.ContainerTest
+  public void testContainerCountPerVolume(ContainerLayoutVersion layout)
+      throws StorageContainerException {
+    setLayoutVersion(layout);
+    HddsVolume vol1 = mockHddsVolume("uuid-1");
+    HddsVolume vol2 = mockHddsVolume("uuid-2");
+    HddsVolume vol3 = mockHddsVolume("uuid-3");
+
+    ContainerSet containerSet = newContainerSet();
+
+    // Add 100 containers to vol1, 50 to vol2, 0 to vol3
+    for (int i = 0; i < 100; i++) {
+      KeyValueContainerData kvData = new KeyValueContainerData(i,
+          layout,
+          (long) StorageUnit.GB.toBytes(5), UUID.randomUUID().toString(),
+          UUID.randomUUID().toString());
+      kvData.setVolume(vol1);
+      kvData.setState(ContainerProtos.ContainerDataProto.State.CLOSED);
+      containerSet.addContainer(new KeyValueContainer(kvData, new OzoneConfiguration()));
+    }
+
+    for (int i = 100; i < 150; i++) {
+      KeyValueContainerData kvData = new KeyValueContainerData(i,
+          layout,
+          (long) StorageUnit.GB.toBytes(5), UUID.randomUUID().toString(),
+          UUID.randomUUID().toString());
+      kvData.setVolume(vol2);
+      kvData.setState(ContainerProtos.ContainerDataProto.State.CLOSED);
+      containerSet.addContainer(new KeyValueContainer(kvData, new OzoneConfiguration()));
+    }
+
+    // Verify counts
+    assertEquals(100, containerSet.containerCount(vol1));
+    assertEquals(50, containerSet.containerCount(vol2));
+    assertEquals(0, containerSet.containerCount(vol3));
+
+    // Remove some containers and verify counts are updated
+    containerSet.removeContainer(0);
+    containerSet.removeContainer(1);
+    containerSet.removeContainer(100);
+    assertEquals(98, containerSet.containerCount(vol1));
+    assertEquals(49, containerSet.containerCount(vol2));
+  }
+
+  /**
+   * Test that per-volume iterator only returns containers from that volume.
+   */
+  @ContainerLayoutTestInfo.ContainerTest
+  public void testContainerIteratorPerVolume(ContainerLayoutVersion layout)
+      throws StorageContainerException {
+    setLayoutVersion(layout);
+    HddsVolume vol1 = mockHddsVolume("uuid-11");
+    HddsVolume vol2 = mockHddsVolume("uuid-12");
+
+    ContainerSet containerSet = newContainerSet();
+
+    // Add containers with specific IDs to each volume
+    List<Long> vol1Ids = new ArrayList<>();
+    List<Long> vol2Ids = new ArrayList<>();
+
+    for (int i = 0; i < 20; i++) {
+      KeyValueContainerData kvData = new KeyValueContainerData(i,
+          layout,
+          (long) StorageUnit.GB.toBytes(5), UUID.randomUUID().toString(),
+          UUID.randomUUID().toString());
+      if (i % 2 == 0) {
+        kvData.setVolume(vol1);
+        vol1Ids.add((long) i);
+      } else {
+        kvData.setVolume(vol2);
+        vol2Ids.add((long) i);
+      }
+      kvData.setState(ContainerProtos.ContainerDataProto.State.CLOSED);
+      containerSet.addContainer(new KeyValueContainer(kvData, new OzoneConfiguration()));
+    }
+
+    // Verify iterator only returns containers from vol1
+    Iterator<Container<?>> iter1 = containerSet.getContainerIterator(vol1);
+    List<Long> foundVol1Ids = new ArrayList<>();
+    while (iter1.hasNext()) {
+      foundVol1Ids.add(iter1.next().getContainerData().getContainerID());
+    }
+    assertEquals(vol1Ids.size(), foundVol1Ids.size());
+    assertTrue(foundVol1Ids.containsAll(vol1Ids));
+
+    // Verify iterator only returns containers from vol2
+    Iterator<Container<?>> iter2 = containerSet.getContainerIterator(vol2);
+    List<Long> foundVol2Ids = new ArrayList<>();
+    while (iter2.hasNext()) {
+      foundVol2Ids.add(iter2.next().getContainerData().getContainerID());
+    }
+    assertEquals(vol2Ids.size(), foundVol2Ids.size());
+    assertTrue(foundVol2Ids.containsAll(vol2Ids));
+  }
+
 }

hadoop-hdds/container-service/src/test/java/org/apache/hadoop/ozone/container/common/statemachine/commandhandler/TestDeleteBlocksCommandHandler.java

@@ -49,6 +49,7 @@
 import java.util.Map;
 import java.util.UUID;
 import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.ConcurrentSkipListSet;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.Future;
 import java.util.concurrent.TimeUnit;
@@ -104,13 +105,31 @@ private void prepareTest(ContainerTestVersionInfo versionInfo)
     setup();
   }
 
+  /**
+   * Create a mock {@link HddsVolume} to track container IDs.
+   */
+  private HddsVolume mockHddsVolume(String storageId) {
+    HddsVolume volume = mock(HddsVolume.class);
+    when(volume.getStorageID()).thenReturn(storageId);
+
+    ConcurrentSkipListSet<Long> containerIds = new ConcurrentSkipListSet<>();
+
+    doAnswer(inv -> {
+      Long containerId = inv.getArgument(0);
+      containerIds.add(containerId);
+      return null;
+    }).when(volume).addContainer(any(Long.class));
+
+    when(volume.getContainerIterator()).thenAnswer(inv -> containerIds.iterator());
+    return volume;
+  }
+
   private void setup() throws Exception {
     OzoneConfiguration conf = new OzoneConfiguration();
     ContainerLayoutVersion layout = ContainerLayoutVersion.FILE_PER_BLOCK;
     OzoneContainer ozoneContainer = mock(OzoneContainer.class);
     containerSet = newContainerSet();
-    volume1 = mock(HddsVolume.class);
-    when(volume1.getStorageID()).thenReturn("uuid-1");
+    volume1 = mockHddsVolume("uuid-1");
     for (int i = 0; i <= 10; i++) {
       KeyValueContainerData data =
           new KeyValueContainerData(i,