io.grpc.stub.CallStreamObserver Java Examples

/**
 * Implements a bidirectional stream → stream call as {@link Flux} → {@link Flux}, where both the client
 * and the server independently stream to each other.
 */
@SuppressWarnings("unchecked")
public static <TRequest, TResponse> Flux<TResponse> manyToMany(
        Flux<TRequest> fluxSource,
        Function<StreamObserver<TResponse>, StreamObserver<TRequest>> delegate,
        CallOptions options) {
    try {

        final int prefetch = ReactorCallOptions.getPrefetch(options);
        final int lowTide = ReactorCallOptions.getLowTide(options);

        ReactorSubscriberAndClientProducer<TRequest> subscriberAndGRPCProducer =
            fluxSource.subscribeWith(new ReactorSubscriberAndClientProducer<>());
        ReactorClientStreamObserverAndPublisher<TResponse> observerAndPublisher =
            new ReactorClientStreamObserverAndPublisher<>(
                s -> subscriberAndGRPCProducer.subscribe((CallStreamObserver<TRequest>) s),
                subscriberAndGRPCProducer::cancel, prefetch, lowTide
            );
        delegate.apply(observerAndPublisher);

        return Flux.from(observerAndPublisher);
    } catch (Throwable throwable) {
        return Flux.error(throwable);
    }
}
 

/**
 * Implements a stream → unary call as {@link Flux} → {@link Mono}, where the client transits a stream of
 * messages.
 */
@SuppressWarnings("unchecked")
public static <TRequest, TResponse> Mono<TResponse> manyToOne(
        Flux<TRequest> fluxSource,
        Function<StreamObserver<TResponse>, StreamObserver<TRequest>> delegate,
        CallOptions options) {
    try {
        ReactorSubscriberAndClientProducer<TRequest> subscriberAndGRPCProducer =
                fluxSource.subscribeWith(new ReactorSubscriberAndClientProducer<>());
        ReactorClientStreamObserverAndPublisher<TResponse> observerAndPublisher =
                new ReactorClientStreamObserverAndPublisher<>(
                    s -> subscriberAndGRPCProducer.subscribe((CallStreamObserver<TRequest>) s),
                    subscriberAndGRPCProducer::cancel
                );
        delegate.apply(observerAndPublisher);

        return Flux.from(observerAndPublisher)
                   .singleOrEmpty();
    } catch (Throwable throwable) {
        return Mono.error(throwable);
    }
}
 

@Test
public void statusRuntimeExceptionTriggersHandler() {
    CallStreamObserver delegate = mock(CallStreamObserver.class);
    final AtomicBoolean called = new AtomicBoolean(false);

    AbstractStreamObserverAndPublisher observer = new AbstractStreamObserverAndPublisher(new ArrayBlockingQueue(1), null, new Runnable() {
        @Override
        public void run() {
            called.set(true);
        }
    }) { };

    observer.onSubscribe(delegate);

    TestSubscriber test = Flowable.fromPublisher(observer)
                                  .test();

    StatusRuntimeException exception = Status.CANCELLED.asRuntimeException();
    observer.onError(exception);

    test.awaitTerminalEvent();
    test.assertError(exception);

    assertThat(called.get()).isTrue();
    assertThat(observer.outputFused).isFalse();
}
 

@Test
public void statusExceptionTriggersHandler() {
    CallStreamObserver delegate = mock(CallStreamObserver.class);
    final AtomicBoolean called = new AtomicBoolean(false);

    AbstractStreamObserverAndPublisher observer = new AbstractStreamObserverAndPublisher(new ArrayBlockingQueue(1), null, new Runnable() {
        @Override
        public void run() {
            called.set(true);
        }
    }) { };

    observer.onSubscribe(delegate);

    TestSubscriber test = Flowable.fromPublisher(observer)
                                  .test();

    StatusException exception = Status.CANCELLED.asException();
    observer.onError(exception);

    test.awaitTerminalEvent();
    test.assertError(exception);

    assertThat(called.get()).isTrue();
    assertThat(observer.outputFused).isFalse();
}
 

@Override
public void onNext(T t) {
    if (sourceMode == ASYNC || sourceMode == NOT_FUSED) {
        drain();
        return;
    }

    if (!isCanceled()) {
        checkNotNull(t);

        final CallStreamObserver<T> subscriber = downstream;

        try {
            subscriber.onNext(t);
            isRequested = false;
            drain();
        } catch (Throwable throwable) {
            cancel();
            try {
                subscriber.onError(prepareError(throwable));
            } catch (Throwable ignore) { }
        }
    }
}
 

public AbstractServerStreamObserverAndPublisher(
    ServerCallStreamObserver<?> serverCallStreamObserver,
    Queue<T> queue,
    Consumer<CallStreamObserver<?>> onSubscribe) {
    super(queue, onSubscribe);
    super.onSubscribe(serverCallStreamObserver);
}
 

BlockingStreamObserver(CallStreamObserver<T> delegate) {
  this.delegate = delegate;
  final Runnable notifyAll = () -> {
    synchronized (lock) {
      lock.notifyAll(); // wake up our thread
    }
  };
  this.delegate.setOnReadyHandler(notifyAll);
  if (delegate instanceof ServerCallStreamObserver) {
    ((ServerCallStreamObserver<T>) delegate).setOnCancelHandler(notifyAll);
  }
}
 

ReactorClientStreamObserverAndPublisher(
        Consumer<CallStreamObserver<?>> onSubscribe,
        Runnable onTerminate,
        int prefetch,
        int lowTide) {
    super(Queues.<T>get(DEFAULT_CHUNK_SIZE).get(), onSubscribe, onTerminate, prefetch, lowTide);
}
 

ReactorServerStreamObserverAndPublisher(
        ServerCallStreamObserver<?> serverCallStreamObserver,
        Consumer<CallStreamObserver<?>> onSubscribe,
        int prefetch,
        int lowTide) {
    super(serverCallStreamObserver, Queues.<T>get(DEFAULT_CHUNK_SIZE).get(), onSubscribe, prefetch, lowTide);
}
 

@Test
public void statusRuntimeExceptionTriggersHandlerFuseable() {
    CallStreamObserver delegate = mock(CallStreamObserver.class);
    final AtomicBoolean called = new AtomicBoolean(false);

    AbstractStreamObserverAndPublisher observer = new TestStreamObserverAndPublisherWithFusion(new ArrayBlockingQueue(1), null, new Runnable() {
        @Override
        public void run() {
            called.set(true);
        }
    });

    observer.onSubscribe(delegate);

    TestSubscriber test = Flowable.fromPublisher(observer)
                                  .observeOn(Schedulers.trampoline())
                                  .test();

    StatusRuntimeException exception = Status.CANCELLED.asRuntimeException();
    observer.onError(exception);

    test.awaitTerminalEvent();
    test.assertError(exception);

    assertThat(called.get()).isTrue();
    assertThat(observer.outputFused).isTrue();
}
 

@Test
public void statusExceptionTriggersHandlerFuseable() {
    CallStreamObserver delegate = mock(CallStreamObserver.class);
    final AtomicBoolean called = new AtomicBoolean(false);

    AbstractStreamObserverAndPublisher observer = new TestStreamObserverAndPublisherWithFusion(new ArrayBlockingQueue(1), null, new Runnable() {
        @Override
        public void run() {
            called.set(true);
        }
    });

    observer.onSubscribe(delegate);

    TestSubscriber test = Flowable.fromPublisher(observer)
                                  .observeOn(Schedulers.trampoline())
                                  .test();

    StatusException exception = Status.CANCELLED.asException();
    observer.onError(exception);

    test.awaitTerminalEvent();
    test.assertError(exception);

    assertThat(called.get()).isTrue();
    assertThat(observer.outputFused).isTrue();
}
 

public AbstractServerStreamObserverAndPublisher(
        ServerCallStreamObserver<?> serverCallStreamObserver,
        Queue<T> queue,
        Consumer<CallStreamObserver<?>> onSubscribe,
        int prefetch,
        int lowTide) {
    super(queue, prefetch, lowTide, onSubscribe);
    super.onSubscribe(serverCallStreamObserver);
}
 

/**
 * Implements a stream → unary call as {@link Flowable} → {@link Single}, where the client transits a stream of
 * messages.
 */
@SuppressWarnings("unchecked")
public static <TRequest, TResponse> Single<TResponse> manyToOne(
        final Flowable<TRequest> flowableSource,
        final Function<StreamObserver<TResponse>, StreamObserver<TRequest>> delegate,
        final CallOptions options) {
    try {
        final RxSubscriberAndClientProducer<TRequest> subscriberAndGRPCProducer =
                flowableSource.subscribeWith(new RxSubscriberAndClientProducer<TRequest>());
        final RxClientStreamObserverAndPublisher<TResponse> observerAndPublisher =
            new RxClientStreamObserverAndPublisher<TResponse>(
                new com.salesforce.reactivegrpc.common.Consumer<CallStreamObserver<?>>() {
                    @Override
                    public void accept(CallStreamObserver<?> observer) {
                        subscriberAndGRPCProducer.subscribe((CallStreamObserver<TRequest>) observer);
                    }
                },
                new Runnable() {
                    @Override
                    public void run() {
                        subscriberAndGRPCProducer.cancel();
                    }
                }
            );
        delegate.apply(observerAndPublisher);

        return Flowable.fromPublisher(observerAndPublisher)
                       .singleOrError();
    } catch (Throwable throwable) {
        return Single.error(throwable);
    }
}
 

@Override
public void subscribe(CallStreamObserver<T> downstream) {
    super.subscribe(downstream);
    ((ServerCallStreamObserver<?>) downstream).setOnCancelHandler(new Runnable() {
        @Override
        public void run() {
            AbstractSubscriberAndServerProducer.super.cancel();
        }
    });
}
 

public AbstractClientStreamObserverAndPublisher(
        Queue<T> queue,
        Consumer<CallStreamObserver<?>> onSubscribe,
        Runnable onTerminate,
        int prefetch,
        int lowTide) {
    super(queue, prefetch, lowTide, onSubscribe, onTerminate);
}
 

/**
 * Implements a bidirectional stream → stream call as {@link Flowable} → {@link Flowable}, where both the client
 * and the server independently stream to each other.
 */
@SuppressWarnings("unchecked")
public static <TRequest, TResponse> Flowable<TResponse> manyToMany(
        final Flowable<TRequest> flowableSource,
        final Function<StreamObserver<TResponse>, StreamObserver<TRequest>> delegate,
        final CallOptions options) {

    final int prefetch = RxCallOptions.getPrefetch(options);
    final int lowTide = RxCallOptions.getLowTide(options);

    try {
        final RxSubscriberAndClientProducer<TRequest> subscriberAndGRPCProducer =
                flowableSource.subscribeWith(new RxSubscriberAndClientProducer<TRequest>());
        final RxClientStreamObserverAndPublisher<TResponse> observerAndPublisher =
            new RxClientStreamObserverAndPublisher<TResponse>(
                new com.salesforce.reactivegrpc.common.Consumer<CallStreamObserver<?>>() {
                    @Override
                    public void accept(CallStreamObserver<?> observer) {
                        subscriberAndGRPCProducer.subscribe((CallStreamObserver<TRequest>) observer);
                    }
                },
                new Runnable() {
                    @Override
                    public void run() {
                        subscriberAndGRPCProducer.cancel();
                    }
                },
                prefetch, lowTide);
        delegate.apply(observerAndPublisher);

        return Flowable.fromPublisher(observerAndPublisher);
    } catch (Throwable throwable) {
        return Flowable.error(throwable);
    }
}
 

RxClientStreamObserverAndPublisher(
        Consumer<CallStreamObserver<?>> onSubscribe,
        Runnable onTerminate,
        int prefetch,
        int lowTide) {
    super(new SimpleQueueAdapter<T>(new SpscArrayQueue<T>(prefetch)), onSubscribe, onTerminate, prefetch, lowTide);
}
 

RxServerStreamObserverAndPublisher(
        ServerCallStreamObserver<?> serverCallStreamObserver,
        Consumer<CallStreamObserver<?>> onSubscribe,
        int prefetch,
        int lowTide) {
    super(serverCallStreamObserver, new SimpleQueueAdapter<T>(new SpscArrayQueue<T>(prefetch)), onSubscribe, prefetch, lowTide);
}
 

public void subscribe(final CallStreamObserver<T> downstream) {
    checkNotNull(downstream);

    if (this.downstream == null && DOWNSTREAM.compareAndSet(this, null, downstream)) {
        downstream.setOnReadyHandler(this);
        return;
    }

    throw new IllegalStateException(getClass().getSimpleName() + " does not support multiple subscribers");
}
 

AbstractStreamObserverAndPublisher(
        Queue<T> queue,
        int prefetch,
        int lowTide,
        Consumer<CallStreamObserver<?>> onSubscribe,
        Runnable onTerminate) {
    this.prefetch = prefetch;
    this.limit = lowTide;
    this.queue = queue;
    this.onSubscribe = onSubscribe;
    this.onTerminate = onTerminate;
}
 

AbstractStreamObserverAndPublisher(
        Queue<T> queue,
        int prefetch,
        int lowTide,
        Consumer<CallStreamObserver<?>> onSubscribe) {
    this(queue, prefetch, lowTide, onSubscribe, null);
}
 

protected void onSubscribe(final CallStreamObserver<?> upstream) {
    if (subscription == null && SUBSCRIPTION.compareAndSet(this, null, upstream)) {
        upstream.disableAutoInboundFlowControl();
        if (onSubscribe != null) {
            onSubscribe.accept(upstream);
        }
        return;
    }

    throw new IllegalStateException(getClass().getSimpleName() + " supports only a single subscription");
}
 

public AbstractClientStreamObserverAndPublisher(
        Queue<T> queue,
        Consumer<CallStreamObserver<?>> onSubscribe) {
    super(queue, onSubscribe);
}
 

private void drainRegular(final Subscriber<? super T> subscriber) {
    int missed = 1;

    final CallStreamObserver<?> s = subscription;
    final Queue<T> q = queue;
    int sent = produced;
    long r;

    for (;;) {
        r = requested;
        while (r != sent) {
            boolean d = done;

            T t = q.poll();
            boolean empty = t == null;

            if (checkTerminated(d, empty, subscriber, q)) {
                return;
            }

            if (empty) {
                break;
            }

            subscriber.onNext(t);

            sent++;

            if (sent == limit) {
                if (r != Long.MAX_VALUE) {
                    r = REQUESTED.addAndGet(this, -sent);
                }

                s.request(sent);
                sent = 0;
            }
        }

        if (r == sent) {
            if (checkTerminated(done, q.isEmpty(), subscriber, q)) {
                return;
            }
        }

        int w = wip;
        if (missed == w) {
            produced = sent;
            missed = WIP.addAndGet(this, -missed);
            if (missed == 0) {
                break;
            }
        } else {
            missed = w;
        }
    }
}
 

@Override
public synchronized StreamObserver<Update> streamUpdates(StreamObserver<Update> _outgoingUpdates) {
  // this is kind of odd, but we don't know the right session for this
  // call until we get the first streaming update (grpc doesn't allow
  // a method call with both unary+streaming arguments).
  final AtomicReference<MirrorSession> session = new AtomicReference<>();
  final StreamObserver<Update> outgoingChanges = new BlockingStreamObserver<Update>((CallStreamObserver<Update>) _outgoingUpdates);
  // make an observable for when the client sends in new updates
  return new StreamObserver<Update>() {
    @Override
    public void onNext(Update value) {
      if (session.get() == null) {
        // this is the first update, which is a dummy value with our session id
        MirrorSession ms = sessions.get(value.getPath());
        session.set(ms);
        // look for file system updates to send back to the client
        ms.diffAndStartPolling(new OutgoingConnectionImpl(outgoingChanges));
      } else {
        session.get().addRemoteUpdate(value);
      }
    }

    @Override
    public void onError(Throwable t) {
      Utils.logConnectionError(log, t);
      stopSession();
    }

    @Override
    public void onCompleted() {
      log.info("Connection completed");
      stopSession();
    }

    private void stopSession() {
      if (session.get() != null) {
        session.get().stop();
      }
    }
  };
}
 

void readFrom(InputStream in, long limit, CallStreamObserver<ReadResponse> target) {
  final class ReadFromOnReadyHandler implements Runnable {
    private final byte[] buf = new byte[CHUNK_SIZE];
    private final boolean unlimited = limit == 0;
    private long remaining = limit;
    private boolean complete = false;

    ReadResponse next() throws IOException {
      int readBytes = in.read(buf, 0, (int) Math.min(remaining, buf.length));
      if (readBytes <= 0) {
        if (readBytes == -1) {
          if (!unlimited) {
            throw new UnexpectedEndOfStreamException(remaining, limit);
          }
          complete = true;
        }
        return ReadResponse.getDefaultInstance();
      }

      if (readBytes > remaining) {
        logger.log(Level.WARNING, format("read %d bytes, expected %d", readBytes, remaining));
        readBytes = (int) remaining;
      }
      remaining -= readBytes;
      complete = remaining == 0;
      return ReadResponse.newBuilder().setData(ByteString.copyFrom(buf, 0, readBytes)).build();
    }

    @Override
    public void run() {
      if (!complete) {
        copy();
      }
    }

    void copy() {
      try {
        while (target.isReady() && !complete) {
          ReadResponse response = next();
          if (response.getData().size() != 0) {
            target.onNext(response);
          }
        }

        if (complete) {
          in.close();
          target.onCompleted();
        }
      } catch (Exception e) {
        complete = true;
        try {
          in.close();
        } catch (IOException closeEx) {
          e.addSuppressed(closeEx);
        }
        if (e instanceof UnexpectedEndOfStreamException) {
          e = Status.UNAVAILABLE.withCause(e).asException();
        }
        target.onError(e);
      }
    }
  }
  target.setOnReadyHandler(new ReadFromOnReadyHandler());
}
 

RxClientStreamObserverAndPublisher(
        Consumer<CallStreamObserver<?>> onSubscribe,
        Runnable onTerminate) {
    super(new SimpleQueueAdapter<T>(new SpscArrayQueue<T>(DEFAULT_CHUNK_SIZE)), onSubscribe, onTerminate);
}
 

ReactorClientStreamObserverAndPublisher(
        Consumer<CallStreamObserver<?>> onSubscribe,
        Runnable onTerminate) {
    super(Queues.<T>get(DEFAULT_CHUNK_SIZE).get(), onSubscribe, onTerminate);
}
 

ReactorClientStreamObserverAndPublisher(Consumer<CallStreamObserver<?>> onSubscribe) {
    super(Queues.<T>get(DEFAULT_CHUNK_SIZE).get(), onSubscribe);
}
 

private void drain() {
    if (WIP.getAndIncrement(this) != 0) {
        return;
    }

    int mode = sourceMode;

    int missed = 1;
    final CallStreamObserver<? super T> subscriber = downstream;


    if (mode == NOT_FUSED) {
        final Subscription s = subscription;

        if (s instanceof FusionModeAwareSubscription) {
            mode = ((FusionModeAwareSubscription) s).mode();
            sourceMode = mode;

            if (mode == SYNC) {
                done = true;
            } else {
                s.request(1);
            }
        } else {
            mode = NONE;
            sourceMode = mode;
        }
    }


    for (;;) {
        if (subscriber != null) {
            if (mode == SYNC) {
                drainSync();
            } else if (mode == ASYNC) {
                drainAsync();
            } else {
                drainRegular();
            }

            return;
        }

        missed = WIP.addAndGet(this, -missed);

        if (missed == 0) {
            break;
        }
    }
}