Add akka shard metrics. (#142)

vjkoskela · web-flow · commit 3f927d8c20ae · 2021-01-08T10:24:40.000-08:00
* Add akka shard metrics.

* Addressed code review feedback.
diff --git a/src/main/java/com/arpnetworking/clusteraggregator/ClusterStatusCache.java b/src/main/java/com/arpnetworking/clusteraggregator/ClusterStatusCache.java
@@ -18,21 +18,29 @@
 
 import akka.actor.AbstractActor;
 import akka.actor.ActorRef;
+import akka.actor.ActorSystem;
+import akka.actor.Address;
 import akka.actor.Cancellable;
 import akka.actor.Props;
 import akka.actor.Scheduler;
 import akka.cluster.Cluster;
 import akka.cluster.ClusterEvent;
+import akka.cluster.sharding.ClusterSharding;
+import akka.cluster.sharding.ShardRegion;
 import com.arpnetworking.clusteraggregator.models.ShardAllocation;
 import com.arpnetworking.metrics.Metrics;
 import com.arpnetworking.metrics.MetricsFactory;
+import com.arpnetworking.steno.Logger;
+import com.arpnetworking.steno.LoggerFactory;
 import com.arpnetworking.utility.ParallelLeastShardAllocationStrategy;
 import com.google.common.collect.ArrayListMultimap;
+import com.google.common.collect.Maps;
 import com.google.common.collect.Multimaps;
 import com.google.common.collect.Sets;
 import edu.umd.cs.findbugs.annotations.SuppressFBWarnings;
 import scala.compat.java8.OptionConverters;
 import scala.concurrent.duration.Duration;
+import scala.concurrent.duration.FiniteDuration;
 
 import java.io.Serializable;
 import java.util.Collection;
@@ -55,22 +63,33 @@ public class ClusterStatusCache extends AbstractActor {
     /**
      * Creates a {@link akka.actor.Props} for use in Akka.
      *
-     * @param cluster The cluster to reference.
+     *
+     * @param system The Akka {@link ActorSystem}.
+     * @param pollInterval The {@link java.time.Duration} for polling state.
      * @param metricsFactory A {@link MetricsFactory} to use for metrics creation.
      * @return A new {@link akka.actor.Props}
      */
-    public static Props props(final Cluster cluster, final MetricsFactory metricsFactory) {
-        return Props.create(ClusterStatusCache.class, cluster, metricsFactory);
+    public static Props props(
+            final ActorSystem system,
+            final java.time.Duration pollInterval,
+            final MetricsFactory metricsFactory) {
+        return Props.create(ClusterStatusCache.class, system, pollInterval, metricsFactory);
     }
 
     /**
      * Public constructor.
      *
-     * @param cluster {@link akka.cluster.Cluster} whose state is cached
+     * @param system The Akka {@link ActorSystem}.
+     * @param pollInterval The {@link java.time.Duration} for polling state.
      * @param metricsFactory A {@link MetricsFactory} to use for metrics creation.
      */
-    public ClusterStatusCache(final Cluster cluster, final MetricsFactory metricsFactory) {
-        _cluster = cluster;
+    public ClusterStatusCache(
+            final ActorSystem system,
+            final java.time.Duration pollInterval,
+            final MetricsFactory metricsFactory) {
+        _cluster = Cluster.get(system);
+        _sharding = ClusterSharding.get(system);
+        _pollInterval = pollInterval;
         _metricsFactory = metricsFactory;
     }
 
@@ -81,7 +100,7 @@ public void preStart() {
                 .scheduler();
         _pollTimer = scheduler.schedule(
                 Duration.apply(0, TimeUnit.SECONDS),
-                Duration.apply(10, TimeUnit.SECONDS),
+                Duration.apply(_pollInterval.toMillis(), TimeUnit.MILLISECONDS),
                 getSelf(),
                 POLL,
                 getContext().system().dispatcher(),
@@ -109,13 +128,50 @@ public Receive createReceive() {
                         }
                     }
                 })
+                .match(ShardRegion.ClusterShardingStats.class, shardingStats -> {
+                    LOGGER.debug()
+                            .setMessage("Received shard statistics")
+                            .addData("regionCount", shardingStats.getRegions().size())
+                            .log();
+                    final Map<String, Integer> shardsPerAddress = Maps.newHashMap();
+                    final Map<String, Long> actorsPerAddress = Maps.newHashMap();
+                    for (final Map.Entry<Address, ShardRegion.ShardRegionStats> entry : shardingStats.getRegions().entrySet()) {
+                        final String address = entry.getKey().hostPort();
+                        shardsPerAddress.put(address, entry.getValue().getStats().size());
+                        for (final Object stat : entry.getValue().getStats().values()) {
+                            if (stat instanceof Number) {
+                                final long currentActorCount = actorsPerAddress.getOrDefault(address, 0L);
+                                actorsPerAddress.put(
+                                        address,
+                                        ((Number) stat).longValue() + currentActorCount);
+                            }
+                        }
+                    }
+                    for (final Map.Entry<String, Integer> entry : shardsPerAddress.entrySet()) {
+                        try (Metrics metrics = _metricsFactory.create()) {
+                            final Long actorCount = actorsPerAddress.get(entry.getKey());
+                            metrics.addAnnotation("address", entry.getKey());
+                            metrics.setGauge("akka/cluster/shards", entry.getValue());
+                            metrics.setGauge("akka/cluster/actors", actorCount);
+                        }
+                    }
+                })
                 .match(GetRequest.class, message -> sendResponse(getSender()))
                 .match(ParallelLeastShardAllocationStrategy.RebalanceNotification.class, rebalanceNotification -> {
                     _rebalanceState = Optional.of(rebalanceNotification);
                 })
                 .matchEquals(POLL, message -> {
                     if (self().equals(sender())) {
                         _cluster.sendCurrentClusterState(getSelf());
+                        for (final String shardTypeName : _sharding.getShardTypeNames()) {
+                            LOGGER.debug()
+                                    .setMessage("Requesting shard statistics")
+                                    .addData("shardType", shardTypeName)
+                                    .log();
+                            _sharding.shardRegion(shardTypeName).tell(
+                                    new ShardRegion.GetClusterShardingStats(FiniteDuration.fromNanos(_pollInterval.toNanos())),
+                                    self());
+                        }
                     } else {
                         unhandled(message);
                     }
@@ -131,7 +187,6 @@ private void sendResponse(final ActorRef sender) {
     }
 
     private static String hostFromActorRef(final ActorRef shardRegion) {
-
         return OptionConverters.toJava(
                 shardRegion.path()
                         .address()
@@ -140,13 +195,16 @@ private static String hostFromActorRef(final ActorRef shardRegion) {
     }
 
     private final Cluster _cluster;
+    private final ClusterSharding _sharding;
+    private final java.time.Duration _pollInterval;
     private final MetricsFactory _metricsFactory;
     private Optional<ClusterEvent.CurrentClusterState> _clusterState = Optional.empty();
     @Nullable
     private Cancellable _pollTimer;
     private Optional<ParallelLeastShardAllocationStrategy.RebalanceNotification> _rebalanceState = Optional.empty();
 
     private static final String POLL = "poll";
+    private static final Logger LOGGER = LoggerFactory.getLogger(ClusterStatusCache.class);
 
     /**
      * Request to get a cluster status.
diff --git a/src/main/java/com/arpnetworking/clusteraggregator/GuiceModule.java b/src/main/java/com/arpnetworking/clusteraggregator/GuiceModule.java
@@ -245,7 +245,12 @@ private ActorRef provideStatusCache(
             @Named("periodic-statistics") final ActorRef periodicStats,
             final MetricsFactory metricsFactory) {
         final Cluster cluster = Cluster.get(system);
-        final ActorRef clusterStatusCache = system.actorOf(ClusterStatusCache.props(cluster, metricsFactory), "cluster-status");
+        final ActorRef clusterStatusCache = system.actorOf(
+                ClusterStatusCache.props(
+                        system,
+                        _configuration.getClusterStatusInterval(),
+                        metricsFactory),
+                "cluster-status");
         return system.actorOf(Status.props(cluster, clusterStatusCache, periodicStats), "status");
     }
 
diff --git a/src/main/java/com/arpnetworking/clusteraggregator/configuration/ClusterAggregatorConfiguration.java b/src/main/java/com/arpnetworking/clusteraggregator/configuration/ClusterAggregatorConfiguration.java
@@ -162,6 +162,10 @@ public boolean getCalculateClusterAggregations() {
         return _calculateClusterAggregations;
     }
 
+    public Duration getClusterStatusInterval() {
+        return _clusterStatusInterval;
+    }
+
     @Override
     public String toString() {
         return MoreObjects.toStringHelper(this)
@@ -223,6 +227,7 @@ private ClusterAggregatorConfiguration(final Builder builder) {
         _clusterHostSuffix = builder._clusterHostSuffix;
         _calculateClusterAggregations = builder._calculateClusterAggregations;
         _databaseConfigurations = Maps.newHashMap(builder._databaseConfigurations);
+        _clusterStatusInterval = builder._clusterStatusInterval;
     }
 
     private final String _monitoringCluster;
@@ -252,6 +257,7 @@ private ClusterAggregatorConfiguration(final Builder builder) {
     private final String _clusterHostSuffix;
     private final boolean _calculateClusterAggregations;
     private final Map<String, DatabaseConfiguration> _databaseConfigurations;
+    private final Duration _clusterStatusInterval;
 
     private static final ObjectMapper OBJECT_MAPPER = ObjectMapperFactory.getInstance();
 
@@ -614,6 +620,17 @@ public boolean validateAggregatorLivelinessTimeout(final Duration aggregatorLive
             return aggregatorLivelinessTimeout.compareTo(_reaggregationTimeout) > 0;
         }
 
+        /**
+         * Interval for polling cluster status. Optional. Defaults to 10 seconds.
+         *
+         * @param value interval for polling cluster status.
+         * @return This instance of {@link Builder}.
+         */
+        public Builder setClusterStatusInterval(final Duration value) {
+            _clusterStatusInterval = value;
+            return this;
+        }
+
         @NotNull
         @NotEmpty
         private String _monitoringCluster;
@@ -679,5 +696,7 @@ public boolean validateAggregatorLivelinessTimeout(final Duration aggregatorLive
         private Boolean _calculateClusterAggregations = true;
         @NotNull
         private Map<String, DatabaseConfiguration> _databaseConfigurations = Maps.newHashMap();
+        @NotNull
+        private Duration _clusterStatusInterval = Duration.ofSeconds(10);
     }
 }
diff --git a/src/test/java/com/arpnetworking/tsdcore/sinks/KairosDbSinkTest.java b/src/test/java/com/arpnetworking/tsdcore/sinks/KairosDbSinkTest.java
@@ -133,7 +133,7 @@ public void testPostFailure() throws InterruptedException, IOException {
         _kairosDbSinkBuilder.setMaximumAttempts(2).setBaseBackoff(Duration.ofMillis(1)).build()
                 .recordAggregateData(createPeriodicData(10L));
 
-        Awaitility.await().atMost(1, TimeUnit.SECONDS).untilAsserted(
+        Awaitility.await().atMost(2, TimeUnit.SECONDS).untilAsserted(
                 () -> _wireMock.verifyThat(2, WireMock.postRequestedFor(WireMock.urlEqualTo(PATH)))
         );
 
