Benchmarking YSQL Connection Time in YugabyteDB with the JDBC Smart Driver

				
					package com.yb.smart;

import java.sql.Connection;
import java.sql.DriverManager;
import java.sql.ResultSet;
import java.sql.Statement;
import java.text.MessageFormat;
import java.text.SimpleDateFormat;
import java.util.Arrays;
import java.util.Date;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.atomic.AtomicLong;
import java.util.logging.ConsoleHandler;
import java.util.logging.Handler;
import java.util.logging.Level;
import java.util.logging.LogManager;
import java.util.logging.LogRecord;
import java.util.logging.Logger;

public class App {

    public static void main(String[] args) throws Exception {
        if (args.length < 1) {
            System.err.println("Usage: java -jar yb-smart-driver-test.jar <JDBC_URL> [--quiet|--verbose]");
            System.exit(1);
        }

        final String jdbcUrl = args[0];
        final boolean quiet   = Arrays.stream(args).anyMatch("--quiet"::equalsIgnoreCase);
        final boolean verbose = Arrays.stream(args).anyMatch("--verbose"::equalsIgnoreCase);

        // Primary timing markers based on "Created connection to ..."
        final AtomicLong tStart     = new AtomicLong(-1); // first "Trying to establish..."
        final AtomicLong tCreated1  = new AtomicLong(-1); // first "Created connection to ..."
        final AtomicLong tCreated2  = new AtomicLong(-1); // second "Created connection to ..." (load-balanced hop)

        // Fallback timing markers based on "AuthenticationOk"
        final AtomicLong tAuth1     = new AtomicLong(-1);
        final AtomicLong tAuth2     = new AtomicLong(-1);

        // Count "Trying..." to detect load balancing
        final AtomicInteger tryingCount = new AtomicInteger(0);

        setupLogging(quiet, verbose, tStart, tCreated1, tCreated2, tAuth1, tAuth2, tryingCount);

        // Measure wall-clock around the entire getConnection call
        final long t0 = System.nanoTime();
        String host = null, cloud = null, region = null, zone = null;
        try (Connection conn = DriverManager.getConnection(jdbcUrl)) {
            final long t1 = System.nanoTime();
            final double connectionMs = (t1 - t0) / 1_000_000.0;
            System.out.printf("connection_time_ms=%.3f%n", connectionMs);

            // Decide which markers to use for first/final
            final boolean createdAvailable = (tCreated1.get() > 0);
            final long firstEnd  = createdAvailable ? tCreated1.get() : tAuth1.get();
            final long secondEnd = createdAvailable ? tCreated2.get() : tAuth2.get();

            // Names reflect the semantics explained in the blog
            final String firstLabel  = "driver_first_attempt_ms";
            final String finalLabel  = "driver_final_attempt_ms";

            if (tStart.get() > 0 && firstEnd > 0 && firstEnd >= tStart.get()) {
                System.out.printf("%s=%.3f%n", firstLabel, (firstEnd - tStart.get()) * 1.0);
            } else {
                System.out.printf("%s=N/A%n", firstLabel);
            }

            if (tStart.get() > 0 && secondEnd > 0 && secondEnd >= tStart.get()) {
                System.out.printf("%s=%.3f%n", finalLabel, (secondEnd - tStart.get()) * 1.0);
            }

            // Print load-balanced flag (true if we saw 2+ "Trying..." lines)
            System.out.printf("load_balanced=%b%n", tryingCount.get() >= 2);

            // Show where we landed
            try (Statement stmt = conn.createStatement()) {
                final String sql = "SELECT host(inet_server_addr()) AS host, " +
                        "yb_server_cloud() AS cloud, " +
                        "yb_server_region() AS region, " +
                        "yb_server_zone() AS zone";
                try (ResultSet rs = stmt.executeQuery(sql)) {
                    if (rs.next()) {
                        host   = rs.getString("host");
                        cloud  = rs.getString("cloud");
                        region = rs.getString("region");
                        zone   = rs.getString("zone");
                    }
                }
            }

            if (host != null) {
                System.out.printf("host=%s / cloud=%s / region=%s / zone=%s%n",
                        host, cloud, region, zone);
            }
        }
    }

    private static void setupLogging(
            boolean quiet,
            boolean verbose,
            AtomicLong tStart,
            AtomicLong tCreated1,
            AtomicLong tCreated2,
            AtomicLong tAuth1,
            AtomicLong tAuth2,
            AtomicInteger tryingCount
    ) {
        // Reset global logging configuration
        LogManager.getLogManager().reset();

        // Our single handler: captures markers; prints depending on flags
        Handler h = new ConsoleHandler() {
            private final SimpleDateFormat sdf = new SimpleDateFormat("yyyy-MM-dd HH:mm:ss.SSS");
            { setLevel(Level.ALL); }

            @Override
            public synchronized void publish(LogRecord r) {
                String raw = r.getMessage();
                if (raw == null) return;

                // Expand {0} placeholders (driver uses MessageFormat)
                String msg = (r.getParameters() != null && r.getParameters().length > 0)
                        ? MessageFormat.format(raw, r.getParameters())
                        : raw;

                // Capture start and end markers
                if (msg.contains("Trying to establish a protocol version 3 connection to")) {
                    tryingCount.incrementAndGet();
                    tStart.compareAndSet(-1, r.getMillis()); // only first start matters as our zero
                }
                if (msg.contains("Created connection to")) {
                    if (tCreated1.compareAndSet(-1, r.getMillis())) {
                        // first Created
                    } else {
                        tCreated2.compareAndSet(-1, r.getMillis());
                    }
                }
                if (msg.contains("AuthenticationOk")) {
                    if (tAuth1.compareAndSet(-1, r.getMillis())) {
                        // first Auth
                    } else {
                        tAuth2.compareAndSet(-1, r.getMillis());
                    }
                }

                // Printing policy:
                //  - --quiet   : print nothing from the driver
                //  - --verbose : print ALL log entries
                //  - default   : print only key lines (start + Created + AuthenticationOk)
                if (quiet) {
                    return;
                }

                boolean isKey =
                        msg.contains("Trying to establish a protocol version 3 connection to") ||
                        msg.contains("Created connection to") ||
                        msg.contains("AuthenticationOk");

                if (verbose || isKey) {
                    String ts = sdf.format(new Date(r.getMillis()));
                    String loggerName = r.getLoggerName() != null ? r.getLoggerName() : "";
                    System.out.printf("%s %s %s%n", ts, loggerName, msg);
                }
            }
        };

        // Attach handler to Yugabyte driver loggers (and only those), at FINEST
        Level lvl = Level.FINEST;

        Logger core = Logger.getLogger("com.yugabyte.core");
        core.setUseParentHandlers(false);
        core.setLevel(lvl);
        core.addHandler(h);

        Logger drv = Logger.getLogger("com.yugabyte.Driver");
        drv.setUseParentHandlers(false);
        drv.setLevel(lvl);
        drv.addHandler(h);

        Logger jdbc = Logger.getLogger("com.yugabyte.jdbc");
        jdbc.setUseParentHandlers(false);
        jdbc.setLevel(lvl);
        jdbc.addHandler(h);

        // Silence everything else (root)
        Logger root = Logger.getLogger("");
        root.setUseParentHandlers(false);
        root.setLevel(Level.OFF);
    }
}