diff --git a/hadoop-hdds/container-service/src/main/java/org/apache/hadoop/ozone/container/common/impl/ContainerSet.java b/hadoop-hdds/container-service/src/main/java/org/apache/hadoop/ozone/container/common/impl/ContainerSet.java
index 69a9ab138293..13bf7789929d 100644
--- a/hadoop-hdds/container-service/src/main/java/org/apache/hadoop/ozone/container/common/impl/ContainerSet.java
+++ b/hadoop-hdds/container-service/src/main/java/org/apache/hadoop/ozone/container/common/impl/ContainerSet.java
@@ -254,7 +254,18 @@ public Container updateContainer(Container<?> container) throws
     } else {
       LOG.debug("Container with container Id {} is updated to containerMap",
           containerId);
-      return containerMap.put(containerId, container);
+      Container<?> oldContainer = containerMap.put(containerId, container);
+      HddsVolume volume = container.getContainerData().getVolume();
+      if (volume != null) {
+        volume.addContainer(container.getContainerData().getContainerID());
+      }
+      if (oldContainer != null) {
+        volume = oldContainer.getContainerData().getVolume();
+        if (volume != null) {
+          volume.removeContainer(oldContainer.getContainerData().getContainerID());
+        }
+      }
+      return oldContainer;
     }
   }
 
@@ -446,7 +457,7 @@ public Iterator<Container<?>> getContainerIterator(HddsVolume volume) {
     while (containerIdIterator.hasNext()) {
       Long containerId = containerIdIterator.next();
       Container<?> container = containerMap.get(containerId);
-      if (container != null) {
+      if (container != null && container.getContainerData().getVolume() == volume) {
         containers.add(container);
       }
     }
diff --git a/hadoop-hdds/container-service/src/test/java/org/apache/hadoop/ozone/container/diskbalancer/TestDiskBalancerTask.java b/hadoop-hdds/container-service/src/test/java/org/apache/hadoop/ozone/container/diskbalancer/TestDiskBalancerTask.java
index 5e2df4eb3926..af9efdf75049 100644
--- a/hadoop-hdds/container-service/src/test/java/org/apache/hadoop/ozone/container/diskbalancer/TestDiskBalancerTask.java
+++ b/hadoop-hdds/container-service/src/test/java/org/apache/hadoop/ozone/container/diskbalancer/TestDiskBalancerTask.java
@@ -40,6 +40,7 @@
 import java.nio.file.Path;
 import java.nio.file.Paths;
 import java.util.HashMap;
+import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
 import java.util.UUID;
@@ -282,6 +283,11 @@ public void moveSuccess(State containerState) throws IOException {
     State originalState = container.getContainerState();
     assertEquals(initialSourceUsed + CONTAINER_SIZE, sourceVolume.getCurrentUsage().getUsedSpace());
     assertEquals(initialDestUsed, destVolume.getCurrentUsage().getUsedSpace());
+    Iterator<Long> containerIterator = sourceVolume.getContainerIterator();
+    assertTrue(containerIterator.hasNext());
+    assertEquals(CONTAINER_ID, containerIterator.next());
+    containerIterator = destVolume.getContainerIterator();
+    assertFalse(containerIterator.hasNext());
 
     String oldContainerPath = container.getContainerData().getContainerPath();
     if (containerState == State.QUASI_CLOSED) {
@@ -305,6 +311,12 @@ public void moveSuccess(State containerState) throws IOException {
     assertEquals(CONTAINER_SIZE, diskBalancerService.getMetrics().getSuccessBytes());
     assertEquals(initialDestCommitted, destVolume.getCommittedBytes());
     assertEquals(initialSourceDelta, diskBalancerService.getDeltaSizes().get(sourceVolume));
+
+    containerIterator = sourceVolume.getContainerIterator();
+    assertFalse(containerIterator.hasNext());
+    containerIterator = destVolume.getContainerIterator();
+    assertTrue(containerIterator.hasNext());
+    assertEquals(CONTAINER_ID, containerIterator.next());
   }
 
   @ContainerTestVersionInfo.ContainerTest