com.transferwise.common.baseutils.concurrency.DiscardingQueueProcessor Maven / Gradle / Ivy

Go to download
package com.transferwise.common.baseutils.concurrency;

import com.transferwise.common.baseutils.clock.ClockHolder;
import lombok.Data;
import lombok.Setter;
import lombok.experimental.Accessors;
import lombok.extern.slf4j.Slf4j;

import java.lang.reflect.UndeclaredThrowableException;
import java.time.Duration;
import java.util.concurrent.ConcurrentLinkedQueue;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.locks.Condition;
import java.util.concurrent.locks.Lock;
import java.util.concurrent.locks.ReentrantLock;
import java.util.function.Consumer;
import java.util.function.Function;
import java.util.function.Predicate;

@Slf4j
@Accessors(chain = true)
public class DiscardingQueueProcessor {
    @Setter
    /**
     * If returns `true`, the soft limit will be applied.
     */
    private Predicate softLimitPredicate;
    @Setter
    /**
     * Transforms `T` to `K`, before element `K` is added to queue.
     */
    private Function dataTransformer;
    @Setter
    private ExecutorService executorService;
    @Setter
    private Consumer> processor;
    @Setter
    private Consumer errorProcessor;
    @Setter
    private int maxConcurrency = Runtime.getRuntime().availableProcessors();
    @Setter
    private int hardQueueLimit = 2000;
    @Setter
    private int softQueueLimit = 500;
    @Setter
    private Duration queueTimeout;
    @Setter
    private Consumer> timeoutsHander;

    private ConcurrentLinkedQueue> queue = new ConcurrentLinkedQueue<>();
    private AtomicInteger queueSize = new AtomicInteger(0);
    private AtomicInteger concurrency = new AtomicInteger(0);

    private AtomicBoolean stopRequested = new AtomicBoolean();
    private AtomicBoolean started = new AtomicBoolean();

    private Lock genericLock = new ReentrantLock();
    private Condition genericCondition = genericLock.newCondition();
    private Runnable onStop;

    public DiscardingQueueProcessor(ExecutorService executorService, Consumer> processor) {
        this.executorService = executorService;
        this.processor = processor;
    }

    public ScheduleResult schedule(T data) {
        genericLock.lock();
        try {
            if (queueSize.get() >= hardQueueLimit) {
                return new ScheduleResult().setScheduled(false).setDiscardReason(DiscardReason.HARD_LIMIT);
            } else if (queueSize.get() >= softQueueLimit) {
                if (softLimitPredicate != null && Boolean.TRUE.equals(softLimitPredicate.test(data))) {
                    return new ScheduleResult().setScheduled(false).setDiscardReason(DiscardReason.SOFT_LIMIT);
                }
            }

            K transformedData;
            if (dataTransformer != null) {
                transformedData = dataTransformer.apply(data);
            } else {
                transformedData = (K) data;
            }

            queueSize.incrementAndGet();
            Payload payload = new Payload<>();
            payload.setData(transformedData);
            queue.add(payload);
            genericCondition.signalAll();

            return new ScheduleResult().setScheduled(true);
        } finally {
            genericLock.unlock();
        }
    }

    public void start() {
        if (!started.compareAndSet(false, true)) {
            throw new IllegalStateException("Can not start. Already started.");
        }
        stopRequested.set(false);
        executorService.submit(() -> {
            AtomicBoolean shouldStop = new AtomicBoolean();
            while (!shouldStop.get()) {
                genericLock.lock();
                try {
                    while (queue.peek() == null && !stopRequested.get()) {
                        boolean ignored = genericCondition.await(5, TimeUnit.SECONDS);
                    }

                    Payload payload = queue.poll();

                    if (payload == null && stopRequested.get()) {
                        shouldStop.set(true);
                        return;
                    }
                    while (concurrency.get() >= maxConcurrency) {
                        boolean ignored = genericCondition.await(5, TimeUnit.SECONDS);
                    }

                    concurrency.incrementAndGet();
                    executorService.submit(() -> {
                        try {
                            if (queueTimeout != null && ClockHolder.getClock().millis() - payload.getSchedulingTimeMillis() > queueTimeout.toMillis()) {
                                if (timeoutsHander != null) {
                                    timeoutsHander.accept(payload);
                                }
                            } else {
                                processor.accept(payload);
                            }
                        } catch (Throwable t) {
                            onErrorRaw(t);
                        }

                        genericLock.lock();
                        try {
                            queueSize.decrementAndGet();
                            concurrency.decrementAndGet();
                            genericCondition.signalAll();
                        } finally {
                            genericLock.unlock();
                        }
                    });
                } catch (Throwable t) {
                    onErrorRaw(t);
                } finally {
                    genericLock.unlock();
                }
            }
            if (onStop != null) {
                try {
                    onStop.run();
                } catch (Throwable t) {
                    onErrorRaw(t);
                }
            }
            started.set(false);
        });
    }

    protected void onErrorRaw(Throwable t) {
        if (t instanceof UndeclaredThrowableException) {
            onError(((UndeclaredThrowableException) t).getUndeclaredThrowable());
        } else {
            onError(t);
        }
    }

    protected void onError(Throwable t) {
        if (errorProcessor == null) {
            log.error(t.getMessage(), t);
        } else {
            errorProcessor.accept(t);
        }
    }

    public void stop(Runnable onStop) {
        genericLock.lock();
        try {
            this.onStop = onStop;
            stop();
        } finally {
            genericLock.unlock();
        }
    }

    public void stop() {
        genericLock.lock();
        try {
            stopRequested.set(true);
            genericCondition.signalAll();
        } finally {
            genericLock.unlock();
        }
    }

    public boolean hasStopped() {
        genericLock.lock();
        try {
            return stopRequested.get() && queueSize.get() == 0;
        } finally {
            genericLock.unlock();
        }
    }

    public int getQueueSize() {
        return queueSize.get();
    }

    public int getConcurrency() {
        return concurrency.get();
    }

    @Data
    @Accessors(chain = true)
    public class Payload {
        private T data;
        private long schedulingTimeMillis = ClockHolder.getClock().millis();
    }

    public enum DiscardReason {
        HARD_LIMIT,
        SOFT_LIMIT
    }

    @Data
    @Accessors(chain = true)
    public static class ScheduleResult {
        private boolean scheduled;
        private DiscardReason discardReason;
    }

}