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

com.github.netty.protocol.nrpc.RpcClientChunkCompletableFuture Maven / Gradle / Ivy

The newest version!
package com.github.netty.protocol.nrpc;

import io.netty.util.concurrent.GlobalEventExecutor;
import org.reactivestreams.Subscriber;
import org.reactivestreams.Subscription;

import java.util.Collection;
import java.util.concurrent.*;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.function.*;

import static com.github.netty.protocol.nrpc.RpcClientAop.CONTEXT_LOCAL;

/**
 * support CompletableFuture async response.
 *
 * @author wangzihao
 * 2020/05/30/019
 */
public class RpcClientChunkCompletableFuture extends CompletableFuture {
    private final Collection> chunkConsumerList = new ConcurrentLinkedQueue<>();
    private final Collection> chunkIndexConsumerList = new ConcurrentLinkedQueue<>();
    private final Collection> chunkIndexAckConsumerList = new ConcurrentLinkedQueue<>();
    private final RpcMethod rpcMethod;
    private final AtomicBoolean chunkBuildEndFlag = new AtomicBoolean();
    private Executor chunkScheduler;
    private Subscription subscription;

    RpcClientChunkCompletableFuture(RpcMethod rpcMethod, RpcClientReactivePublisher source) {
        this.rpcMethod = rpcMethod;
        source.subscribe(new SubscriberAdapter(this));
    }

    public RpcClientChunkCompletableFuture request() {
        subscription.request(1);
        return this;
    }

    @Override
    public boolean cancel(boolean mayInterruptIfRunning) {
        if (subscription != null) {
            subscription.cancel();
        }
        return super.cancel(mayInterruptIfRunning);
    }

    private void chunkBuildEnd() {
        if (chunkBuildEndFlag.compareAndSet(false, true)) {
            request();
        }
    }

    /**
     * 收到chunk数据时用的线程池 (callback默认是执行在IO线程,阻塞IO线程的话,会导致超时)
     *
     * @param chunkScheduler 异步调度器
     * @return this
     */
    public RpcClientChunkCompletableFuture chunkScheduler(Executor chunkScheduler) {
        this.chunkScheduler = chunkScheduler;
        return this;
    }

    public RpcClientChunkCompletableFuture whenChunk(Consumer consumer) {
        getChunkConsumerList().add(consumer);
        return this;
    }

    public RpcClientChunkCompletableFuture whenChunk(BiConsumer consumer) {
        getChunkIndexConsumerList().add(consumer);
        return this;
    }

    public RpcClientChunkCompletableFuture whenChunk(Consumer consumer, int onIndex) {
        getChunkIndexConsumerList().add((chunk, index) -> {
            if (index == onIndex) {
                consumer.accept(chunk);
            }
        });
        return this;
    }

    public RpcClientChunkCompletableFuture whenChunk1(Consumer consumer) {
        whenChunk(consumer, 0);
        return this;
    }

    public RpcClientChunkCompletableFuture whenChunk2(Consumer consumer) {
        whenChunk(consumer, 1);
        return this;
    }

    public RpcClientChunkCompletableFuture whenChunk3(Consumer consumer) {
        whenChunk(consumer, 2);
        return this;
    }

    public RpcClientChunkCompletableFuture whenChunk4(Consumer consumer) {
        whenChunk(consumer, 3);
        return this;
    }

    public RpcClientChunkCompletableFuture whenChunk5(Consumer consumer) {
        whenChunk(consumer, 4);
        return this;
    }

    public RpcClientChunkCompletableFuture whenChunk6(Consumer consumer) {
        whenChunk(consumer, 5);
        return this;
    }

    public RpcClientChunkCompletableFuture whenChunkAck(Consumer3 consumer) {
        getChunkIndexAckConsumerList().add(consumer);
        return this;
    }

    public RpcClientChunkCompletableFuture whenChunkAck(BiConsumer consumer, int onIndex) {
        getChunkIndexAckConsumerList().add((chunk, index, ack) -> {
            if (index == onIndex) {
                consumer.accept(chunk, ack);
            }
        });
        return this;
    }

    public RpcClientChunkCompletableFuture whenChunk1Ack(BiConsumer consumer) {
        whenChunkAck(consumer, 0);
        return this;
    }

    public RpcClientChunkCompletableFuture whenChunk2Ack(BiConsumer consumer) {
        whenChunkAck(consumer, 1);
        return this;
    }

    public RpcClientChunkCompletableFuture whenChunk3Ack(BiConsumer consumer) {
        whenChunkAck(consumer, 2);
        return this;
    }

    public RpcClientChunkCompletableFuture whenChunk4Ack(BiConsumer consumer) {
        whenChunkAck(consumer, 3);
        return this;
    }

    public RpcClientChunkCompletableFuture whenChunk5Ack(BiConsumer consumer) {
        whenChunkAck(consumer, 4);
        return this;
    }

    public RpcClientChunkCompletableFuture whenChunk6Ack(BiConsumer consumer) {
        whenChunkAck(consumer, 5);
        return this;
    }

    public Collection> getChunkConsumerList() {
        return chunkConsumerList;
    }

    public Collection> getChunkIndexConsumerList() {
        return chunkIndexConsumerList;
    }

    public Collection> getChunkIndexAckConsumerList() {
        return chunkIndexAckConsumerList;
    }

    public void callbackChunkConsumerList(CHUNK chunk, int index, int chunkId, ChunkAck ack) {
        if (!existChunkCallback()) {
            ack.ack();
            return;
        }
        Executor chunkScheduler = this.chunkScheduler;
        if (chunkScheduler == null) {
            chunkScheduler = GlobalEventExecutor.INSTANCE;
        }
        RpcContext rpcContext = CONTEXT_LOCAL.get();
        chunkScheduler.execute(() -> {
            CONTEXT_LOCAL.set(rpcContext);
            try {
                // 1.chunk
                for (Consumer chunkConsumer : getChunkConsumerList()) {
                    try {
                        chunkConsumer.accept(chunk);
                    } catch (Exception e) {
                        rpcMethod.getLog().warn(rpcMethod + " chunkConsumer(chunk) exception = {}", e.toString(), e);
                    }
                }
                // 2.chunk,index
                for (BiConsumer chunkConsumer : getChunkIndexConsumerList()) {
                    try {
                        chunkConsumer.accept(chunk, index);
                    } catch (Exception e) {
                        rpcMethod.getLog().warn(rpcMethod + " chunkConsumer(chunk,index) exception = {}", e.toString(), e);
                    }
                }
                // 3.chunk,index,ack
                for (Consumer3 chunkConsumer : getChunkIndexAckConsumerList()) {
                    try {
                        chunkConsumer.accept(chunk, index, ack);
                    } catch (Exception e) {
                        rpcMethod.getLog().warn(rpcMethod + " chunkConsumer(chunk,index,ack) exception = {}", e.toString(), e);
                    }
                }
            } finally {
                Supplier chunkSupplier = () -> chunk;
                // call aop
                for (RpcClientAop aop : rpcMethod.getInstance().getAopList()) {
                    try {
                        aop.onChunkAfter(rpcContext, chunkSupplier, index, chunkId, ack);
                    } catch (Exception e) {
                        rpcMethod.getLog().warn(rpcMethod + " client.aop.onChunkAfter() exception = {}", e.toString(), e);
                    }
                }
                // ensure ack
                if (!ack.isAck()) {
                    ack.ack();
                }
                CONTEXT_LOCAL.remove();
            }
        });
    }

    public boolean existChunkCallback() {
        if (!getChunkConsumerList().isEmpty()) {
            return true;
        }
        if (!getChunkIndexConsumerList().isEmpty()) {
            return true;
        }
        return !getChunkIndexAckConsumerList().isEmpty();
    }

    @Override
    public CompletableFuture whenComplete(BiConsumer action) {
        chunkBuildEnd();
        return super.whenComplete(action);
    }

    @Override
    public CompletableFuture whenCompleteAsync(BiConsumer action) {
        chunkBuildEnd();
        return super.whenCompleteAsync(action);
    }

    @Override
    public  CompletableFuture thenApply(Function fn) {
        chunkBuildEnd();
        return super.thenApply(fn);
    }

    @Override
    public COMPLETE_RESULT get() throws InterruptedException, ExecutionException {
        chunkBuildEnd();
        return super.get();
    }

    @Override
    public COMPLETE_RESULT get(long timeout, TimeUnit unit) throws InterruptedException, ExecutionException, TimeoutException {
        chunkBuildEnd();
        return super.get(timeout, unit);
    }

    @Override
    public COMPLETE_RESULT join() {
        chunkBuildEnd();
        return super.join();
    }

    @Override
    public COMPLETE_RESULT getNow(COMPLETE_RESULT valueIfAbsent) {
        chunkBuildEnd();
        return super.getNow(valueIfAbsent);
    }

    @Override
    public  CompletableFuture thenApplyAsync(Function fn) {
        chunkBuildEnd();
        return super.thenApplyAsync(fn);
    }

    @Override
    public  CompletableFuture thenApplyAsync(Function fn, Executor executor) {
        chunkBuildEnd();
        return super.thenApplyAsync(fn, executor);
    }

    @Override
    public CompletableFuture thenAccept(Consumer action) {
        chunkBuildEnd();
        return super.thenAccept(action);
    }

    @Override
    public CompletableFuture thenAcceptAsync(Consumer action) {
        chunkBuildEnd();
        return super.thenAcceptAsync(action);
    }

    @Override
    public CompletableFuture thenAcceptAsync(Consumer action, Executor executor) {
        chunkBuildEnd();
        return super.thenAcceptAsync(action, executor);
    }

    @Override
    public CompletableFuture thenRun(Runnable action) {
        chunkBuildEnd();
        return super.thenRun(action);
    }

    @Override
    public CompletableFuture thenRunAsync(Runnable action) {
        chunkBuildEnd();
        return super.thenRunAsync(action);
    }

    @Override
    public CompletableFuture thenRunAsync(Runnable action, Executor executor) {
        chunkBuildEnd();
        return super.thenRunAsync(action, executor);
    }

    @Override
    public  CompletableFuture thenCombine(CompletionStage other, BiFunction fn) {
        chunkBuildEnd();
        return super.thenCombine(other, fn);
    }

    @Override
    public  CompletableFuture thenCombineAsync(CompletionStage other, BiFunction fn) {
        chunkBuildEnd();
        return super.thenCombineAsync(other, fn);
    }

    @Override
    public  CompletableFuture thenCombineAsync(CompletionStage other, BiFunction fn, Executor executor) {
        chunkBuildEnd();
        return super.thenCombineAsync(other, fn, executor);
    }

    @Override
    public  CompletableFuture thenAcceptBoth(CompletionStage other, BiConsumer action) {
        chunkBuildEnd();
        return super.thenAcceptBoth(other, action);
    }

    @Override
    public  CompletableFuture thenAcceptBothAsync(CompletionStage other, BiConsumer action) {
        chunkBuildEnd();
        return super.thenAcceptBothAsync(other, action);
    }

    @Override
    public  CompletableFuture thenAcceptBothAsync(CompletionStage other, BiConsumer action, Executor executor) {
        chunkBuildEnd();
        return super.thenAcceptBothAsync(other, action, executor);
    }

    @Override
    public CompletableFuture runAfterBoth(CompletionStage other, Runnable action) {
        chunkBuildEnd();
        return super.runAfterBoth(other, action);
    }

    @Override
    public CompletableFuture runAfterBothAsync(CompletionStage other, Runnable action) {
        chunkBuildEnd();
        return super.runAfterBothAsync(other, action);
    }

    @Override
    public CompletableFuture runAfterBothAsync(CompletionStage other, Runnable action, Executor executor) {
        chunkBuildEnd();
        return super.runAfterBothAsync(other, action, executor);
    }

    @Override
    public  CompletableFuture applyToEither(CompletionStage other, Function fn) {
        chunkBuildEnd();
        return super.applyToEither(other, fn);
    }

    @Override
    public  CompletableFuture applyToEitherAsync(CompletionStage other, Function fn) {
        chunkBuildEnd();
        return super.applyToEitherAsync(other, fn);
    }

    @Override
    public  CompletableFuture applyToEitherAsync(CompletionStage other, Function fn, Executor executor) {
        chunkBuildEnd();
        return super.applyToEitherAsync(other, fn, executor);
    }

    @Override
    public CompletableFuture acceptEither(CompletionStage other, Consumer action) {
        chunkBuildEnd();
        return super.acceptEither(other, action);
    }

    @Override
    public CompletableFuture acceptEitherAsync(CompletionStage other, Consumer action) {
        chunkBuildEnd();
        return super.acceptEitherAsync(other, action);
    }

    @Override
    public CompletableFuture acceptEitherAsync(CompletionStage other, Consumer action, Executor executor) {
        chunkBuildEnd();
        return super.acceptEitherAsync(other, action, executor);
    }

    @Override
    public CompletableFuture runAfterEither(CompletionStage other, Runnable action) {
        chunkBuildEnd();
        return super.runAfterEither(other, action);
    }

    @Override
    public CompletableFuture runAfterEitherAsync(CompletionStage other, Runnable action) {
        chunkBuildEnd();
        return super.runAfterEitherAsync(other, action);
    }

    @Override
    public CompletableFuture runAfterEitherAsync(CompletionStage other, Runnable action, Executor executor) {
        chunkBuildEnd();
        return super.runAfterEitherAsync(other, action, executor);
    }

    @Override
    public  CompletableFuture thenCompose(Function> fn) {
        chunkBuildEnd();
        return super.thenCompose(fn);
    }

    @Override
    public  CompletableFuture thenComposeAsync(Function> fn) {
        chunkBuildEnd();
        return super.thenComposeAsync(fn);
    }

    @Override
    public  CompletableFuture thenComposeAsync(Function> fn, Executor executor) {
        chunkBuildEnd();
        return super.thenComposeAsync(fn, executor);
    }

    @Override
    public CompletableFuture whenCompleteAsync(BiConsumer action, Executor executor) {
        chunkBuildEnd();
        return super.whenCompleteAsync(action, executor);
    }

    @Override
    public  CompletableFuture handle(BiFunction fn) {
        chunkBuildEnd();
        return super.handle(fn);
    }

    @Override
    public  CompletableFuture handleAsync(BiFunction fn) {
        chunkBuildEnd();
        return super.handleAsync(fn);
    }

    @Override
    public  CompletableFuture handleAsync(BiFunction fn, Executor executor) {
        chunkBuildEnd();
        return super.handleAsync(fn, executor);
    }

    @Override
    public CompletableFuture toCompletableFuture() {
        chunkBuildEnd();
        return super.toCompletableFuture();
    }

    @Override
    public CompletableFuture exceptionally(Function fn) {
        chunkBuildEnd();
        return super.exceptionally(fn);
    }

    @Override
    public boolean isDone() {
        chunkBuildEnd();
        return super.isDone();
    }

    @Override
    public boolean isCancelled() {
        chunkBuildEnd();
        return super.isCancelled();
    }

    @Override
    public boolean isCompletedExceptionally() {
        chunkBuildEnd();
        return super.isCompletedExceptionally();
    }

    @FunctionalInterface
    public interface Consumer3 {
        void accept(T1 t1, T2 t2, T3 t3);
    }

    public static class SubscriberAdapter implements Subscriber, RpcDone.ChunkListener {
        private final RpcClientChunkCompletableFuture completableFuture;
        private final AtomicInteger chunkIndex = new AtomicInteger();
        private RESULT result;
        private Throwable throwable;

        private SubscriberAdapter(RpcClientChunkCompletableFuture completableFuture) {
            this.completableFuture = completableFuture;
        }

        @Override
        public void onSubscribe(Subscription s) {
            this.completableFuture.subscription = s;
        }

        @Override
        public void onChunk(CHUNK chunk, int chunkId, ChunkAck ack) {
            completableFuture.callbackChunkConsumerList(chunk, chunkIndex.getAndIncrement(), chunkId, ack);
        }

        @Override
        public void onNext(RESULT o) {
            this.result = o;
        }

        @Override
        public void onError(Throwable t) {
            this.throwable = t;
        }

        @Override
        public void onComplete() {
            Throwable throwable = this.throwable;
            RESULT result = this.result;
            this.throwable = null;
            this.result = null;
            if (throwable != null) {
                completableFuture.completeExceptionally(throwable);
            } else {
                completableFuture.complete(result);
            }
        }
    }
}