spring-messaging-demo/protocol-comparison/src/main/java/com/ankurm/protocols/GrpcQuoteService.java

package com.ankurm.protocols;

import com.ankurm.protocols.grpc.Quote;
import com.ankurm.protocols.grpc.QuoteRequest;
import com.ankurm.protocols.grpc.QuoteServiceGrpc;
import com.ankurm.protocols.grpc.StreamRequest;
import io.grpc.Context;
import io.grpc.stub.ServerCallStreamObserver;
import io.grpc.stub.StreamObserver;
import java.util.concurrent.atomic.AtomicLong;

import org.springframework.stereotype.Service;

/**
 * The gRPC side. No registration code: the generated {@code ImplBase} is a
 * {@code BindableService} and Boot registers every such bean with the server.
 *
 * <p>The cancellation check in {@link #streamQuotes} is not decoration. gRPC does not interrupt
 * your thread when a client goes away &mdash; it sets a flag on the {@link Context} &mdash; so a
 * loop that never looks keeps producing into a stream nobody is reading. That is covered at
 * length in <a href="https://ankurm.com/spring-grpc-spring-boot-4/">the Spring gRPC article</a>;
 * here it matters because the back-pressure benchmark deliberately abandons a stream.
 *
 * @see <a href="../../../../../docs/04-backpressure.md">docs/04-backpressure.md</a>
 */
@Service
public class GrpcQuoteService extends QuoteServiceGrpc.QuoteServiceImplBase {

    /** Counts what the SERVER produced, comparable with the RSocket and WebSocket counters. */
    static final AtomicLong PRODUCED = new AtomicLong();

    private final QuoteSource source;

    GrpcQuoteService(QuoteSource source) {
        this.source = source;
    }

    @Override
    public void getQuote(QuoteRequest request, StreamObserver<Quote> observer) {
        observer.onNext(toProto(source.at(request.getSymbol(), 0)));
        observer.onCompleted();
    }

    /**
     * An effectively unbounded stream, for the back-pressure comparison. A consumer that stops
     * calling {@code next()} closes the HTTP/2 flow-control window, and this loop blocks inside
     * {@code onNext} &mdash; so the server stops producing, but it stops by blocking a thread
     * rather than by being asked to stop.
     */
    @Override
    public void streamUnbounded(QuoteRequest request, StreamObserver<Quote> observer) {
        PRODUCED.set(0);
        while (!Context.current().isCancelled()) {
            observer.onNext(toProto(source.at(request.getSymbol(), PRODUCED.get())));
            PRODUCED.incrementAndGet();
        }
    }

    /**
     * The fix, and the reason the unfixed version is worth measuring.
     *
     * <p>{@code StreamObserver.onNext} on a gRPC server <strong>never blocks</strong>. If the
     * client is not reading, the message is queued in the server's outbound buffer and the loop
     * keeps going &mdash; HTTP/2 flow control governs the wire, not your code. The only thing
     * that connects the two is {@link ServerCallStreamObserver#isReady()}, and using it means
     * restructuring the handler around {@code setOnReadyHandler} rather than writing a loop.
     *
     * <p>The spin here is deliberately the simplest possible demonstration rather than the
     * shape you would ship; a real implementation registers an on-ready handler and returns.
     */
    @Override
    public void streamUnboundedReady(QuoteRequest request, StreamObserver<Quote> observer) {
        PRODUCED.set(0);
        ServerCallStreamObserver<Quote> ready = (ServerCallStreamObserver<Quote>) observer;
        while (!ready.isCancelled()) {
            if (!ready.isReady()) {
                Thread.onSpinWait();
                continue;
            }
            ready.onNext(toProto(source.at(request.getSymbol(), PRODUCED.get())));
            PRODUCED.incrementAndGet();
        }
    }

    @Override
    public void produced(QuoteRequest request, StreamObserver<com.ankurm.protocols.grpc.ProducedCount> observer) {
        observer.onNext(com.ankurm.protocols.grpc.ProducedCount.newBuilder()
                .setCount(PRODUCED.get()).build());
        observer.onCompleted();
    }

    @Override
    public void streamQuotes(StreamRequest request, StreamObserver<Quote> observer) {
        for (int i = 0; i < request.getCount(); i++) {
            if (Context.current().isCancelled()) {
                // Do NOT call onCompleted/onError here: the stream is already closed and
                // touching it throws IllegalStateException. Just return.
                return;
            }
            observer.onNext(toProto(source.at(request.getSymbol(), i)));
            if (request.getDelayMs() > 0) {
                try {
                    Thread.sleep(request.getDelayMs());
                }
                catch (InterruptedException ex) {
                    Thread.currentThread().interrupt();
                    return;
                }
            }
        }
        observer.onCompleted();
    }

    static Quote toProto(com.ankurm.protocols.Quote q) {
        return Quote.newBuilder()
                .setSymbol(q.symbol()).setSeq(q.seq())
                .setBid(q.bid()).setAsk(q.ask())
                .setEpochMicros(q.epochMicros())
                .build();
    }
}