tech.ydb.core.grpc.ServerStreamToObserver Maven / Gradle / Ivy

Go to download
package tech.ydb.core.grpc;

import java.util.function.Consumer;

import javax.annotation.Nullable;

import io.grpc.ClientCall;
import io.grpc.Metadata;
import io.grpc.Status;
import io.grpc.stub.ClientCallStreamObserver;

import tech.ydb.core.rpc.StreamObserver;


/**
 * @author Sergey Polovko
 * @param  type of request
 * @param  type of response
 */
public class ServerStreamToObserver extends ClientCall.Listener {

    private final StreamObserver observer;
    private final CallToStreamObserverAdapter adapter;
    private final Consumer trailersHandler;
    private final Consumer statusHandler;

    public ServerStreamToObserver(StreamObserver observer,
            ClientCall call,
            Consumer trailersHandler,
            Consumer statusHandler) {
        this.observer = observer;
        this.adapter = new CallToStreamObserverAdapter<>(call);
        this.trailersHandler = trailersHandler;
        this.statusHandler = statusHandler;
    }

    @Override
    public void onHeaders(Metadata headers) {
    }

    @Override
    public void onMessage(RespT message) {
        try {
            observer.onNext(message);
            // request delivery of the next inbound message.
            adapter.request(1);
        } catch (Exception ex) {
            adapter.cancel(ex.getMessage(), ex);
        }
    }

    @Override
    public void onClose(Status status, @Nullable Metadata trailers) {
        if (trailersHandler != null && trailers != null) {
            trailersHandler.accept(trailers);
        }
        if (statusHandler != null) {
            statusHandler.accept(status);
        }

        if (status.isOk()) {
            observer.onCompleted();
        } else {
            observer.onError(GrpcStatuses.toStatus(status));
        }
    }

    @Override
    public void onReady() {
    }

    private static final class CallToStreamObserverAdapter extends ClientCallStreamObserver {
        private final ClientCall call;

        CallToStreamObserverAdapter(ClientCall call) {
            this.call = call;
        }

        @Override
        public void onNext(T value) {
            call.sendMessage(value);
        }

        @Override
        public void onError(Throwable t) {
            call.cancel("Cancelled by client with StreamObserver.onError()", t);
        }

        @Override
        public void onCompleted() {
            call.halfClose();
        }

        @Override
        public boolean isReady() {
            return call.isReady();
        }

        @Override
        public void setOnReadyHandler(Runnable onReadyHandler) {
            throw new IllegalStateException("Cannot alter onReadyHandler after call started");
        }

        @Override
        public void disableAutoInboundFlowControl() {
            throw new IllegalStateException("Cannot disable auto flow control call started");
        }

        @Override
        public void request(int count) {
            call.request(count);
        }

        @Override
        public void setMessageCompression(boolean enable) {
            call.setMessageCompression(enable);
        }

        @Override
        public void cancel(@Nullable String message, @Nullable Throwable cause) {
            call.cancel(message, cause);
        }
    }
}