add select by zone logic for read/write pool selection

Author: lens161

modules/dcache/src/main/java/diskCacheV111/poolManager/PoolMonitorV5.java

@@ -259,6 +259,8 @@ public SelectedPool selectWritePool(long preallocated)
                             .map(_costModule::getPoolInfo)
                             .filter(Objects::nonNull)
                             .collect(toList());
+                pools = filterByZone(pools);
+
                 if (!pools.isEmpty()) {
                     Partition partition = _partitionManager.getPartition(level.getTag());
                     try {
@@ -341,6 +343,8 @@ public SelectedPool selectReadPool()
                 LOGGER.debug("[read] Available pools at level {}: {}",
                       prio, pools);
 
+                pools = filterByZone(pools);
+
                 /* If allowed, fallback to next link if current link doesn't point
                  * to any pool holding the file.
                  */
@@ -393,6 +397,17 @@ public SelectedPool selectReadPool()
                   "File is online, but not in read-allowed pool");
         }
 
+        private List<PoolInfo> filterByZone(List<PoolInfo> pools) {
+            if (_zone.isPresent()) {
+                List<PoolInfo> poolsZ =
+                        pools.stream()
+                        .filter(p -> Objects.equals(p.getTags().get("zone"), _zone.get()))
+                        .toList();
+                pools = (poolsZ.isEmpty()) ? pools : poolsZ;
+            }
+            return pools;
+        }
+
         @Nullable
         private String getHostName() {
             if (_protocolInfo instanceof IpProtocolInfo) {

modules/dcache/src/test/java/org/dcache/tests/poolmanager/PoolMonitorTest.java

@@ -29,6 +29,8 @@
 import java.util.Collections;
 import java.util.HashSet;
 import java.util.List;
+import java.util.Map;
+import java.util.Optional;
 import java.util.Set;
 import org.dcache.cells.UniversalSpringCell;
 import org.dcache.pool.classic.IoQueueManager;
@@ -151,7 +153,37 @@ public void testGsonDeserialization() throws Exception {
               .disableHtmlEscaping().create().toJson(obj);
     }
 
+    @Test
+    public void testWritePoolZonePreference() throws Exception {
+        prepareCostModule(false, true);
+
+        FileAttributes attributes = FileAttributes.of().pnfsId(_pnfsId).build();
+        StorageInfos.injectInto(_storageInfo, attributes);
+
+        PoolSelector selector = _poolMonitor.getPoolSelector(attributes, _protocolInfo,
+                null, Optional.of("1"), Collections.EMPTY_SET);
+
+        assertEquals("pool1", selector.selectWritePool(0).name());
+    }
+
+    @Test
+    public void testReadPoolZonePreference() throws Exception {
+        prepareCostModule(false, true);
+
+        FileAttributes attributes = FileAttributes.of().pnfsId(_pnfsId).locations(_pools).build();
+        StorageInfos.injectInto(_storageInfo, attributes);
+
+        PoolSelector selector = _poolMonitor.getPoolSelector(attributes, _protocolInfo,
+                null, Optional.of("1"), Collections.EMPTY_SET);
+
+        assertEquals("pool1", selector.selectReadPool().name());
+    }
+
     private void prepareCostModule(boolean linkPerPool) throws Exception {
+        prepareCostModule(linkPerPool, false);
+    }
+
+    private void prepareCostModule(boolean linkPerPool, boolean withZones) throws Exception {
         if (linkPerPool) {
             PoolMonitorHelper.prepareLinkPerPool(_selectionUnit, _access, _pools);
         } else {
@@ -179,6 +211,11 @@ private void prepareCostModule(boolean linkPerPool) throws Exception {
         pool1UpMessage.setHostName(_localhost);
         pool2UpMessage.setHostName(_localhost);
 
+        if (withZones) {
+            pool1UpMessage.setTagMap(Map.of("zone", "1"));
+            pool2UpMessage.setTagMap(Map.of("zone", "2"));
+        }
+
         CellMessage envelope1 = new CellMessage(new CellAddressCore("PoolManager"), null);
         envelope1.addSourceAddress(new CellAddressCore("pool1"));
         CellMessage envelope2 = new CellMessage(new CellAddressCore("PoolManager"), null);