spring-messaging-demo/protocol-comparison/src/test/java/com/ankurm/protocols/BackPressureTest.java

package com.ankurm.protocols;

import java.util.Iterator;
import java.util.concurrent.atomic.AtomicLong;

import com.ankurm.protocols.grpc.QuoteRequest;
import com.ankurm.protocols.grpc.QuoteServiceGrpc;
import io.grpc.ManagedChannel;
import io.grpc.ManagedChannelBuilder;
import org.junit.jupiter.api.Test;
import org.reactivestreams.Subscription;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.boot.test.context.SpringBootTest;
import org.springframework.boot.test.web.server.LocalServerPort;
import org.springframework.messaging.rsocket.RSocketRequester;
import org.springframework.messaging.rsocket.RSocketStrategies;
import reactor.core.publisher.BaseSubscriber;

import static org.assertj.core.api.Assertions.assertThat;

/**
 * The measurement that decides the article.
 *
 * <p>Each protocol serves an unbounded stream. Each client takes a hundred messages and then does
 * nothing for a while. The question is: <strong>how many did the server produce?</strong>
 *
 * <p>The sweep over quiet periods is what makes the answer conclusive. A number that stays flat
 * as the quiet period grows means something bounded the producer. A number that grows in
 * proportion to the wait means <em>nothing did</em> &mdash; the producer is running flat out and
 * the only limit is the clock.
 */
@SpringBootTest(webEnvironment = SpringBootTest.WebEnvironment.RANDOM_PORT,
        properties = {"spring.rsocket.server.port=7003", "spring.grpc.server.port=9093"})
class BackPressureTest {

    private static final int WANTED = 100;
    private static final long[] QUIET_MS = {500, 1_000, 2_000};

    @LocalServerPort
    int httpPort;

    @Autowired
    RSocketStrategies strategies;

    @Test
    void howManyDidTheServerProduce() throws Exception {
        System.out.println("=== unbounded stream; client takes " + WANTED + ", then goes quiet ===");
        System.out.printf("%-14s %10s %10s %10s%n", "protocol", "quiet 0.5s", "quiet 1s", "quiet 2s");

        long[] rs = new long[QUIET_MS.length];
        long[] gr = new long[QUIET_MS.length];
        long[] ws = new long[QUIET_MS.length];
        for (int i = 0; i < QUIET_MS.length; i++) {
            rs[i] = rsocket(QUIET_MS[i]);
            gr[i] = grpc(QUIET_MS[i]);
            ws[i] = websocket(QUIET_MS[i]);
        }
        row("RSocket", rs);
        row("gRPC", gr);
        row("WebSocket", ws);

        System.out.println();
        System.out.println("RSocket is flat: the server produced exactly what request(n) asked for.");
        System.out.printf("gRPC grows %.1fx between 0.5 s and 2 s; WebSocket grows %.1fx.%n",
                gr[2] / (double) gr[0], ws[2] / (double) ws[0]);
        System.out.println("Growth that tracks the wait, rather than settling at a buffer size,");
        System.out.println("means nothing bounded the producer. (Growth is faster than linear");
        System.out.println("because the loop is still being JIT-compiled during the first run.)");

        // RSocket delivered exactly the demand, at every quiet period.
        assertThat(rs).containsOnly(WANTED);
        // The other two grew with the wait rather than settling at a buffer size.
        assertThat(gr[2]).isGreaterThan(gr[0] * 2);
        assertThat(ws[2]).isGreaterThan(ws[0] * 2);
    }

    private void row(String label, long[] v) {
        System.out.printf("%-14s %,10d %,10d %,10d%n", label, v[0], v[1], v[2]);
    }

    /**
     * RSocket: {@code request(n)} is a frame on the wire. The server is told the number, produces
     * that many, and stops. Nothing blocks and nothing buffers &mdash; the producer is simply not
     * called again.
     */
    private long rsocket(long quietMs) throws Exception {
        RSocketRequester requester = RSocketRequester.builder()
                .rsocketStrategies(strategies).tcp("localhost", 7003);
        try {
            AtomicLong received = new AtomicLong();
            BaseSubscriber<Quote> subscriber = new BaseSubscriber<>() {
                @Override
                protected void hookOnSubscribe(Subscription subscription) {
                    subscription.request(WANTED);   // and never again
                }

                @Override
                protected void hookOnNext(Quote value) {
                    received.incrementAndGet();
                }
            };
            requester.route("quotes.unbounded").data("AAPL")
                    .retrieveFlux(Quote.class).subscribe(subscriber);
            Thread.sleep(quietMs);
            long produced = requester.route("produced").retrieveMono(Long.class).block();
            subscriber.dispose();
            assertThat(received.get()).isEqualTo(WANTED);
            return produced;
        }
        finally {
            requester.dispose();
        }
    }

    /**
     * gRPC: there is no application-level demand signal. The blocking stub's {@code Iterator}
     * requests one message per {@code next()}, and the handler writes with
     * {@code StreamObserver.onNext}, which never blocks.
     */
    private long grpc(long quietMs) throws Exception {
        ManagedChannel channel = ManagedChannelBuilder.forAddress("localhost", 9093)
                .usePlaintext().build();
        try {
            var stub = QuoteServiceGrpc.newBlockingStub(channel);
            QuoteRequest req = QuoteRequest.newBuilder().setSymbol("AAPL").build();
            Iterator<com.ankurm.protocols.grpc.Quote> it = stub.streamUnbounded(req);
            long received = 0;
            while (received < WANTED && it.hasNext()) {
                it.next();
                received++;
            }
            Thread.sleep(quietMs);
            return stub.produced(req).getCount();
        }
        finally {
            channel.shutdownNow();
        }
    }

    /** Raw WebSocket: the frame protocol has no notion of demand at all. */
    private long websocket(long quietMs) throws Exception {
        try (WsClient client = new WsClient(httpPort, WANTED)) {
            client.send("UNBOUNDED AAPL");
            long received = 0;
            while (received < WANTED && client.take(5_000) != null) {
                received++;
            }
            Thread.sleep(quietMs);
            // Ask on a second connection: the first one's server-side thread is busy writing.
            try (WsClient asker = new WsClient(httpPort, 8)) {
                asker.send("PRODUCED");
                return Long.parseLong(asker.take(5_000));
            }
        }
    }
}