Allow to bypass the dynamic bandwidth estimation in NetworkCost

common/src/main/java/org/astraea/common/cost/NetworkCost.java

@@ -205,6 +205,24 @@ public MetricSensor metricSensor() {
             .toList();
   }
 
+  /**
+   * Set {@code partitionNetIn} and {@code partitionNetOut} network usage to the given {@code
+   * handle}. This function offers a way to specify the network usage of each partition directly,
+   * which can be crucial for situations like software simulation, where every partition resource
+   * usage and pattern is deterministic before the simulation run, so there is no need to trigger a
+   * bandwidth estimation.
+   *
+   * @param handle a {@link ClusterBean} object associate with the following partition usage.
+   * @param partitionNetIn the network ingress usage of each partition.
+   * @param partitionNetOut the network egress usage of each partition.
+   */
+  public void setCalculation(
+      ClusterBean handle,
+      Map<TopicPartition, Long> partitionNetIn,
+      Map<TopicPartition, Long> partitionNetOut) {
+    this.calculationCache.put(handle, new CachedCalculation(partitionNetIn, partitionNetOut));
+  }
+
   private Map<BrokerTopic, List<Replica>> mapLeaderAllocation(ClusterInfo clusterInfo) {
     return clusterInfo
         .replicaStream()
@@ -342,6 +360,11 @@ private CachedCalculation(ClusterBean sourceMetric) {
       this.partitionEgressRate =
           estimateRate(metricViewCluster, sourceMetric, ServerMetrics.Topic.BYTES_OUT_PER_SEC);
     }
+
+    private CachedCalculation(Map<TopicPartition, Long> netIn, Map<TopicPartition, Long> netOut) {
+      this.partitionIngressRate = netIn;
+      this.partitionEgressRate = netOut;
+    }
   }
 
   static class NetworkClusterCost implements ClusterCost {

common/src/test/java/org/astraea/common/cost/NetworkCostTest.java

@@ -57,6 +57,7 @@
 import org.junit.jupiter.params.provider.CsvSource;
 import org.junit.jupiter.params.provider.MethodSource;
 import org.junit.jupiter.params.provider.ValueSource;
+import org.mockito.Mockito;
 
 class NetworkCostTest {
 
@@ -490,6 +491,40 @@ void testPartitionCost() {
     Assertions.assertEquals(ingressPartitionCost.get(TopicPartition.of("test-8")), (double) 9 / 18);
   }
 
+  @ParameterizedTest
+  @ValueSource(classes = {NetworkIngressCost.class, NetworkEgressCost.class})
+  void testSetCalculation(Class<? extends NetworkCost> costClass) {
+    var cluster =
+        ClusterInfo.builder()
+            .addNode(Set.of(1, 2))
+            .addFolders(Map.of(1, Set.of("/a")))
+            .addFolders(Map.of(2, Set.of("/a")))
+            .addTopic("topic", 1, (short) 1)
+            .build();
+    var handle = ClusterBean.of(Map.of());
+    var ingress = Mockito.spy(Map.of(TopicPartition.of("topic-0"), 1024L));
+    var egress = Mockito.spy(Map.of(TopicPartition.of("topic-0"), 1024L));
+    var costFunction = Utils.construct(costClass, Configuration.EMPTY);
+
+    costFunction.setCalculation(handle, ingress, egress);
+
+    Assertions.assertDoesNotThrow(
+        () -> {
+          Mockito.verifyNoInteractions(ingress);
+          Mockito.verifyNoInteractions(egress);
+        },
+        "The ingress/egress maps should not be touched");
+    Assertions.assertDoesNotThrow(
+        () -> costFunction.clusterCost(cluster, handle),
+        "Trigger the cost evaluation, this should touch the specified calculations");
+    Assertions.assertDoesNotThrow(
+        () -> {
+          Mockito.verify(ingress, Mockito.atLeastOnce()).get(Mockito.any(TopicPartition.class));
+          Mockito.verify(egress, Mockito.atLeastOnce()).get(Mockito.any(TopicPartition.class));
+        },
+        "The ingress/egress maps should be touched");
+  }
+
   @ParameterizedTest
   @CsvSource({
     "BYTES_IN_PER_SEC, org.astraea.common.cost.NetworkIngressCost",