org.wildfly.clustering.infinispan.affinity.impl.DefaultKeyAffinityService Maven / Gradle / Ivy

Go to download
/*
 * Copyright The WildFly Authors
 * SPDX-License-Identifier: Apache-2.0
 */

package org.wildfly.clustering.infinispan.affinity.impl;

import java.time.Duration;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.TreeMap;
import java.util.concurrent.ArrayBlockingQueue;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.CompletionStage;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.Future;
import java.util.concurrent.RejectedExecutionException;
import java.util.concurrent.ThreadFactory;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicReference;
import java.util.function.Predicate;
import java.util.function.Supplier;

import org.infinispan.Cache;
import org.infinispan.affinity.KeyAffinityService;
import org.infinispan.affinity.KeyGenerator;
import org.infinispan.distribution.ch.ConsistentHash;
import org.infinispan.distribution.ch.KeyPartitioner;
import org.infinispan.notifications.Listener;
import org.infinispan.notifications.Listener.Observation;
import org.infinispan.notifications.cachelistener.annotation.TopologyChanged;
import org.infinispan.notifications.cachelistener.event.TopologyChangedEvent;
import org.infinispan.remoting.transport.Address;
import org.jboss.logging.Logger;
import org.wildfly.clustering.context.DefaultExecutorService;
import org.wildfly.clustering.context.DefaultThreadFactory;
import org.wildfly.clustering.infinispan.distribution.ConsistentHashKeyDistribution;
import org.wildfly.clustering.infinispan.distribution.KeyDistribution;
import org.wildfly.security.manager.WildFlySecurityManager;

/**
 * A custom key affinity service implementation with the following distinct characteristics (as compared to {@link org.infinispan.affinity.impl.KeyAffinityServiceImpl}):
 * 
 * {@link #getKeyForAddress(Address)} will return a random key (instead of throwing an ISE) if the specified address does not own any segments.
 * Uses a worker thread per address for which to generate keys.
 * Minimal CPU utilization when key queues are full.
 * Non-blocking topology change event handler.
 * {@link #getKeyForAddress(Address)} calls will not block during topology change events.
 * 
 * @author Paul Ferraro
 */
@Listener(observation = Observation.POST)
public class DefaultKeyAffinityService implements KeyAffinityService, Supplier> {

    static final int DEFAULT_QUEUE_SIZE = 100;
    private static final Logger LOGGER = Logger.getLogger(DefaultKeyAffinityService.class);
    private static final ThreadFactory THREAD_FACTORY = new DefaultThreadFactory(DefaultKeyAffinityService.class);

    private final Cache cache;
    private final KeyGenerator generator;
    private final AtomicReference> currentState = new AtomicReference<>();
    private final KeyPartitioner partitioner;
    private final Predicate filter;

    private volatile int queueSize = DEFAULT_QUEUE_SIZE;
    private volatile Duration timeout = Duration.ofMillis(100L);
    private volatile ExecutorService executor;

    private interface KeyAffinityState {
        KeyDistribution getDistribution();
        KeyRegistry getRegistry();
        Iterable> getFutures();
    }

    /**
     * Constructs a key affinity service that generates keys hashing to the members matching the specified filter.
     * @param cache the target cache
     * @param generator a key generator
     */
    @SuppressWarnings("deprecation")
    public DefaultKeyAffinityService(Cache cache, KeyGenerator generator, Predicate filter) {
        this(cache, cache.getAdvancedCache().getComponentRegistry().getLocalComponent(KeyPartitioner.class), generator, filter);
    }

    DefaultKeyAffinityService(Cache cache, KeyPartitioner partitioner, KeyGenerator generator, Predicate filter) {
        this.cache = cache;
        this.partitioner = partitioner;
        this.generator = generator;
        this.filter = filter;
    }

    /**
     * Overrides the maximum number of keys with affinity to a given member to pre-generate.
     * @param size a queue size threshold
     */
    public void setQueueSize(int size) {
        this.queueSize = size;
    }

    /**
     * Overrides the duration of time for which calls to {@link #getKeyForAddress(Address)} will wait for an available pre-generated key,
     * after which a random key will be returned.
     * @param timeout a queue poll timeout
     */
    public void setPollTimeout(Duration timeout) {
        this.timeout = timeout;
    }

    @Override
    public BlockingQueue get() {
        return new ArrayBlockingQueue<>(this.queueSize);
    }

    @Override
    public boolean isStarted() {
        ExecutorService executor = this.executor;
        return (executor != null) && !executor.isShutdown();
    }

    @Override
    public void start() {
        this.executor = Executors.newCachedThreadPool(THREAD_FACTORY);
        this.accept(this.cache.getAdvancedCache().getDistributionManager().getCacheTopology().getWriteConsistentHash());
        this.cache.addListener(this);
    }

    @Override
    public void stop() {
        this.cache.removeListener(this);
        WildFlySecurityManager.doUnchecked(this.executor, DefaultExecutorService.SHUTDOWN_NOW_ACTION);
    }

    @Override
    public K getCollocatedKey(K otherKey) {
        KeyAffinityState currentState = this.currentState.get();
        if (currentState == null) {
            // Not yet started!
            throw new IllegalStateException();
        }
        return this.getCollocatedKey(currentState, otherKey);
    }

    private K getCollocatedKey(KeyAffinityState state, K otherKey) {
        K key = this.poll(state.getRegistry(), state.getDistribution().getPrimaryOwner(otherKey));
        if (key != null) {
            return key;
        }
        KeyAffinityState currentState = this.currentState.get();
        // If state is out-dated, retry
        if (state != currentState) {
            return this.getCollocatedKey(currentState, otherKey);
        }
        LOGGER.debugf("Could not obtain pre-generated key with same affinity as %s -- generating random key", otherKey);
        return this.generator.getKey();
    }

    @Override
    public K getKeyForAddress(Address address) {
        if (!this.filter.test(address)) {
            throw new IllegalArgumentException(address.toString());
        }
        KeyAffinityState currentState = this.currentState.get();
        if (currentState == null) {
            // Not yet started!
            throw new IllegalStateException();
        }
        return this.getKeyForAddress(currentState, address);
    }

    private K getKeyForAddress(KeyAffinityState state, Address address) {
        K key = this.poll(state.getRegistry(), address);
        if (key != null) {
            return key;
        }
        KeyAffinityState currentState = this.currentState.get();
        // If state is out-dated, retry
        if (state != currentState) {
            return this.getKeyForAddress(currentState, address);
        }
        LOGGER.debugf("Could not obtain pre-generated key with affinity for %s -- generating random key", address);
        return this.generator.getKey();
    }

    private K poll(KeyRegistry registry, Address address) {
        BlockingQueue keys = registry.getKeys(address);
        if (keys != null) {
            Duration timeout = this.timeout;
            long nanos = (timeout.getSeconds() == 0) ? timeout.getNano() : timeout.toNanos();
            try {
                return keys.poll(nanos, TimeUnit.NANOSECONDS);
            } catch (InterruptedException e) {
                Thread.currentThread().interrupt();
            }
        }
        return null;
    }

    @TopologyChanged
    public CompletionStage topologyChanged(TopologyChangedEvent event) {
        if (!this.getSegments(event.getWriteConsistentHashAtStart()).equals(this.getSegments(event.getWriteConsistentHashAtEnd()))) {
            LOGGER.debugf("Restarting key generation based on new consistent hash for topology %d", event.getNewTopologyId());
            this.accept(event.getWriteConsistentHashAtEnd());
        }
        return CompletableFuture.completedStage(null);
    }

    private Map> getSegments(ConsistentHash hash) {
        Map> segments = new TreeMap<>();
        for (Address address : hash.getMembers()) {
            if (this.filter.test(address)) {
                segments.put(address, hash.getPrimarySegmentsForOwner(address));
            }
        }
        return segments;
    }

    private void accept(ConsistentHash hash) {
        KeyDistribution distribution = new ConsistentHashKeyDistribution(this.partitioner, hash);
        KeyRegistry registry = new ConsistentHashKeyRegistry<>(hash, this.filter, this);
        Set addresses = registry.getAddresses();
        List> futures = !addresses.isEmpty() ? new ArrayList<>(addresses.size()) : Collections.emptyList();
        try {
            for (Address address : addresses) {
                BlockingQueue keys = registry.getKeys(address);
                futures.add(this.executor.submit(new GenerateKeysTask<>(this.generator, distribution, address, keys)));
            }
            KeyAffinityState previousState = this.currentState.getAndSet(new KeyAffinityState() {
                @Override
                public KeyDistribution getDistribution() {
                    return distribution;
                }

                @Override
                public KeyRegistry getRegistry() {
                    return registry;
                }

                @Override
                public Iterable> getFutures() {
                    return futures;
                }
            });
            if (previousState != null) {
                for (Future future : previousState.getFutures()) {
                    future.cancel(true);
                }
            }
        } catch (RejectedExecutionException e) {
            // Executor was shutdown. Cancel any tasks that were already submitted
            for (Future future : futures) {
                future.cancel(true);
            }
        }
    }

    private static class GenerateKeysTask implements Runnable {
        private final KeyGenerator generator;
        private final KeyDistribution distribution;
        private final Address address;
        private final BlockingQueue keys;

        GenerateKeysTask(KeyGenerator generator, KeyDistribution distribution, Address address, BlockingQueue keys) {
            this.generator = generator;
            this.distribution = distribution;
            this.address = address;
            this.keys = keys;
        }

        @Override
        public void run() {
            while (!Thread.currentThread().isInterrupted()) {
                K key = this.generator.getKey();
                if (this.distribution.getPrimaryOwner(key).equals(this.address)) {
                    try {
                        this.keys.put(key);
                    } catch (InterruptedException e) {
                        Thread.currentThread().interrupt();
                    }
                }
            }
        }
    }
}