[jdbc-v2] setNetworkTimeout is not applied to statements created before the call, RowBinary writer inserts, or the writer's schema lookup #3268

import java.io.InputStream;
import java.io.OutputStream;
import java.net.ServerSocket;
import java.net.Socket;
import java.sql.Connection;
import java.sql.DriverManager;
import java.sql.PreparedStatement;
import java.sql.Statement;
import java.util.Properties;
import java.util.concurrent.Executor;
import java.util.concurrent.Executors;
import java.util.concurrent.atomic.AtomicBoolean;

/**
 * Throwaway repro (not committed). A TCP proxy in front of the real server holds the first
 * response chunk of every request for DELAY_MS while delay is enabled. Connection#setNetworkTimeout
 * is set to TIMEOUT_MS < DELAY_MS, so every request that honors it must fail with a socket timeout.
 */
public class NetTimeoutRepro {
    static final int DELAY_MS = 1500;
    static final int TIMEOUT_MS = 500;
    static final AtomicBoolean delay = new AtomicBoolean(false);
    static final Executor EXEC = Executors.newCachedThreadPool(r -> { Thread t = new Thread(r); t.setDaemon(true); return t; });

    public static void main(String[] args) throws Exception {
        Class.forName("com.clickhouse.jdbc.Driver");
        ServerSocket ss = new ServerSocket(18124);
        Thread acceptor = new Thread(() -> {
            while (true) {
                try {
                    Socket c = ss.accept();
                    Socket s = new Socket("clickhouse", 18123);
                    AtomicBoolean requestSent = new AtomicBoolean(false);
                    pump(c.getInputStream(), s.getOutputStream(), () -> requestSent.set(true));
                    pump(s.getInputStream(), c.getOutputStream(), () -> {
                        if (delay.get() && requestSent.getAndSet(false)) sleep(DELAY_MS);
                    });
                } catch (Exception e) { return; }
            }
        });
        acceptor.setDaemon(true);
        acceptor.start();

        String proxyUrl = "jdbc:clickhouse://localhost:18124/default";
        String directUrl = "jdbc:clickhouse://clickhouse:18123/default";

        try (Connection c = conn(directUrl, false); Statement s = c.createStatement()) {
            s.execute("DROP TABLE IF EXISTS nt_repro");
            s.execute("CREATE TABLE nt_repro (a Int32) ENGINE = MergeTree ORDER BY a");
        }

        // A: contrast. Statement created AFTER setNetworkTimeout -> must time out.
        run("A  stmt created AFTER setNetworkTimeout, executeQuery", proxyUrl, false, c -> {
            c.setNetworkTimeout(EXEC, TIMEOUT_MS);
            Statement s = c.createStatement();
            delay.set(true);
            s.executeQuery("SELECT 1").close();
        });
        // B: statement created BEFORE setNetworkTimeout, executeQuery.
        run("B  stmt created BEFORE setNetworkTimeout, executeQuery", proxyUrl, false, c -> {
            Statement s = c.createStatement();
            c.setNetworkTimeout(EXEC, TIMEOUT_MS);
            delay.set(true);
            s.executeQuery("SELECT 1").close();
        });
        // B2: same as B, no proxy: the delay is a real server-side sleep.
        run("B2 stmt created BEFORE setNetworkTimeout, executeQuery SELECT sleep(1.5) (no proxy)", directUrl, false, c -> {
            Statement s = c.createStatement();
            c.setNetworkTimeout(EXEC, TIMEOUT_MS);
            s.executeQuery("SELECT sleep(1.5)").close();
        });
        // A2: contrast for B2.
        run("A2 stmt created AFTER setNetworkTimeout, executeQuery SELECT sleep(1.5) (no proxy)", directUrl, false, c -> {
            c.setNetworkTimeout(EXEC, TIMEOUT_MS);
            Statement s = c.createStatement();
            s.executeQuery("SELECT sleep(1.5)").close();
        });
        // C: statement created BEFORE setNetworkTimeout (never set before), executeUpdate.
        run("C  stmt created BEFORE setNetworkTimeout (none set before), executeUpdate", proxyUrl, false, c -> {
            Statement s = c.createStatement();
            c.setNetworkTimeout(EXEC, TIMEOUT_MS);
            delay.set(true);
            s.executeUpdate("INSERT INTO nt_repro VALUES (1)");
        });
        // D: setNetworkTimeout(60s) -> createStatement -> setNetworkTimeout(500ms) -> executeUpdate.
        run("D  setNetworkTimeout(60000) -> stmt -> setNetworkTimeout(500), executeUpdate", proxyUrl, false, c -> {
            c.setNetworkTimeout(EXEC, 60_000);
            Statement s = c.createStatement();
            c.setNetworkTimeout(EXEC, TIMEOUT_MS);
            delay.set(true);
            s.executeUpdate("INSERT INTO nt_repro VALUES (2)");
        });
        // D2: same as D, PreparedStatement.
        run("D2 setNetworkTimeout(60000) -> prepareStatement -> setNetworkTimeout(500), executeQuery", proxyUrl, false, c -> {
            c.setNetworkTimeout(EXEC, 60_000);
            PreparedStatement ps = c.prepareStatement("SELECT ?");
            c.setNetworkTimeout(EXEC, TIMEOUT_MS);
            ps.setInt(1, 1);
            delay.set(true);
            ps.executeQuery().close();
        });
        // H: reverse direction. setNetworkTimeout(500) -> stmt -> setNetworkTimeout(0) (no limit), executeQuery.
        run("H  setNetworkTimeout(500) -> stmt -> setNetworkTimeout(0), executeQuery", proxyUrl, false, c -> {
            c.setNetworkTimeout(EXEC, TIMEOUT_MS);
            Statement s = c.createStatement();
            c.setNetworkTimeout(EXEC, 0);
            delay.set(true);
            s.executeQuery("SELECT 1").close();
        });
        // H2: contrast for H. Fresh statement after setNetworkTimeout(0) -> must succeed.
        run("H2 setNetworkTimeout(500) -> setNetworkTimeout(0) -> stmt, executeQuery", proxyUrl, false, c -> {
            c.setNetworkTimeout(EXEC, TIMEOUT_MS);
            c.setNetworkTimeout(EXEC, 0);
            Statement s = c.createStatement();
            delay.set(true);
            s.executeQuery("SELECT 1").close();
        });
        // E: contrast. Plain PreparedStatementImpl insert (writer off), stmt created after -> must time out.
        run("E  writer OFF: setNetworkTimeout -> prepareStatement INSERT, executeUpdate", proxyUrl, false, c -> {
            c.setNetworkTimeout(EXEC, TIMEOUT_MS);
            PreparedStatement ps = c.prepareStatement("INSERT INTO nt_repro VALUES (?)");
            ps.setInt(1, 3);
            delay.set(true);
            ps.executeUpdate();
        });
        // F: writer ON (beta.row_binary_for_simple_insert=true), stmt created after setNetworkTimeout.
        run("F  writer ON:  setNetworkTimeout -> prepareStatement INSERT, executeUpdate", proxyUrl, true, c -> {
            c.setNetworkTimeout(EXEC, TIMEOUT_MS);
            PreparedStatement ps = c.prepareStatement("INSERT INTO nt_repro VALUES (?)");
            System.out.println("     statement class: " + ps.getClass().getSimpleName());
            ps.setInt(1, 4);
            delay.set(true);
            ps.executeUpdate();
        });
        // F2: writer ON, executeBatch.
        run("F2 writer ON:  setNetworkTimeout -> prepareStatement INSERT, executeBatch", proxyUrl, true, c -> {
            c.setNetworkTimeout(EXEC, TIMEOUT_MS);
            PreparedStatement ps = c.prepareStatement("INSERT INTO nt_repro VALUES (?)");
            ps.setInt(1, 5); ps.addBatch();
            ps.setInt(1, 6); ps.addBatch();
            delay.set(true);
            ps.executeBatch();
        });
        // F3/E3: no proxy. Insert into a Null table whose materialized view sleeps 1.5 s.
        try (Connection c = conn(directUrl, false); Statement s = c.createStatement()) {
            s.execute("DROP TABLE IF EXISTS nt_mv");
            s.execute("DROP TABLE IF EXISTS nt_src");
            s.execute("DROP TABLE IF EXISTS nt_dst");
            s.execute("CREATE TABLE nt_src (a Int32) ENGINE = Null");
            s.execute("CREATE TABLE nt_dst (a Int32, s UInt8) ENGINE = MergeTree ORDER BY a");
            s.execute("CREATE MATERIALIZED VIEW nt_mv TO nt_dst AS SELECT a, sleep(1.5) AS s FROM nt_src");
        }
        run("E3 writer OFF: setNetworkTimeout -> prepareStatement INSERT, MV sleep(1.5) (no proxy)", directUrl, false, c -> {
            c.setNetworkTimeout(EXEC, TIMEOUT_MS);
            PreparedStatement ps = c.prepareStatement("INSERT INTO nt_src VALUES (?)");
            ps.setInt(1, 7);
            ps.executeUpdate();
        });
        run("F3 writer ON:  setNetworkTimeout -> prepareStatement INSERT, MV sleep(1.5) (no proxy)", directUrl, true, c -> {
            c.setNetworkTimeout(EXEC, TIMEOUT_MS);
            PreparedStatement ps = c.prepareStatement("INSERT INTO nt_src VALUES (?)");
            System.out.println("     statement class: " + ps.getClass().getSimpleName());
            ps.setInt(1, 8);
            ps.executeUpdate();
        });
        // G: writer ON, the DESCRIBE TABLE schema lookup inside prepareStatement.
        run("G  writer ON:  setNetworkTimeout -> prepareStatement INSERT (DESCRIBE lookup)", proxyUrl, true, c -> {
            c.setNetworkTimeout(EXEC, TIMEOUT_MS);
            delay.set(true);
            c.prepareStatement("INSERT INTO nt_repro VALUES (?)");
        });

        try (Connection c = conn(directUrl, false); Statement s = c.createStatement()) {
            java.sql.ResultSet rs = s.executeQuery("SELECT groupArray(a) FROM (SELECT a FROM nt_repro ORDER BY a)");
            rs.next();
            System.out.println("rows in nt_repro: " + rs.getString(1));
            s.execute("DROP TABLE IF EXISTS nt_repro");
        }
    }

    interface Body { void run(Connection c) throws Exception; }

    static Connection conn(String url, boolean writer) throws Exception {
        Properties p = new Properties();
        p.setProperty("user", "default");
        p.setProperty("password", "test_default_password");
        if (writer) p.setProperty("beta.row_binary_for_simple_insert", "true");
        return new com.clickhouse.jdbc.Driver().connect(url, p);
    }

    static void run(String name, String url, boolean writer, Body body) {
        delay.set(false);
        long t0 = System.nanoTime();
        String outcome;
        try (Connection c = conn(url, writer)) {
            t0 = System.nanoTime();
            body.run(c);
            outcome = "SUCCEEDED (no timeout)";
        } catch (Throwable e) {
            Throwable r = e;
            while (r.getCause() != null && r.getCause() != r) r = r.getCause();
            outcome = "FAILED: " + e.getClass().getSimpleName() + " / root " + r.getClass().getName() + ": " + r.getMessage();
        }
        long ms = (System.nanoTime() - t0) / 1_000_000;
        delay.set(false);
        System.out.printf("%-92s %5d ms  %s%n", name, ms, outcome);
    }

    static void pump(InputStream in, OutputStream out, Runnable beforeWrite) {
        Thread t = new Thread(() -> {
            byte[] b = new byte[8192];
            try {
                int n;
                while ((n = in.read(b)) > 0) {
                    beforeWrite.run();
                    out.write(b, 0, n);
                    out.flush();
                }
            } catch (Exception ignored) {
            } finally {
                try { in.close(); } catch (Exception ignored) { }
                try { out.close(); } catch (Exception ignored) { }
            }
        });
        t.setDaemon(true);
        t.start();
    }

    static void sleep(long ms) {
        try { Thread.sleep(ms); } catch (InterruptedException ignored) { }
    }
}