shz.EtlHelp Maven / Gradle / Ivy

Go to download
package shz;

import java.util.LinkedList;
import java.util.List;
import java.util.Objects;
import java.util.Queue;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ConcurrentLinkedQueue;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicReference;
import java.util.function.Consumer;
import java.util.function.Predicate;

public final class EtlHelp {
    private EtlHelp() {
        throw new IllegalStateException();
    }

    /**
     * 生产及消费
     */
    public static  void execute(Consumer> runnable, Consumer> consumer, Predicate filter, int fetchSize) {
        Objects.requireNonNull(runnable);
        Objects.requireNonNull(consumer);
        int permits = fetchSize <= 0 || fetchSize >= Integer.MAX_VALUE >> 1 ? Integer.MAX_VALUE : fetchSize << 1;
        Queue queue = new ConcurrentLinkedQueue<>();
        AtomicReference latch = new AtomicReference<>();
        AtomicBoolean waitFrom = new AtomicBoolean();
        AtomicBoolean waitTo = new AtomicBoolean();
        CompletableFuture future = CompletableFuture.runAsync(() -> {
                    try {
                        runnable.accept(t -> {
                            if (filter != null && !filter.test(t)) return;
                            queue.offer(t);
                            if (waitTo.get() && queue.size() >= permits >>> 1) latch.get().countDown();
                            while (queue.size() >= permits) wait0(latch, waitFrom, waitTo);
                        });
                    } catch (Throwable t) {
                        throw PRException.of(t);
                    } finally {
                        if (waitTo.get()) latch.get().countDown();
                    }
                }
        );
        wait0(latch, waitTo, waitFrom);
        List values = new LinkedList<>();
        while (!future.isDone() && !future.isCancelled()) {
            if (waitFrom.get() && queue.size() <= permits >>> 1) latch.get().countDown();
            while (queue.size() > 0 && values.size() < fetchSize) values.add(queue.remove());
            if (values.size() > 0) {
                consumer.accept(values);
                values.clear();
            } else wait0(latch, waitTo, waitFrom);
        }
        while (queue.size() > 0) values.add(queue.remove());
        if (values.size() > 0) consumer.accept(values);
    }

    private static void wait0(AtomicReference latch, AtomicBoolean wait, AtomicBoolean down) {
        if (down.get()) latch.get().countDown();
        latch.set(new CountDownLatch(1));
        try {
            wait.set(true);
            latch.get().await();
            wait.set(false);
        } catch (InterruptedException e) {
            try {
                TimeUnit.MICROSECONDS.sleep(500L);
            } catch (InterruptedException ignored1) {
            }
        }
    }

    public static  void execute(Consumer> runnable, Consumer> consumer, Predicate filter) {
        execute(runnable, consumer, filter, 2000);
    }
}