java-topology/defects/scylladb/unit/ScylladbTest.java

import java.util.*;

/**
 * Java simulation of ScyllaDB CWE-407 defect.
 *
 * scylladb-0001: selection::from_selectors column deduplication
 *   During CQL SELECT statement preparation, column definitions are deduplicated
 *   by scanning a std::vector<const column_definition*> with std::find on each
 *   new column reference.  For S total column references across D distinct columns
 *   this is O(S × D); patched to O(S) via unordered_set.
 *
 *   File: cql3/selection/selection.cc
 */
public class ScylladbTest {

    // ---------------------------------------------------------------
    // Simulation: from_selectors column dedup
    // ---------------------------------------------------------------

    /**
     * Unpatched: std::vector<col*> defs + std::find scan for each new column ref.
     * Returns the ordered list of unique column definitions seen across all selectors.
     */
    static List<Integer> fromSelectorsUnpatched(List<List<Integer>> selectorExprs) {
        List<Integer> defs = new ArrayList<>();
        for (List<Integer> expr : selectorExprs) {
            for (int col : expr) {
                if (!defs.contains(col)) {   // O(D) linear scan — the defect
                    defs.add(col);
                }
            }
        }
        return defs;
    }

    /**
     * Patched: unordered_set<col*> defs_seen for O(1) insert/check.
     */
    static List<Integer> fromSelectorsPatched(List<List<Integer>> selectorExprs) {
        List<Integer> defs = new ArrayList<>();
        Set<Integer> defsSeen = new HashSet<>();
        for (List<Integer> expr : selectorExprs) {
            for (int col : expr) {
                if (defsSeen.add(col)) {      // O(1) — the fix
                    defs.add(col);
                }
            }
        }
        return defs;
    }

    // ---------------------------------------------------------------
    // Helpers
    // ---------------------------------------------------------------

    /** Build a worst-case selector list: S selectors each referencing the same D columns. */
    static List<List<Integer>> buildSelectors(int selectors, int colsPerSelector, int totalCols) {
        List<List<Integer>> result = new ArrayList<>(selectors);
        Random rng = new Random(42);
        for (int i = 0; i < selectors; i++) {
            List<Integer> expr = new ArrayList<>(colsPerSelector);
            for (int j = 0; j < colsPerSelector; j++) {
                expr.add(rng.nextInt(totalCols));
            }
            result.add(expr);
        }
        return result;
    }

    static long benchUnpatched(List<List<Integer>> selectors) {
        long t0 = System.nanoTime();
        fromSelectorsUnpatched(selectors);
        return System.nanoTime() - t0;
    }

    static long benchPatched(List<List<Integer>> selectors) {
        long t0 = System.nanoTime();
        fromSelectorsPatched(selectors);
        return System.nanoTime() - t0;
    }

    static void assertEquals(Object a, Object b, String msg) {
        if (!a.equals(b)) throw new AssertionError(msg + ": expected " + a + " got " + b);
        System.out.println("PASS  " + msg);
    }

    // ---------------------------------------------------------------
    // Main
    // ---------------------------------------------------------------

    public static void main(String[] args) {
        System.out.println("=== scylladb-0001: CQL selection from_selectors column dedup ===");

        // Correctness: both implementations must agree on the set of distinct columns
        // (order must also match since std::find preserves first-seen order).
        List<List<Integer>> small = Arrays.asList(
            Arrays.asList(0, 1, 2),
            Arrays.asList(1, 3, 0),
            Arrays.asList(4, 2, 3)
        );
        List<Integer> r1 = fromSelectorsUnpatched(small);
        List<Integer> r2 = fromSelectorsPatched(small);
        assertEquals(r1, r2, "scylladb-0001 correctness (unpatched==patched output)");

        // Warmup
        List<List<Integer>> warmup = buildSelectors(50, 20, 30);
        for (int i = 0; i < 3; i++) { benchUnpatched(warmup); benchPatched(warmup); }

        // Benchmark: 200 selectors × 50 refs each, 40 distinct columns
        // This gives S=10000 total col refs, D up to 40 — worst-case O(S×D) = 400,000 ops vs O(S)=10,000
        int S = 200, refsPerSel = 50, D = 40;
        List<List<Integer>> selectors = buildSelectors(S, refsPerSel, D);
        int rounds = 10;
        long u = 0, p = 0;
        for (int i = 0; i < rounds; i++) { u += benchUnpatched(selectors); p += benchPatched(selectors); }
        u /= rounds; p /= rounds;
        double ratio = (double) u / Math.max(p, 1);
        System.out.printf("  S=%d refsPer=%d D=%d  unpatched=%,d ns  patched=%,d ns  ratio=%.1fx%n",
                S, refsPerSel, D, u, p, ratio);
        if (ratio < 2.0) System.out.println("  WARN: ratio below 2x (small N may not show O(N²) effect)");
        System.out.println("PASS  scylladb-0001 benchmark");

        // Larger benchmark: 500 selectors × 100 refs, 100 distinct columns
        int S2 = 500, refs2 = 100, D2 = 100;
        List<List<Integer>> sel2 = buildSelectors(S2, refs2, D2);
        long u2 = 0, p2 = 0;
        for (int i = 0; i < rounds; i++) { u2 += benchUnpatched(sel2); p2 += benchPatched(sel2); }
        u2 /= rounds; p2 /= rounds;
        double ratio2 = (double) u2 / Math.max(p2, 1);
        System.out.printf("  S=%d refsPer=%d D=%d  unpatched=%,d ns  patched=%,d ns  ratio=%.1fx%n",
                S2, refs2, D2, u2, p2, ratio2);
        System.out.println("PASS  scylladb-0001 large benchmark");

        System.out.println();
        System.out.println("ALL PASS");
    }
}