HDDS-13639. Optimize container iterator for frequent operation

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

@@ -34,12 +34,15 @@
 import java.util.Map;
 import java.util.Objects;
 import java.util.Set;
+import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ConcurrentNavigableMap;
 import java.util.concurrent.ConcurrentSkipListMap;
 import java.util.concurrent.ConcurrentSkipListSet;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.atomic.AtomicLong;
 import java.util.function.ToLongFunction;
+import java.util.stream.Collectors;
 import org.apache.hadoop.hdds.protocol.datanode.proto.ContainerProtos;
 import org.apache.hadoop.hdds.protocol.datanode.proto.ContainerProtos.ContainerDataProto.State;
 import org.apache.hadoop.hdds.protocol.proto.StorageContainerDatanodeProtocolProtos.ContainerReportsProto;
@@ -75,6 +78,14 @@ public class ContainerSet implements Iterable<Container<?>> {
   private final WitnessedContainerMetadataStore containerMetadataStore;
   // Handler that will be invoked when a scan of a container in this set is requested.
   private OnDemandContainerScanner containerScanner;
+
+  // Maps volume storage ID to a set of container IDs
+  private final ConcurrentHashMap<String, ConcurrentSkipListSet<Long>> volumeToContainersMap =
+      new ConcurrentHashMap<>();
+
+  // Maps volume storage ID to container count
+  private final ConcurrentHashMap<String, AtomicLong> volumeContainerCountCache =
+      new ConcurrentHashMap<>();
 
   public static ContainerSet newReadOnlyContainerSet(long recoveringTimeout) {
     return new ContainerSet(null, recoveringTimeout);
@@ -205,6 +216,9 @@ private boolean addContainer(Container<?> container, boolean overwrite) throws
         recoveringContainerMap.put(
             clock.millis() + recoveringTimeout, containerId);
       }
+      // Update per-volume index and count cache
+      updateVolumeIndexOnAdd(container);
+
       return true;
     } else {
       LOG.warn("Container already exists with container Id {}", containerId);
@@ -213,6 +227,29 @@ private boolean addContainer(Container<?> container, boolean overwrite) throws
           ContainerProtos.Result.CONTAINER_EXISTS);
     }
   }
+
+  /**
+   * Updates the per-volume container index when a container is added.
+   *
+   * @param container the container being added
+   */
+  private void updateVolumeIndexOnAdd(Container<?> container) {
+    HddsVolume volume = container.getContainerData().getVolume();
+    if (volume != null && volume.getStorageID() != null) {
+      String volumeUuid = volume.getStorageID();
+      long containerId = container.getContainerData().getContainerID();
+
+      // Add container ID to volume's container set
+      volumeToContainersMap
+          .computeIfAbsent(volumeUuid, k -> new ConcurrentSkipListSet<>())
+          .add(containerId);
+
+      // Increment volume's container count
+      volumeContainerCountCache
+          .computeIfAbsent(volumeUuid, k -> new AtomicLong(0))
+          .incrementAndGet();
+    }
+  }
 
   private void updateContainerIdTable(long containerId, ContainerData containerData) throws StorageContainerException {
     if (null != containerMetadataStore) {
@@ -299,11 +336,49 @@ private boolean removeContainer(long containerId, boolean markMissing, boolean r
           "containerMap", containerId);
       return false;
     } else {
+      // Update per-volume index and count cache
+      updateVolumeIndexOnRemove(removed);
+
       LOG.debug("Container with containerId {} is removed from containerMap",
           containerId);
       return true;
     }
   }
+
+  /**
+   * Updates the per-volume container index when a container is removed.
+   *
+   * @param container the container being removed
+   */
+  private void updateVolumeIndexOnRemove(Container<?> container) {
+    HddsVolume volume = container.getContainerData().getVolume();
+    if (volume != null && volume.getStorageID() != null) {
+      String volumeUuid = volume.getStorageID();
+      long containerId = container.getContainerData().getContainerID();
+
+      // Remove container ID from volume's container set
+      ConcurrentSkipListSet<Long> containerSet = volumeToContainersMap.get(volumeUuid);
+      if (containerSet != null) {
+        containerSet.remove(containerId);
+
+        // If the set is now empty, remove it from the map to save memory
+        if (containerSet.isEmpty()) {
+          volumeToContainersMap.remove(volumeUuid);
+        }
+      }
+
+      // Decrement volume's container count
+      AtomicLong count = volumeContainerCountCache.get(volumeUuid);
+      if (count != null) {
+        long newCount = count.decrementAndGet();
+
+        // If count reaches zero, remove from cache to save memory
+        if (newCount <= 0) {
+          volumeContainerCountCache.remove(volumeUuid);
+        }
+      }
+    }
+  }
 
   private void deleteFromContainerTable(long containerId) throws StorageContainerException {
     if (null != containerMetadataStore) {
@@ -411,11 +486,18 @@ 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()))
+
+    ConcurrentSkipListSet<Long> containerIds = volumeToContainersMap.get(volumeUuid);
+    if (containerIds == null || containerIds.isEmpty()) {
+      return Collections.emptyIterator();
+    }
+    List<Container<?>> containers = containerIds.stream()
+        .map(containerMap::get)
+        .filter(Objects::nonNull)
         .sorted(ContainerDataScanOrder.INSTANCE)
-        .iterator();
+        .collect(Collectors.toList());
+
+    return containers.iterator();
   }
 
   /**
@@ -428,9 +510,8 @@ 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();
+    AtomicLong count = volumeContainerCountCache.get(volumeUuid);
+    return count != null ? count.get() : 0;
   }
 
   /**

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

@@ -375,4 +375,107 @@ 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 = mock(HddsVolume.class);
+    when(vol1.getStorageID()).thenReturn("uuid-1");
+    HddsVolume vol2 = mock(HddsVolume.class);
+    when(vol2.getStorageID()).thenReturn("uuid-2");
+    HddsVolume vol3 = mock(HddsVolume.class);
+    when(vol3.getStorageID()).thenReturn("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 = mock(HddsVolume.class);
+    when(vol1.getStorageID()).thenReturn("uuid-11");
+    HddsVolume vol2 = mock(HddsVolume.class);
+    when(vol2.getStorageID()).thenReturn("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));
+  }
+
 }