com.github.phantomthief.failover.util.ConcurrencyAware Maven / Gradle / Ivy

Go to download
package com.github.phantomthief.failover.util;

import static com.github.phantomthief.failover.util.RandomListUtils.getRandom;
import static com.google.common.base.Preconditions.checkNotNull;
import static com.google.common.collect.Multimaps.asMap;
import static com.google.common.collect.Multimaps.newListMultimap;

import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.NavigableMap;
import java.util.NoSuchElementException;
import java.util.TreeMap;
import java.util.concurrent.ConcurrentHashMap;
import java.util.function.BiFunction;

import javax.annotation.Nonnull;
import javax.annotation.Nullable;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import com.github.phantomthief.util.ThrowableConsumer;
import com.github.phantomthief.util.ThrowableFunction;
import com.google.common.collect.ListMultimap;

/**
 * @author w.vela
 * Created on 2018-01-22.
 */
public class ConcurrencyAware {

    private static final Logger logger = LoggerFactory.getLogger(ConcurrencyAware.class);

    private final Map concurrency = new ConcurrentHashMap<>();
    private final BiFunction concurrencyEvaluator;
    
    private final List> illegalStateHandlers = new ArrayList<>();

    private ConcurrencyAware(@Nonnull BiFunction concurrencyEvaluator) {
        this.concurrencyEvaluator = checkNotNull(concurrencyEvaluator);
    }

    public static  ConcurrencyAware
            create(@Nonnull BiFunction concurrencyEvaluator) {
        return new ConcurrencyAware<>(concurrencyEvaluator);
    }

    public static  ConcurrencyAware create() {
        return create((b, i) -> i);
    }

    @Nullable
    private T selectIdlest(@Nonnull Iterable candidates) {
        checkNotNull(candidates);
        ListMultimap map = null;
        for (T obj : candidates) {
            int c = concurrency.getOrDefault(obj, 0);
            if (map == null) {
                map = newListMultimap(new TreeMap<>(), ArrayList::new);
            }
            map.put(concurrencyEvaluator.apply(obj, c), obj);
        }
        if (map == null) {
            return null;
        }
        NavigableMap> asMap = (NavigableMap>) asMap(map);
        T result = getRandom(asMap.firstEntry().getValue());
        assert result != null;
        return result;
    }

    /**
     * @throws X, or {@link NoSuchElementException} if candidates is empty
     */
    public  void run(@Nonnull Iterable candidates,
            @Nonnull ThrowableConsumer func) throws X {
        checkNotNull(func);
        supply(candidates, it -> {
            func.accept(it);
            return null;
        });
    }

    /**
     * @throws X, or {@link NoSuchElementException} if candidates is empty
     */
    public  E supply(@Nonnull Iterable candidates,
            @Nonnull ThrowableFunction func) throws X {
        checkNotNull(func);
        T obj = begin(candidates);
        try {
            return func.apply(obj);
        } finally {
            end(obj);
        }
    }

    /**
     * better use {@link #supply} or {@link #run} unless need to control begin and end in special situations.
     * @throws NoSuchElementException if candidates is empty
     */
    @Nonnull
    public T begin(@Nonnull Iterable candidates) {
        T obj = beginWithoutRecordConcurrency(candidates);
        recordBeginConcurrency(obj);
        return obj;
    }

    /**
     * this is a low level api, for special purpose or mock.
     */
    @Nonnull
    public T beginWithoutRecordConcurrency(@Nonnull Iterable candidates) {
        T obj = selectIdlest(candidates);
        if (obj == null) {
            throw new NoSuchElementException();
        }
        return obj;
    }

    /**
     * this is a low level api, for special purpose or mock.
     */
    public void recordBeginConcurrency(@Nonnull T obj) {
        concurrency.merge(obj, 1, Integer::sum);
    }

    /**
     * @param obj from {@link #begin}'s return
     * @see #begin
     */
    public void end(@Nonnull T obj) {
        concurrency.compute(obj, (thisKey, oldValue) -> {
            if (oldValue == null) {
                logger.warn("illegal state found, obj:{}", thisKey);
                for (ThrowableConsumer handler : illegalStateHandlers) {
                    try {
                        handler.accept(thisKey);
                    } catch (Throwable e) {
                        logger.error("", e);
                    }
                }
                return null;
            }
            int result = oldValue - 1;
            if (result == 0) {
                return null;
            } else {
                return result;
            }
        });
    }

    public ConcurrencyAware
            addIllegalStateHandler(@Nonnull ThrowableConsumer handler) {
        illegalStateHandlers.add(checkNotNull(handler));
        return this;
    }
}