// Repro: grpc-gcp channel affinity collapses onto a few channels under
// concurrent warmup with multiplexed sessions, capping per-client throughput.
//
// For each (grpcGcp, concurrency) pair, builds a FRESH Spanner client,
// fires a high-concurrency warmup burst (this is when the race binds keys),
// then measures qps/p50/p99 of SELECT 1 at that concurrency. Expect parity
// at C<=100 and a large divergence (lower qps, higher p50) for grpc-gcp=ON
// at C>=200.
//
// javac -cp 'lib/*' SpannerAffinityRepro.java
// java  -cp '.:lib/*' SpannerAffinityRepro PROJECT INSTANCE DATABASE
import com.google.cloud.spanner.*;
import java.util.*;
import java.util.concurrent.*;
import java.util.concurrent.atomic.AtomicInteger;

public class SpannerAffinityRepro {
  static final int NUM_CHANNELS = 32;
  static final int[] SWEEP = {50, 100, 200, 400};

  record Row(double qps, double p50, double p99) {}

  static Row one(String proj, String inst, String dbId, boolean gcp, int c) throws Exception {
    SpannerOptions.Builder b = SpannerOptions.newBuilder()
        .setProjectId(proj).setNumChannels(NUM_CHANNELS);
    if (!gcp) b.disableGrpcGcpExtension();
    try (Spanner s = b.build().getService()) {
      DatabaseClient db = s.getDatabaseClient(DatabaseId.of(proj, inst, dbId));
      Statement stmt = Statement.of("SELECT 1");
      Runnable q = () -> { try (ResultSet rs = db.singleUse().executeQuery(stmt)) { while (rs.next()) {} } };
      ExecutorService ex = Executors.newVirtualThreadPerTaskExecutor();
      // Warmup AT the target concurrency: this is the burst during which the
      // ~63 affinity keys race through pickLeastBusyChannel and bind.
      runAt(ex, q, c, Math.max(500, c * 4), null);
      int iters = Math.max(2000, c * 30);
      long[] lats = new long[iters];
      long elapsed = runAt(ex, q, c, iters, lats);
      ex.shutdown();
      Arrays.sort(lats);
      double ms = 1e-6;
      return new Row(iters / (elapsed * 1e-9), lats[iters / 2] * ms, lats[(int) (iters * 0.99)] * ms);
    }
  }

  static long runAt(ExecutorService ex, Runnable q, int c, int iters, long[] lats) throws Exception {
    AtomicInteger left = new AtomicInteger(iters), idx = new AtomicInteger();
    CountDownLatch done = new CountDownLatch(c);
    long t0 = System.nanoTime();
    for (int i = 0; i < c; i++) ex.submit(() -> {
      while (true) {
        int s = left.decrementAndGet(); if (s < 0) break;
        long t = System.nanoTime(); q.run();
        if (lats != null) lats[idx.getAndIncrement()] = System.nanoTime() - t;
      }
      done.countDown();
    });
    done.await();
    return System.nanoTime() - t0;
  }

  public static void main(String[] a) throws Exception {
    if (a.length < 3) { System.err.println("usage: PROJECT INSTANCE DATABASE"); System.exit(1); }
    System.out.printf("numChannels=%d, multiplexed sessions (default), fresh client per cell%n", NUM_CHANNELS);
    System.out.printf("%5s | %22s | %22s%n", "", "grpc-gcp ON (default)", "grpc-gcp OFF");
    System.out.printf("%5s | %6s %6s %6s | %6s %6s %6s%n", "C", "qps", "p50ms", "p99ms", "qps", "p50ms", "p99ms");
    for (int c : SWEEP) {
      Row on = one(a[0], a[1], a[2], true, c);
      Row off = one(a[0], a[1], a[2], false, c);
      System.out.printf("%5d | %6.0f %6.1f %6.1f | %6.0f %6.1f %6.1f%n",
          c, on.qps, on.p50, on.p99, off.qps, off.p50, off.p99);
    }
  }
}