Repository: skiptests/astraea
Pull request #261 (2022-03-01 14:00:47)

Introduce Fetcher
新增`Fetcher`方便後續cost function擴增。之前直接使用java function會限制我們之後加上更多方便的函式


Patch:
diff --git a/app/src/main/java/org/astraea/metrics/collector/BeanCollector.java b/app/src/main/java/org/astraea/metrics/collector/BeanCollector.java
index 540b47d3..36154ecd 100644
--- a/app/src/main/java/org/astraea/metrics/collector/BeanCollector.java
+++ b/app/src/main/java/org/astraea/metrics/collector/BeanCollector.java
@@ -1,23 +1,20 @@
 package org.astraea.metrics.collector;
 
 import java.time.Duration;
-import java.util.ArrayList;
-import java.util.Collection;
 import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
 import java.util.Set;
-import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ConcurrentMap;
 import java.util.concurrent.ConcurrentSkipListMap;
 import java.util.concurrent.locks.Lock;
 import java.util.concurrent.locks.ReentrantLock;
 import java.util.function.BiFunction;
-import java.util.function.Function;
 import org.astraea.Utils;
 import org.astraea.metrics.HasBeanObject;
 import org.astraea.metrics.jmx.MBeanClient;
+import org.astraea.metrics.kafka.KafkaMetrics;
 
 public class BeanCollector {
 
@@ -72,7 +69,7 @@ public class BeanCollector {
     return new Register() {
       private String host;
       private int port = -1;
-      private final List<Function<MBeanClient, HasBeanObject>> getters = new ArrayList<>();
+      private Fetcher fetcher = client -> List.of(KafkaMetrics.Host.jvmMemory(client));
 
       @Override
       public Register host(String host) {
@@ -87,8 +84,8 @@ public class BeanCollector {
       }
 
       @Override
-      public Register metricsGetters(Collection<Function<MBeanClient, HasBeanObject>> getters) {
-        this.getters.addAll(Objects.requireNonNull(getters));
+      public Register fetcher(Fetcher fetcher) {
+        this.fetcher = Objects.requireNonNull(fetcher);
         return this;
       }
 
@@ -98,7 +95,7 @@ public class BeanCollector {
         var node = nodes.computeIfAbsent(nodeKey, ignored -> new Node(host, port));
         var receiver =
             new Receiver() {
-              private final Map<Long, HasBeanObject> objects = new ConcurrentHashMap<>();
+              private final Map<Long, HasBeanObject> objects = new ConcurrentSkipListMap<>();
 
               @Override
               public String host() {
@@ -138,14 +135,14 @@ public class BeanCollector {
                   try {
                     if (node.mBeanClient == null)
                       node.mBeanClient = clientCreator.apply(host, port);
-                    getters.forEach(
-                        getter -> {
-                          if (objects.size() >= numberOfObjectsPerNode)
-                            objects.keySet().stream()
-                                .min((Long::compare))
-                                .ifPresent(objects::remove);
-                          objects.put(System.currentTimeMillis(), getter.apply(node.mBeanClient));
-                        });
+                    var beans = fetcher.fetch(node.mBeanClient);
+                    // remove old beans if the queue is full
+                    for (var t : objects.keySet()) {
+                      if (objects.size() + beans.size() <= numberOfObjectsPerNode) break;
+                      objects.remove(t);
+                    }
+                    long now = System.currentTimeMillis();
+                    for (var bean : beans) objects.put(now++, bean);
                   } finally {
                     node.lock.unlock();
                   }
diff --git a/app/src/main/java/org/astraea/metrics/collector/Fetcher.java b/app/src/main/java/org/astraea/metrics/collector/Fetcher.java
new file mode 100644
index 00000000..9e1b3677
--- /dev/null
+++ b/app/src/main/java/org/astraea/metrics/collector/Fetcher.java
@@ -0,0 +1,29 @@
+package org.astraea.metrics.collector;
+
+import java.util.Collection;
+import java.util.stream.Collectors;
+import org.astraea.metrics.HasBeanObject;
+import org.astraea.metrics.jmx.MBeanClient;
+
+@FunctionalInterface
+public interface Fetcher {
+
+  /**
+   * merge multiples fetchers to single one.
+   *
+   * @param fs fetchers
+   * @return single fetcher
+   */
+  static Fetcher of(Collection<Fetcher> fs) {
+    return client ->
+        fs.stream().flatMap(f -> f.fetch(client).stream()).collect(Collectors.toUnmodifiableList());
+  }
+
+  /**
+   * fetch to specify metrics from remote JMX server
+   *
+   * @param client mbean client (don't close it!)
+   * @return java metrics
+   */
+  Collection<HasBeanObject> fetch(MBeanClient client);
+}
diff --git a/app/src/main/java/org/astraea/metrics/collector/Register.java b/app/src/main/java/org/astraea/metrics/collector/Register.java
index a7bf78c5..cce58c98 100644
--- a/app/src/main/java/org/astraea/metrics/collector/Register.java
+++ b/app/src/main/java/org/astraea/metrics/collector/Register.java
@@ -1,11 +1,5 @@
 package org.astraea.metrics.collector;
 
-import java.util.Collection;
-import java.util.List;
-import java.util.function.Function;
-import org.astraea.metrics.HasBeanObject;
-import org.astraea.metrics.jmx.MBeanClient;
-
 public interface Register {
 
   /**
@@ -20,19 +14,11 @@ public interface Register {
    */
   Register port(int port);
 
-  /**
-   * @param getter to get metrics from MBeanClient
-   * @return this register
-   */
-  default Register metricsGetter(Function<MBeanClient, HasBeanObject> getter) {
-    return metricsGetters(List.of(getter));
-  }
-
   /**
    * @param getters to get metrics from MBeanClient
    * @return this register
    */
-  Register metricsGetters(Collection<Function<MBeanClient, HasBeanObject>> getters);
+  Register fetcher(Fetcher fetcher);
 
   /** @return create a Receiver */
   Receiver build();
diff --git a/app/src/main/java/org/astraea/partitioner/StrictCostDispatcher.java b/app/src/main/java/org/astraea/partitioner/StrictCostDispatcher.java
index e55073f8..d5dbb3e1 100644
--- a/app/src/main/java/org/astraea/partitioner/StrictCostDispatcher.java
+++ b/app/src/main/java/org/astraea/partitioner/StrictCostDispatcher.java
@@ -10,6 +10,7 @@ import java.util.TreeMap;
 import java.util.stream.Collectors;
 import org.astraea.Utils;
 import org.astraea.metrics.collector.BeanCollector;
+import org.astraea.metrics.collector.Fetcher;
 import org.astraea.metrics.collector.Receiver;
 import org.astraea.partitioner.cost.CostFunction;
 
@@ -42,6 +43,11 @@ public class StrictCostDispatcher implements Dispatcher {
     this.functions = functions;
   }
 
+  private static Fetcher of(Collection<Fetcher> fs) {
+    return client ->
+        fs.stream().flatMap(f -> f.fetch(client).stream()).collect(Collectors.toUnmodifiableList());
+  }
+
   @Override
   public int partition(String topic, byte[] key, byte[] value, ClusterInfo clusterInfo) {
     var partitions = clusterInfo.availablePartitions(topic);
@@ -59,10 +65,11 @@ public class StrictCostDispatcher implements Dispatcher {
                         .register()
                         .host(p.leader().host())
                         .port(jmxPort(p.leader().id()))
-                        .metricsGetters(
-                            functions.stream()
-                                .flatMap(c -> c.metricsGetters().stream())
-                                .collect(Collectors.toUnmodifiableList()))
+                        .fetcher(
+                            Fetcher.of(
+                                functions.stream()
+                                    .map(CostFunction::fetcher)
+                                    .collect(Collectors.toUnmodifiableList())))
                         .build()));
 
     // get latest beans for each node
diff --git a/app/src/main/java/org/astraea/partitioner/cost/CostFunction.java b/app/src/main/java/org/astraea/partitioner/cost/CostFunction.java
index 8a630150..c22909da 100644
--- a/app/src/main/java/org/astraea/partitioner/cost/CostFunction.java
+++ b/app/src/main/java/org/astraea/partitioner/cost/CostFunction.java
@@ -1,11 +1,9 @@
 package org.astraea.partitioner.cost;
 
-import java.util.Collection;
 import java.util.List;
 import java.util.Map;
-import java.util.function.Function;
 import org.astraea.metrics.HasBeanObject;
-import org.astraea.metrics.jmx.MBeanClient;
+import org.astraea.metrics.collector.Fetcher;
 import org.astraea.partitioner.ClusterInfo;
 import org.astraea.partitioner.NodeId;
 
@@ -25,5 +23,5 @@ public interface CostFunction {
   Map<NodeId, Double> cost(Map<NodeId, List<HasBeanObject>> beans, ClusterInfo clusterInfo);
 
   /** @return the metrics getters. Those getters are used to fetch mbeans. */
-  Collection<Function<MBeanClient, HasBeanObject>> metricsGetters();
+  Fetcher fetcher();
 }
diff --git a/app/src/main/java/org/astraea/partitioner/cost/ThroughputCost.java b/app/src/main/java/org/astraea/partitioner/cost/ThroughputCost.java
index bf879b03..46d20994 100644
--- a/app/src/main/java/org/astraea/partitioner/cost/ThroughputCost.java
+++ b/app/src/main/java/org/astraea/partitioner/cost/ThroughputCost.java
@@ -1,12 +1,10 @@
 package org.astraea.partitioner.cost;
 
-import java.util.Collection;
 import java.util.List;
 import java.util.Map;
-import java.util.function.Function;
 import java.util.stream.Collectors;
 import org.astraea.metrics.HasBeanObject;
-import org.astraea.metrics.jmx.MBeanClient;
+import org.astraea.metrics.collector.Fetcher;
 import org.astraea.metrics.kafka.BrokerTopicMetricsResult;
 import org.astraea.metrics.kafka.KafkaMetrics;
 import org.astraea.partitioner.ClusterInfo;
@@ -39,9 +37,10 @@ public class ThroughputCost implements CostFunction {
   }
 
   @Override
-  public Collection<Function<MBeanClient, HasBeanObject>> metricsGetters() {
-    return List.of(
-        KafkaMetrics.BrokerTopic.BytesInPerSec::fetch,
-        KafkaMetrics.BrokerTopic.BytesOutPerSec::fetch);
+  public Fetcher fetcher() {
+    return client ->
+        List.of(
+            KafkaMetrics.BrokerTopic.BytesInPerSec.fetch(client),
+            KafkaMetrics.BrokerTopic.BytesOutPerSec.fetch(client));
   }
 }
diff --git a/app/src/main/java/org/astraea/partitioner/nodeLoadMetric/NodeLoadClient.java b/app/src/main/java/org/astraea/partitioner/nodeLoadMetric/NodeLoadClient.java
index 32c30388..34281d9c 100644
--- a/app/src/main/java/org/astraea/partitioner/nodeLoadMetric/NodeLoadClient.java
+++ b/app/src/main/java/org/astraea/partitioner/nodeLoadMetric/NodeLoadClient.java
@@ -154,7 +154,7 @@ public class NodeLoadClient {
           broker.lastOutPut = currentOutput;
           broker.input = Integer.parseInt(String.valueOf((currentInput - broker.lastInput) / 1024));
           broker.lastInput = currentInput;
-
+          Objects.requireNonNull(broker.metrics.get("Memory"));
           broker.jvm =
               (HasJvmMemory)
                   broker.metrics.get("Memory").current().stream().findAny().orElse(broker.jvm);
diff --git a/app/src/main/java/org/astraea/partitioner/nodeLoadMetric/ReceiverFactory.java b/app/src/main/java/org/astraea/partitioner/nodeLoadMetric/ReceiverFactory.java
index eb230376..bdc8bc98 100644
--- a/app/src/main/java/org/astraea/partitioner/nodeLoadMetric/ReceiverFactory.java
+++ b/app/src/main/java/org/astraea/partitioner/nodeLoadMetric/ReceiverFactory.java
@@ -42,21 +42,23 @@ public class ReceiverFactory {
                       .register()
                       .host(host)
                       .port(port)
-                      .metricsGetter(KafkaMetrics.BrokerTopic.BytesInPerSec::fetch)
+                      .fetcher(
+                          client -> List.of(KafkaMetrics.BrokerTopic.BytesInPerSec.fetch(client)))
                       .build());
               receiversList.add(
                   beanCollector
                       .register()
                       .host(host)
                       .port(port)
-                      .metricsGetter(KafkaMetrics.BrokerTopic.BytesOutPerSec::fetch)
+                      .fetcher(
+                          client -> List.of(KafkaMetrics.BrokerTopic.BytesOutPerSec.fetch(client)))
                       .build());
               receiversList.add(
                   beanCollector
                       .register()
                       .host(host)
                       .port(port)
-                      .metricsGetter(KafkaMetrics.Host::jvmMemory)
+                      .fetcher(client -> List.of(KafkaMetrics.Host.jvmMemory(client)))
                       .build());
               count.put(nodeKey(host, port), 1);
             } else {


Test patch:
diff --git a/app/src/test/java/org/astraea/metrics/collector/BeanCollectorTest.java b/app/src/test/java/org/astraea/metrics/collector/BeanCollectorTest.java
index 7526f7e7..801268e9 100644
--- a/app/src/test/java/org/astraea/metrics/collector/BeanCollectorTest.java
+++ b/app/src/test/java/org/astraea/metrics/collector/BeanCollectorTest.java
@@ -38,7 +38,7 @@ public class BeanCollectorTest {
             .register()
             .host("unknown")
             .port(100)
-            .metricsGetter(client -> Mockito.mock(HasBeanObject.class))
+            .fetcher(client -> List.of(Mockito.mock(HasBeanObject.class)))
             .build();
     Assertions.assertEquals("unknown", receiver.host());
     Assertions.assertEquals(100, receiver.port());
@@ -68,7 +68,7 @@ public class BeanCollectorTest {
             .register()
             .host("unknown")
             .port(100)
-            .metricsGetter(client -> Mockito.mock(HasBeanObject.class))
+            .fetcher(client -> List.of(Mockito.mock(HasBeanObject.class)))
             .build();
 
     var c0 = receiver.current();
@@ -103,8 +103,7 @@ public class BeanCollectorTest {
     var collector = BeanCollector.builder().clientCreator(clientCreator).build();
     Assertions.assertThrows(NullPointerException.class, () -> collector.register().host(null));
     Assertions.assertThrows(IllegalArgumentException.class, () -> collector.register().port(-1));
-    Assertions.assertThrows(
-        NullPointerException.class, () -> collector.register().metricsGetter(null));
+    Assertions.assertThrows(NullPointerException.class, () -> collector.register().fetcher(null));
   }
 
   private List<Receiver> receivers(BeanCollector collector) {
@@ -117,7 +116,7 @@ public class BeanCollectorTest {
                     .register()
                     .host("unknown")
                     .port(100)
-                    .metricsGetter(client -> mbean)
+                    .fetcher(client -> List.of(mbean))
                     .build();
             synchronized (receivers) {
               receivers.add(receiver);
@@ -193,10 +192,10 @@ public class BeanCollectorTest {
             .register()
             .host("unknown")
             .port(100)
-            .metricsGetter(
-                c -> {
+            .fetcher(
+                client -> {
                   count.incrementAndGet();
-                  return mbean;
+                  return List.of(mbean);
                 })
             .build();
 
@@ -226,10 +225,10 @@ public class BeanCollectorTest {
                           .register()
                           .host("unknown")
                           .port(100)
-                          .metricsGetter(
-                              c -> {
+                          .fetcher(
+                              client -> {
                                 count.incrementAndGet();
-                                return mbean;
+                                return List.of(mbean);
                               })
                           .build());
                 })
diff --git a/app/src/test/java/org/astraea/metrics/collector/FetcherTest.java b/app/src/test/java/org/astraea/metrics/collector/FetcherTest.java
new file mode 100644
index 00000000..74383c17
--- /dev/null
+++ b/app/src/test/java/org/astraea/metrics/collector/FetcherTest.java
@@ -0,0 +1,27 @@
+package org.astraea.metrics.collector;
+
+import java.util.List;
+import org.astraea.metrics.HasBeanObject;
+import org.astraea.metrics.jmx.MBeanClient;
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.Test;
+import org.mockito.Mockito;
+
+public class FetcherTest {
+
+  @Test
+  void testMultipleFetchers() {
+    var mbean0 = Mockito.mock(HasBeanObject.class);
+    Fetcher fetcher0 = client -> List.of(mbean0);
+    var mbean1 = Mockito.mock(HasBeanObject.class);
+    Fetcher fetcher1 = client -> List.of(mbean1);
+
+    var fetcher = Fetcher.of(List.of(fetcher1, fetcher0));
+
+    var result = fetcher.fetch(Mockito.mock(MBeanClient.class));
+
+    Assertions.assertEquals(2, result.size());
+    Assertions.assertTrue(result.contains(mbean0));
+    Assertions.assertTrue(result.contains(mbean1));
+  }
+}