p2p: select dst by zone in WassPartition

Author: lens161

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

@@ -509,6 +509,7 @@ public Partition.P2pPair selectPool2Pool(String poolGroup,
                           Lists.newArrayList(sources.values()),
                           pools,
                           _fileAttributes,
+                          _zone,
                           force);
                 }
             }

modules/dcache/src/main/java/org/dcache/poolmanager/BufferPartition.java

@@ -145,7 +145,7 @@ public SelectedPool selectReadPool(CostModule cm, List<PoolInfo> pools,
 
     @Override
     public P2pPair selectPool2Pool(CostModule cm, List<PoolInfo> src, List<PoolInfo> dst,
-          FileAttributes attributes, boolean force) throws CacheException {
+          FileAttributes attributes, Optional<String> zone, boolean force) throws CacheException {
         return new P2pPair(selectPool(src), selectPool(dst, attributes.getSize()));
     }
 

modules/dcache/src/main/java/org/dcache/poolmanager/LruPartition.java

@@ -133,6 +133,7 @@ public P2pPair selectPool2Pool(CostModule cm,
           List<PoolInfo> src,
           List<PoolInfo> dst,
           FileAttributes attributes,
+          Optional<String> zone,
           boolean force)
           throws CacheException {
         return new P2pPair(new SelectedPool(select(src, _lastRead)),

modules/dcache/src/main/java/org/dcache/poolmanager/Partition.java

@@ -335,6 +335,7 @@ public abstract P2pPair selectPool2Pool(CostModule cm,
           List<PoolInfo> src,
           List<PoolInfo> dst,
           FileAttributes attributes,
+          Optional<String> zone,
           boolean force)
           throws CacheException;
 

modules/dcache/src/main/java/org/dcache/poolmanager/RandomPartition.java

@@ -82,6 +82,7 @@ public SelectedPool selectReadPool(CostModule cm,
           List<PoolInfo> src,
           List<PoolInfo> dst,
           FileAttributes attributes,
+          Optional<String> zone,
           boolean force)
           throws CacheException {
         return new P2pPair(new SelectedPool(select(src)),

modules/dcache/src/main/java/org/dcache/poolmanager/WRandomPartition.java

@@ -51,7 +51,7 @@ public String getType() {
     }
 
     @Override
-    public P2pPair selectPool2Pool(CostModule cm, List<PoolInfo> src, List<PoolInfo> dst, FileAttributes attributes, boolean force) throws CacheException {
+    public P2pPair selectPool2Pool(CostModule cm, List<PoolInfo> src, List<PoolInfo> dst, FileAttributes attributes, Optional<String> zone, boolean force) throws CacheException {
 
         Collections.shuffle(src);
 

modules/dcache/src/main/java/org/dcache/poolmanager/WassPartition.java

@@ -83,6 +83,7 @@ public P2pPair selectPool2Pool(CostModule cm,
           List<PoolInfo> src,
           List<PoolInfo> dst,
           FileAttributes attributes,
+          Optional<String> zone,
           boolean force)
           throws CacheException {
         checkState(!src.isEmpty());
@@ -139,6 +140,11 @@ public P2pPair selectPool2Pool(CostModule cm,
               : new AvailableSpaceAssumption(filesize).and(PerformanceCostAssumption.of(
                     _error, maxTargetCost));
 
+        if (zone.isPresent()) {
+            List<PoolInfo> dstZ = filterByZone(dst, zone);
+            dst = (dstZ.isEmpty()) ? dst : dstZ;
+        }
+
         if (_allowSameHostCopy != SameHost.NOTCHECKED) {
             /* Loop over all sources and find the most appropriate
              * destination such that same host constraints are
@@ -177,6 +183,13 @@ public P2pPair selectPool2Pool(CostModule cm,
               new SelectedPool(destination, destinationAssumption));
     }
 
+    private List<PoolInfo> filterByZone(List<PoolInfo> pools, Optional<String> zone) {
+        pools = pools.stream()
+                .filter(p -> Objects.equals(p.getTags().get("zone"), zone.get()))
+                .toList();
+        return pools;
+    }
+
     private PoolInfo selectByPrevious(List<PoolInfo> pools,
           Optional<PoolInfo> previous,
           FileAttributes attributes) {

modules/dcache/src/test/java/org/dcache/mock/SimplePartition.java

@@ -66,7 +66,7 @@ public SelectedPool selectReadPool(CostModule cm, List<PoolInfo> pools,
 
     @Override
     public P2pPair selectPool2Pool(CostModule cm, List<PoolInfo> src,
-          List<PoolInfo> dst, FileAttributes attributes, boolean force)
+          List<PoolInfo> dst, FileAttributes attributes, Optional<String> zone, boolean force)
           throws CacheException {
         throw new UnsupportedOperationException("Not supported yet.");
     }

modules/dcache/src/test/java/org/dcache/poolmanager/WassPartitionTest.java

@@ -0,0 +1,74 @@
+package org.dcache.poolmanager;
+
+import com.google.common.collect.ImmutableMap;
+import diskCacheV111.poolManager.CostModule;
+import diskCacheV111.pools.PoolCostInfo;
+import diskCacheV111.util.CacheException;
+import dmg.cells.nucleus.CellAddressCore;
+import org.dcache.vehicles.FileAttributes;
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.jupiter.api.RepeatedTest;
+
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+import java.util.stream.IntStream;
+import java.util.stream.Stream;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertTrue;
+
+public class WassPartitionTest {
+
+    WassPartition wass;
+    FileAttributes attributes;
+    List<PoolInfo> dst;
+    List<PoolInfo> src;
+    Optional<String> zone;
+
+    @Before
+    public void setUp() {
+        wass = new WassPartition(Map.of());
+        attributes = FileAttributes.ofSize(1000L);
+        zone = Optional.of("A");
+
+        src = IntStream.range(0, 1).mapToObj(i -> {
+                    var cost = new PoolCostInfo("pool" + i, "default-queue");
+                    cost.setSpaceUsage(10_000L, 5000L, 0L, 0L);
+                    cost.getSpaceInfo().setParameter(0.0d, 2500L);
+
+                    return new PoolInfo(new CellAddressCore("pool" + i), cost, ImmutableMap.of("zone", "A"));
+                }
+        ).toList();
+
+        List<PoolInfo>dstCorrect = IntStream.range(2, 4).mapToObj(i -> {
+                    var cost = new PoolCostInfo("pool" + i, "default-queue");
+                    cost.setSpaceUsage(10_000L, 5000L, 0L, 0L);
+                    cost.getSpaceInfo().setParameter(0.0d, 2500L);
+
+                    return new PoolInfo(new CellAddressCore("pool" + i), cost, ImmutableMap.of("zone", "A"));
+                }
+        ).toList();
+
+        List<PoolInfo>dstWrong = IntStream.range(5, 7).mapToObj(i -> {
+                    var cost = new PoolCostInfo("pool" + i, "default-queue");
+                    cost.setSpaceUsage(10_000L, 5000L, 0L, 0L);
+                    cost.getSpaceInfo().setParameter(0.0d, 2500L);
+
+                    return new PoolInfo(new CellAddressCore("pool" + i), cost, ImmutableMap.of("zone", "B"));
+                }
+        ).toList();
+
+        dst = Stream.concat(dstCorrect.stream(), dstWrong.stream()).toList();
+    }
+
+    @Test
+    public void testCorrectZoneDst() throws CacheException {
+        for (int i = 0; i < 9; i++){
+            PoolInfo info = wass.selectPool2Pool(null, src, dst, attributes, zone, false).destination.info();
+            assertEquals(zone.get(), info.getTags().get("zone"));
+        }
+    }
+}