All Downloads are FREE. Search and download functionalities are using the official Maven repository.

com.arpnetworking.clusteraggregator.ClusterStatusCache Maven / Gradle / Ivy

/*
 * Copyright 2014 Groupon.com
 *
 * Licensed 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 com.arpnetworking.clusteraggregator;

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 org.apache.pekko.actor.AbstractActor;
import org.apache.pekko.actor.ActorRef;
import org.apache.pekko.actor.ActorSystem;
import org.apache.pekko.actor.Address;
import org.apache.pekko.actor.Cancellable;
import org.apache.pekko.actor.Props;
import org.apache.pekko.actor.Scheduler;
import org.apache.pekko.cluster.Cluster;
import org.apache.pekko.cluster.ClusterEvent;
import org.apache.pekko.cluster.sharding.ClusterSharding;
import org.apache.pekko.cluster.sharding.ShardRegion;
import scala.concurrent.duration.Duration;
import scala.concurrent.duration.FiniteDuration;
import scala.jdk.javaapi.CollectionConverters;
import scala.jdk.javaapi.OptionConverters;

import java.io.Serializable;
import java.util.ArrayList;
import java.util.Collection;
import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.Set;
import java.util.concurrent.TimeUnit;
import java.util.stream.Collectors;
import javax.annotation.Nullable;

/**
 * Caches the cluster state.
 *
 * @author Brandon Arp (brandon dot arp at inscopemetrics dot com)
 */
public class ClusterStatusCache extends AbstractActor {

    /**
     * Creates a {@link org.apache.pekko.actor.Props} for use in Pekko.
     *
     *
     * @param system The Pekko {@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 org.apache.pekko.actor.Props}
     */
    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 system The Pekko {@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 ActorSystem system,
            final java.time.Duration pollInterval,
            final MetricsFactory metricsFactory) {
        _cluster = Cluster.get(system);
        _sharding = ClusterSharding.get(system);
        _pollInterval = pollInterval;
        _metricsFactory = metricsFactory;
    }

    @Override
    public void preStart() {
        final Scheduler scheduler = getContext()
                .system()
                .scheduler();
        _pollTimer = scheduler.scheduleAtFixedRate(
                Duration.apply(0, TimeUnit.SECONDS),
                Duration.apply(_pollInterval.toMillis(), TimeUnit.MILLISECONDS),
                getSelf(),
                POLL,
                getContext().system().dispatcher(),
                getSelf());
    }

    @Override
    public void postStop() throws Exception {
        if (_pollTimer != null) {
            _pollTimer.cancel();
        }
    }

    @Override
    public Receive createReceive() {
        return receiveBuilder()
                .match(ClusterEvent.CurrentClusterState.class, clusterState -> {
                    _clusterState = Optional.of(clusterState);
                    try (Metrics metrics = _metricsFactory.create()) {
                        metrics.setGauge("pekko/members_count", CollectionConverters.asJava(clusterState.members()).size());
                        if (_cluster.selfAddress().equals(clusterState.getLeader())) {
                            metrics.setGauge("pekko/is_leader", 1);
                        } else {
                            metrics.setGauge("pekko/is_leader", 0);
                        }
                    }
                })
                .match(ShardRegion.ClusterShardingStats.class, this::processShardingStats)
                .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());
                        }
                        int rebalanceInflight = 0;
                        int rebalancePending = 0;
                        if (_rebalanceState.isPresent()) {
                            rebalanceInflight = _rebalanceState.get().getInflightRebalances().size();
                            rebalancePending = _rebalanceState.get().getPendingRebalances().size();
                        }
                        try (Metrics metrics = _metricsFactory.create()) {
                            metrics.setGauge("pekko/cluster/rebalance/inflight", rebalanceInflight);
                            metrics.setGauge("pekko/cluster/rebalance/pending", rebalancePending);
                        }
                    } else {
                        unhandled(message);
                    }
                })
                .build();
    }

    private void processShardingStats(final ShardRegion.ClusterShardingStats shardingStats) {
        LOGGER.debug()
                .setMessage("Processing shard statistics")
                .addData("regionCount", shardingStats.getRegions().size())
                .log();
        final Map shardsPerAddress = Maps.newHashMap();
        final Map actorsPerAddress = Maps.newHashMap();
        for (final Map.Entry 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 entry : shardsPerAddress.entrySet()) {
            try (Metrics metrics = _metricsFactory.create()) {
                metrics.addAnnotation("address", entry.getKey());
                metrics.setGauge("pekko/cluster/shards", entry.getValue());
                @Nullable final Long actorCount = actorsPerAddress.get(entry.getKey());
                if (actorCount != null) {
                    metrics.setGauge("pekko/cluster/actors", actorCount);
                }
            }
        }
    }

    private void sendResponse(final ActorRef sender) {
        final StatusResponse response = new StatusResponse(
                _clusterState.orElse(_cluster.state()),
                _rebalanceState);
        sender.tell(response, self());
    }

    private static String hostFromActorRef(final ActorRef shardRegion) {
        return OptionConverters.toJava(
                shardRegion.path()
                        .address()
                        .host())
                .orElse("localhost");
    }

    private final Cluster _cluster;
    private final ClusterSharding _sharding;
    private final java.time.Duration _pollInterval;
    private final MetricsFactory _metricsFactory;
    private Optional _clusterState = Optional.empty();
    @Nullable
    private Cancellable _pollTimer;
    private Optional _rebalanceState = Optional.empty();

    private static final String POLL = "poll";
    private static final Logger LOGGER = LoggerFactory.getLogger(ClusterStatusCache.class);

    /**
     * Request to get a cluster status.
     */
    public static final class GetRequest implements Serializable {
        private static final long serialVersionUID = 2804853560963013618L;
    }

    /**
     * Response to a cluster status request.
     */
    public static final class StatusResponse implements Serializable {

        /**
         * Public constructor.
         *
         * @param clusterState the cluster state
         * @param rebalanceNotification the last rebalance data
         */
        public StatusResponse(
                @Nullable final ClusterEvent.CurrentClusterState clusterState,
                final Optional rebalanceNotification) {
            _clusterState = clusterState;

            if (rebalanceNotification.isPresent()) {
                final ParallelLeastShardAllocationStrategy.RebalanceNotification notification = rebalanceNotification.get();

                // There may be a shard joining the cluster that is not in the currentAllocations list yet, but will
                // have pending rebalances to it.  Compute the set of all shard regions by unioning the current allocation list
                // with the destinations of the rebalances.
                final Set allRefs = Sets.union(
                        notification.getCurrentAllocations().keySet(),
                        Sets.newHashSet(notification.getPendingRebalances().values()));

                final Map pendingRebalances = notification.getPendingRebalances();

                final Map> currentAllocations = notification.getCurrentAllocations();

                _allocations =
                        allRefs.stream()
                                .map(shardRegion -> computeShardAllocation(pendingRebalances, currentAllocations, shardRegion))
                                .collect(Collectors.toCollection(ArrayList::new));
            } else {
                _allocations = null;
            }
        }

        private ShardAllocation computeShardAllocation(
                final Map pendingRebalances,
                final Map> currentAllocations,
                final ActorRef shardRegion) {
            // Setup the map of current shard allocations
            final Set currentShards = currentAllocations.getOrDefault(shardRegion, Collections.emptySet());


            // Setup the list of incoming shard allocations
            final Map> invertPending = Multimaps
                    .invertFrom(Multimaps.forMap(pendingRebalances), ArrayListMultimap.create())
                    .asMap();
            final Set incomingShards = Sets.newHashSet(invertPending.getOrDefault(shardRegion, Collections.emptyList()));

            // Setup the list of outgoing shard allocations
            final Set outgoingShards = Sets.intersection(currentShards, pendingRebalances.keySet()).immutableCopy();

            // Remove the outgoing shards from the currentShards list
            currentShards.removeAll(outgoingShards);

            return new ShardAllocation.Builder()
                    .setCurrentShards(currentShards)
                    .setIncomingShards(incomingShards)
                    .setOutgoingShards(outgoingShards)
                    .setHost(hostFromActorRef(shardRegion))
                    .setShardRegion(shardRegion)
                    .build();
        }

        @Nullable
        public ClusterEvent.CurrentClusterState getClusterState() {
            return _clusterState;
        }

        public Optional> getAllocations() {
            return Optional.ofNullable(_allocations);
        }

        @Nullable
        private final ClusterEvent.CurrentClusterState _clusterState;
        @Nullable
        private final ArrayList _allocations;
        private static final long serialVersionUID = 603308359721162702L;
    }
}




© 2015 - 2025 Weber Informatics LLC | Privacy Policy