revise BalancerHandler and add MetricSensorHandler

app/src/main/java/org/astraea/app/web/BalancerHandler.java

@@ -28,15 +28,13 @@
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.CompletionStage;
 import java.util.concurrent.ConcurrentHashMap;
-import java.util.concurrent.ConcurrentLinkedQueue;
 import java.util.function.Function;
 import java.util.function.Supplier;
 import java.util.stream.Collectors;
 import org.astraea.common.Configuration;
 import org.astraea.common.Utils;
 import org.astraea.common.admin.Admin;
 import org.astraea.common.admin.ClusterInfo;
-import org.astraea.common.admin.NodeInfo;
 import org.astraea.common.admin.Replica;
 import org.astraea.common.balancer.AlgorithmConfig;
 import org.astraea.common.balancer.Balancer;
@@ -48,8 +46,6 @@
 import org.astraea.common.cost.HasMoveCost;
 import org.astraea.common.cost.MigrationCost;
 import org.astraea.common.json.TypeRef;
-import org.astraea.common.metrics.MBeanClient;
-import org.astraea.common.metrics.collector.MetricSensor;
 import org.astraea.common.metrics.collector.MetricStore;
 
 class BalancerHandler implements Handler, AutoCloseable {
@@ -61,35 +57,12 @@ class BalancerHandler implements Handler, AutoCloseable {
       new ConcurrentHashMap<>();
   private final Map<String, CompletionStage<Void>> planExecutions = new ConcurrentHashMap<>();
 
-  private final Collection<MetricSensor> sensors = new ConcurrentLinkedQueue<>();
-
   private final MetricStore metricStore;
 
-  BalancerHandler(Admin admin, Function<Integer, Integer> jmxPortMapper) {
+  BalancerHandler(Admin admin, MetricStore metricStore) {
     this.admin = admin;
     this.balancerConsole = BalancerConsole.create(admin);
-    Supplier<CompletionStage<Map<Integer, MBeanClient>>> clientSupplier =
-        () ->
-            admin
-                .brokers()
-                .thenApply(
-                    brokers ->
-                        brokers.stream()
-                            .collect(
-                                Collectors.toUnmodifiableMap(
-                                    NodeInfo::id,
-                                    b -> MBeanClient.jndi(b.host(), jmxPortMapper.apply(b.id())))));
-    this.metricStore =
-        MetricStore.builder()
-            .beanExpiration(Duration.ofSeconds(90))
-            .localReceiver(clientSupplier)
-            .sensorsSupplier(
-                () ->
-                    sensors.stream()
-                        .collect(
-                            Collectors.toUnmodifiableMap(
-                                Function.identity(), ignored -> (id, ee) -> {})))
-            .build();
+    this.metricStore = metricStore;
   }
 
   @Override
@@ -117,9 +90,6 @@ public CompletionStage<Response> post(Channel channel) {
         Utils.construct(request.balancerClasspath, Balancer.class, request.balancerConfig);
     synchronized (this) {
       var taskId = UUID.randomUUID().toString();
-      request.algorithmConfig.clusterCostFunction().metricSensor().ifPresent(sensors::add);
-      request.algorithmConfig.moveCostFunction().metricSensor().ifPresent(sensors::add);
-
       var task =
           balancerConsole
               .launchRebalancePlanGeneration()
@@ -262,7 +232,6 @@ static PostRequestWrapper parsePostRequestWrapper(
   static class BalancerPostRequest implements Request {
 
     String balancer = GreedyBalancer.class.getName();
-
     Map<String, String> balancerConfig = Map.of();
     Map<String, String> costConfig = Map.of();
     Duration timeout = Duration.ofSeconds(3);

app/src/main/java/org/astraea/app/web/MetricSensorHandler.java

@@ -0,0 +1,77 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.astraea.app.web;
+
+import java.util.Set;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionStage;
+import java.util.stream.Collectors;
+import org.astraea.app.web.WebService.Sensors;
+import org.astraea.common.Configuration;
+import org.astraea.common.Utils;
+import org.astraea.common.cost.CostFunction;
+import org.astraea.common.json.TypeRef;
+
+public class MetricSensorHandler implements Handler {
+
+  private final Sensors sensors;
+  private static final Set<String> DEFAULT_COSTS =
+      Set.of(
+          "org.astraea.common.cost.ReplicaLeaderCost",
+          "org.astraea.common.cost.NetworkIngressCost");
+
+  MetricSensorHandler(Sensors sensors) {
+    this.sensors = sensors;
+  }
+
+  @Override
+  public CompletionStage<Response> get(Channel channel) {
+    var costs =
+        sensors.metricSensors().isEmpty()
+            ? DEFAULT_COSTS
+            : sensors.metricSensors().stream()
+                .map(x -> x.getClass().getName())
+                .collect(Collectors.toSet());
+    return CompletableFuture.completedFuture(new Response(costs));
+  }
+
+  @Override
+  public CompletionStage<Response> post(Channel channel) {
+    var metricSensorPostRequest = channel.request(TypeRef.of(MetricSensorPostRequest.class));
+    var costs = costs(metricSensorPostRequest.costs);
+    sensors.clearSensors();
+    costs.forEach(costFunction -> costFunction.metricSensor().ifPresent(sensors::addSensors));
+    return CompletableFuture.completedFuture(new Response(metricSensorPostRequest.costs));
+  }
+
+  private static Set<CostFunction> costs(Set<String> costs) {
+    if (costs.isEmpty()) throw new IllegalArgumentException("costs is not specified");
+    return Utils.costFunctions(costs, CostFunction.class, Configuration.EMPTY);
+  }
+
+  static class MetricSensorPostRequest implements Request {
+    Set<String> costs = DEFAULT_COSTS;
+  }
+
+  static class Response implements org.astraea.app.web.Response {
+    final Set<String> costs;
+
+    Response(Set<String> costs) {
+      this.costs = costs;
+    }
+  }
+}

app/src/main/java/org/astraea/app/web/WebService.java

@@ -21,22 +21,61 @@
 import com.sun.net.httpserver.HttpServer;
 import java.net.InetSocketAddress;
 import java.time.Duration;
+import java.util.Collection;
 import java.util.Map;
+import java.util.concurrent.CompletionStage;
+import java.util.concurrent.ConcurrentLinkedQueue;
 import java.util.concurrent.TimeUnit;
 import java.util.function.Function;
+import java.util.function.Supplier;
+import java.util.stream.Collectors;
 import org.astraea.app.argument.DurationField;
 import org.astraea.app.argument.IntegerMapField;
 import org.astraea.app.argument.NonNegativeIntegerField;
 import org.astraea.common.Utils;
 import org.astraea.common.admin.Admin;
+import org.astraea.common.admin.NodeInfo;
+import org.astraea.common.metrics.MBeanClient;
+import org.astraea.common.metrics.collector.MetricSensor;
+import org.astraea.common.metrics.collector.MetricStore;
 
 public class WebService implements AutoCloseable {
 
   private final HttpServer server;
   private final Admin admin;
+  private final Sensors sensors = new Sensors();
 
-  public WebService(Admin admin, int port, Function<Integer, Integer> brokerIdToJmxPort) {
+  public WebService(
+      Admin admin,
+      int port,
+      Function<Integer, Integer> brokerIdToJmxPort,
+      Duration beanExpiration) {
     this.admin = admin;
+    Supplier<CompletionStage<Map<Integer, MBeanClient>>> clientSupplier =
+        () ->
+            admin
+                .brokers()
+                .thenApply(
+                    brokers ->
+                        brokers.stream()
+                            .collect(
+                                Collectors.toUnmodifiableMap(
+                                    NodeInfo::id,
+                                    b ->
+                                        MBeanClient.jndi(
+                                            b.host(), brokerIdToJmxPort.apply(b.id())))));
+    var metricStore =
+        MetricStore.builder()
+            .beanExpiration(beanExpiration)
+            .localReceiver(clientSupplier)
+            .sensorsSupplier(
+                () ->
+                    sensors.metricSensors().stream()
+                        .distinct()
+                        .collect(
+                            Collectors.toUnmodifiableMap(
+                                Function.identity(), ignored -> (id, ee) -> {})))
+            .build();
     server = Utils.packException(() -> HttpServer.create(new InetSocketAddress(port), 0));
     server.createContext("/topics", to(new TopicHandler(admin)));
     server.createContext("/groups", to(new GroupHandler(admin)));
@@ -45,9 +84,10 @@ public WebService(Admin admin, int port, Function<Integer, Integer> brokerIdToJm
     server.createContext("/quotas", to(new QuotaHandler(admin)));
     server.createContext("/transactions", to(new TransactionHandler(admin)));
     server.createContext("/beans", to(new BeanHandler(admin, brokerIdToJmxPort)));
+    server.createContext("/metricSensors", to(new MetricSensorHandler(sensors)));
     server.createContext("/records", to(new RecordHandler(admin)));
     server.createContext("/reassignments", to(new ReassignmentHandler(admin)));
-    server.createContext("/balancer", to(new BalancerHandler(admin, brokerIdToJmxPort)));
+    server.createContext("/balancer", to(new BalancerHandler(admin, metricStore)));
     server.createContext("/throttles", to(new ThrottleHandler(admin)));
     server.start();
   }
@@ -64,7 +104,9 @@ public void close() {
 
   public static void main(String[] args) throws Exception {
     var arg = org.astraea.app.argument.Argument.parse(new Argument(), args);
-    try (var service = new WebService(Admin.of(arg.configs()), arg.port, arg::jmxPortMapping)) {
+    try (var service =
+        new WebService(
+            Admin.of(arg.configs()), arg.port, arg::jmxPortMapping, arg.beanExpiration)) {
       if (arg.ttl == null) {
         System.out.println("enter ctrl + c to terminate web service");
         TimeUnit.MILLISECONDS.sleep(Long.MAX_VALUE);
@@ -118,5 +160,32 @@ int jmxPortMapping(int brokerId) {
         validateWith = DurationField.class,
         converter = DurationField.class)
     Duration ttl = null;
+
+    @Parameter(
+        names = {"--bean.expiration"},
+        description = "Duration: the life of collected metrics",
+        validateWith = DurationField.class,
+        converter = DurationField.class)
+    Duration beanExpiration = Duration.ofHours(1);
+  }
+
+  static class Sensors {
+    private final Collection<MetricSensor> sensors;
+
+    Sensors() {
+      sensors = new ConcurrentLinkedQueue<>();
+    }
+
+    Collection<MetricSensor> metricSensors() {
+      return sensors;
+    }
+
+    void clearSensors() {
+      sensors.clear();
+    }
+
+    void addSensors(MetricSensor metricSensor) {
+      sensors.add(metricSensor);
+    }
   }
 }