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

package com.ankurm.protocols;

import java.util.Iterator;

import com.ankurm.protocols.grpc.StreamRequest;
import com.ankurm.protocols.grpc.QuoteServiceGrpc;
import io.grpc.ManagedChannel;
import io.grpc.ManagedChannelBuilder;
import org.junit.jupiter.api.Test;
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 static org.assertj.core.api.Assertions.assertThat;

/**
 * Server streaming: how fast can one connection carry N messages one way?
 *
 * <p>This is the shape most of these protocols are actually chosen for &mdash; a price feed, a
 * log tail, a progress stream &mdash; and it separates them differently from request/response,
 * because the per-message cost stops being dominated by a round trip.
 */
@SpringBootTest(webEnvironment = SpringBootTest.WebEnvironment.RANDOM_PORT,
        properties = {"spring.rsocket.server.port=7002", "spring.grpc.server.port=9092"})
class StreamThroughputBenchmarkTest {

    private static final int WARMUP = 5_000;
    private static final int N = 50_000;

    @LocalServerPort
    int httpPort;

    @Autowired
    RSocketStrategies strategies;

    @Test
    void fiftyThousandMessagesEachWay() throws Exception {
        System.out.println("=== server streaming, " + N + " messages on one connection, loopback ===");
        System.out.println(grpc());
        System.out.println(rsocket());
        System.out.println(websocket());
    }

    private String grpc() {
        ManagedChannel channel = ManagedChannelBuilder.forAddress("localhost", 9092)
                .usePlaintext().build();
        try {
            var stub = QuoteServiceGrpc.newBlockingStub(channel);
            drain(stub, WARMUP);
            long t0 = System.nanoTime();
            int seen = drain(stub, N);
            return report("gRPC", seen, System.nanoTime() - t0);
        }
        finally {
            channel.shutdownNow();
        }
    }

    private int drain(QuoteServiceGrpc.QuoteServiceBlockingStub stub, int count) {
        // The blocking stub returns an Iterator. Errors surface mid-iteration rather than at the
        // call, which is the first thing that surprises people about gRPC server streaming.
        Iterator<com.ankurm.protocols.grpc.Quote> it = stub.streamQuotes(
                StreamRequest.newBuilder().setSymbol("AAPL").setCount(count).build());
        int seen = 0;
        while (it.hasNext()) {
            it.next();
            seen++;
        }
        return seen;
    }

    private String rsocket() {
        RSocketRequester requester = RSocketRequester.builder()
                .rsocketStrategies(strategies).tcp("localhost", 7002);
        try {
            var warm = new RSocketQuoteController.StreamSpec("AAPL", WARMUP, 0);
            requester.route("quotes").data(warm).retrieveFlux(Quote.class).blockLast();

            var spec = new RSocketQuoteController.StreamSpec("AAPL", N, 0);
            long t0 = System.nanoTime();
            long seen = requester.route("quotes").data(spec)
                    .retrieveFlux(Quote.class).count().block();
            return report("RSocket", (int) seen, System.nanoTime() - t0);
        }
        finally {
            requester.dispose();
        }
    }

    private String websocket() throws Exception {
        // A generous inbox, because there is nothing else to do with messages that arrive faster
        // than they are read. This IS the WebSocket back-pressure story, in one constant.
        try (WsClient client = new WsClient(httpPort, N + WARMUP + 16)) {
            client.send("STREAM AAPL " + WARMUP + " 0");
            int warm = 0;
            String m;
            while ((m = client.take(20_000)) != null && !"END".equals(m)) {
                warm++;
            }
            assertThat(warm).isEqualTo(WARMUP);

            long t0 = System.nanoTime();
            client.send("STREAM AAPL " + N + " 0");
            int seen = 0;
            while ((m = client.take(20_000)) != null && !"END".equals(m)) {
                seen++;
            }
            return report("WebSocket", seen, System.nanoTime() - t0);
        }
    }

    private String report(String label, int seen, long elapsedNanos) {
        assertThat(seen).isEqualTo(N);
        double seconds = elapsedNanos / 1e9;
        return String.format("%-12s %,7d msgs in %6.3f s  =  %,10.0f msgs/s  (%6.2f us each)",
                label, seen, seconds, seen / seconds, elapsedNanos / 1000.0 / seen);
    }
}