[METRIC] Add wait in MetricStore interface

common/src/main/java/org/astraea/common/metrics/collector/MetricStore.java

@@ -19,6 +19,7 @@
 import java.time.Duration;
 import java.util.ArrayList;
 import java.util.Collection;
+import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
@@ -27,18 +28,21 @@
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ConcurrentLinkedQueue;
 import java.util.concurrent.ConcurrentSkipListSet;
+import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.function.BiConsumer;
+import java.util.function.Predicate;
 import java.util.function.Supplier;
 import java.util.stream.Collectors;
 import org.astraea.common.Utils;
 import org.astraea.common.consumer.Consumer;
 import org.astraea.common.consumer.ConsumerConfigs;
 import org.astraea.common.consumer.Deserializer;
 import org.astraea.common.consumer.Record;
+import org.astraea.common.cost.NoSufficientMetricsException;
 import org.astraea.common.metrics.BeanObject;
 import org.astraea.common.metrics.BeanQuery;
 import org.astraea.common.metrics.ClusterBean;
@@ -67,6 +71,21 @@ static Builder builder() {
    */
   Map<MetricSensor, BiConsumer<Integer, Exception>> sensors();
 
+  /** Wait for the checker to be true or timeout. */
+  default void wait(Predicate<ClusterBean> checker, Duration timeout) {
+    // Keep checking the checker until it is true or timeout.
+    var start = System.currentTimeMillis();
+    while (System.currentTimeMillis() - start < timeout.toMillis()) {
+      try {
+        if (checker.test(clusterBean())) return;
+      } catch (NoSufficientMetricsException noe) {
+        // keep trying
+      }
+      Utils.sleep(Duration.ofSeconds(1));
+    }
+    throw new IllegalStateException("Timeout waiting for the checker");
+  }
+
   @Override
   void close();
 
@@ -193,6 +212,7 @@ class MetricStoreImpl implements MetricStore {
     private final Set<Integer> identities = new ConcurrentSkipListSet<>();
 
     private volatile Map<MetricSensor, BiConsumer<Integer, Exception>> lastSensors = Map.of();
+    private final Map<CountDownLatch, Predicate<ClusterBean>> waitingList = new HashMap<>();
 
     private MetricStoreImpl(
         Supplier<Map<MetricSensor, BiConsumer<Integer, Exception>>> sensorsSupplier,
@@ -242,8 +262,11 @@ private MetricStoreImpl(
                             }
                           });
                     });
-                // generate new cluster bean
-                if (!allBeans.isEmpty()) updateClusterBean();
+                if (!allBeans.isEmpty()) {
+                  // generate new cluster bean
+                  updateClusterBean();
+                  checkWaitingList();
+                }
               } catch (Exception e) {
                 // TODO: it needs better error handling
                 e.printStackTrace();
@@ -277,6 +300,24 @@ public void close() {
       receiver.close();
     }
 
+    /** User thread will "wait" until being awakened by the metric store or being timeout. */
+    @Override
+    public void wait(Predicate<ClusterBean> checker, Duration timeout) {
+      if (checker.test(clusterBean())) return;
+
+      var latch = new CountDownLatch(1);
+      try {
+        waitingList.put(latch, checker);
+        if (!latch.await(timeout.toMillis(), TimeUnit.MILLISECONDS)) {
+          throw new IllegalStateException("Timeout waiting for the checker");
+        }
+      } catch (InterruptedException ie) {
+        throw new IllegalStateException("Interrupted while waiting for the checker");
+      } finally {
+        waitingList.remove(latch);
+      }
+    }
+
     private void updateClusterBean() {
       lastClusterBean =
           ClusterBean.of(
@@ -286,5 +327,15 @@ private void updateClusterBean() {
                       Collectors.toUnmodifiableMap(
                           Map.Entry::getKey, e -> List.copyOf(e.getValue()))));
     }
+
+    /** Check the checkers in the waiting list. If the checker returns true, awake the thread. */
+    private void checkWaitingList() {
+      waitingList.forEach(
+          (latch, checker) -> {
+            if (checker.test(clusterBean())) {
+              latch.countDown();
+            }
+          });
+    }
   }
 }

common/src/test/java/org/astraea/common/metrics/collector/MetricStoreTest.java

@@ -80,4 +80,32 @@ void testBeanExpiration() {
       Assertions.assertEquals(0, store.clusterBean().all().size());
     }
   }
+
+  @Test
+  void testWait() {
+    var queue = new LinkedBlockingQueue<Map<Integer, Collection<BeanObject>>>();
+
+    try (var store =
+        MetricStore.builder()
+            .receiver(timeout -> Utils.packException(queue::take))
+            .sensorsSupplier(
+                // Metric sensor provide fake hasBeanObject
+                () ->
+                    Map.of(
+                        (client, bean) ->
+                            List.of(() -> new BeanObject(Utils.randomString(), Map.of(), Map.of())),
+                        (id, exception) -> {}))
+            .build()) {
+      Assertions.assertThrows(
+          IllegalStateException.class, () -> store.wait((ignore) -> false, Duration.ofSeconds(1)));
+      Assertions.assertDoesNotThrow(() -> store.wait((ignore) -> true, Duration.ofSeconds(1)));
+
+      Assertions.assertThrows(
+          IllegalStateException.class,
+          () -> store.wait((clusterBean) -> !clusterBean.all().isEmpty(), Duration.ofSeconds(1)));
+      queue.add(Map.of(1000, List.of(new BeanObject(Utils.randomString(), Map.of(), Map.of()))));
+      Assertions.assertDoesNotThrow(
+          () -> store.wait((clusterBean) -> !clusterBean.all().isEmpty(), Duration.ofSeconds(1)));
+    }
+  }
 }