java-topology/defects/memcached/unit/MemcachedTest.java

package unit;
import java.util.*;

/**
 * MemcachedTest — CWE-407 benchmark for memcached-0001
 *
 * memcached-0001: slabs_clsid() O(n) linear scan over sorted slabclass[] array
 *   SLOW: O(n) — while loop walking up to 63 slab classes
 *   FAST: O(log n) — binary search, ⌈log₂63⌉ = 6 comparisons max
 *
 * slabs_clsid is called on every item allocation (do_item_alloc) making it
 * a hot path under write load.
 */
public class MemcachedTest {

    // -------------------------------------------------------------------------
    // Benchmark harness
    // -------------------------------------------------------------------------

    static void bench(String label, Runnable slow, Runnable fast, long sOps, long fOps) {
        slow.run(); fast.run();
        long t0 = System.nanoTime(); slow.run(); long sMs = (System.nanoTime() - t0) / 1_000_000;
        long t1 = System.nanoTime(); fast.run(); long fMs = (System.nanoTime() - t1) / 1_000_000;
        System.out.printf("  %-56s  slow:%4dms (%,d ops)  fast:%4dms (%,d ops)  speedup:%.0fx%n",
                label, sMs, sOps, fMs, fOps, fOps > 0 ? (double) sOps / fOps : 0);
    }

    // =========================================================================
    // Model: slab class array (mirrors slabclass[].size in memcached)
    //
    // Sizes are generated with factor=1.25 starting at 96 bytes (memcached
    // default: chunk_size=48, align to 8 → starts at 96 with use_cas).
    // MAX_NUMBER_OF_SLAB_CLASSES = 64 (1..63 + power_largest).
    // =========================================================================

    static final int POWER_SMALLEST = 1;
    static final int MAX_CLASSES = 64;

    static int[] buildSlabSizes() {
        int[] sizes = new int[MAX_CLASSES];
        double size = 96.0;
        double factor = 1.25;
        int chunkMax = 1024 * 1024;  // 1 MB item_size_max
        for (int i = POWER_SMALLEST; i < MAX_CLASSES - 1; i++) {
            // align to 8 bytes
            int aligned = ((int) size + 7) & ~7;
            sizes[i] = aligned;
            if (aligned >= chunkMax / factor) {
                // fill remaining classes with chunkMax
                for (int j = i + 1; j < MAX_CLASSES; j++) sizes[j] = chunkMax;
                break;
            }
            size *= factor;
        }
        sizes[MAX_CLASSES - 1] = chunkMax;
        return sizes;
    }

    // -------------------------------------------------------------------------
    // SLOW: linear scan (current memcached code)
    // -------------------------------------------------------------------------

    /** Returns number of comparisons performed. */
    static long slabs_clsid_slow(int[] sizes, int powerLargest, int querySize) {
        int res = POWER_SMALLEST;
        long ops = 0;
        while (querySize > sizes[res]) {
            ops++;
            if (res++ == powerLargest) return ops;
        }
        ops++;  // final comparison that passes
        return ops;
    }

    // -------------------------------------------------------------------------
    // FAST: binary search (the fix)
    // -------------------------------------------------------------------------

    /** Returns number of comparisons performed. */
    static long slabs_clsid_fast(int[] sizes, int powerLargest, int querySize) {
        int lo = POWER_SMALLEST, hi = powerLargest;
        long ops = 0;
        while (lo < hi) {
            int mid = lo + (hi - lo) / 2;
            ops++;
            if (sizes[mid] < querySize)
                lo = mid + 1;
            else
                hi = mid;
        }
        ops++;  // final check lo==hi
        return ops;
    }

    // =========================================================================
    // Main
    // =========================================================================

    public static void main(String[] args) {
        System.out.println("MemcachedTest — CWE-407");
        System.out.println();

        int[] sizes = buildSlabSizes();
        // Find actual power_largest
        int powerLargest = MAX_CLASSES - 1;
        for (int i = POWER_SMALLEST; i < MAX_CLASSES; i++) {
            if (sizes[i] == 0) { powerLargest = i - 1; break; }
        }
        System.out.printf("  Slab classes: POWER_SMALLEST=%d power_largest=%d%n",
                POWER_SMALLEST, powerLargest);
        System.out.printf("  size[1]=%d  size[%d]=%d%n",
                sizes[1], powerLargest, sizes[powerLargest]);
        System.out.println();

        int passed = 0, total = 0;

        // --- Scenario 1: worst case — query fits only in the largest class ---
        {
            int querySize = sizes[powerLargest];  // must scan all the way
            final long[] sOps = {0}, fOps = {0};
            final int[] ps = {powerLargest};
            final int[] qs = {querySize};
            Runnable slow = () -> {
                long ops = 0;
                for (int r = 0; r < 5_000_000; r++)
                    ops += slabs_clsid_slow(sizes, ps[0], qs[0]);
                sOps[0] = ops;
            };
            Runnable fast = () -> {
                long ops = 0;
                for (int r = 0; r < 5_000_000; r++)
                    ops += slabs_clsid_fast(sizes, ps[0], qs[0]);
                fOps[0] = ops;
            };
            slow.run(); fast.run();
            bench("memcached-0001 slabs_clsid worst-case size=max (5M allocs)", slow, fast, sOps[0], fOps[0]);
            total++;
            // Expect linear >> log; linear = ~63 ops, log = ~6 ops => ~10x
            boolean ok = sOps[0] >= fOps[0] * 5;
            System.out.printf("    [%s] slow=%,d fast=%,d ratio=%.1fx (need >=5x)%n",
                    ok ? "PASS" : "FAIL", sOps[0], fOps[0], (double) sOps[0] / fOps[0]);
            if (ok) passed++;
        }

        // --- Scenario 2: mixed query sizes (realistic workload) ---
        {
            // Queries uniformly distributed over all slab sizes
            int numSizes = powerLargest - POWER_SMALLEST + 1;
            int[] queries = new int[numSizes];
            for (int i = 0; i < numSizes; i++) queries[i] = sizes[POWER_SMALLEST + i];

            final long[] sOps = {0}, fOps = {0};
            final int[] ps = {powerLargest};
            Runnable slow = () -> {
                long ops = 0;
                for (int r = 0; r < 200_000; r++)
                    for (int q : queries)
                        ops += slabs_clsid_slow(sizes, ps[0], q);
                sOps[0] = ops;
            };
            Runnable fast = () -> {
                long ops = 0;
                for (int r = 0; r < 200_000; r++)
                    for (int q : queries)
                        ops += slabs_clsid_fast(sizes, ps[0], q);
                fOps[0] = ops;
            };
            slow.run(); fast.run();
            bench("memcached-0001 slabs_clsid mixed sizes (200k×all)", slow, fast, sOps[0], fOps[0]);
            total++;
            // Average linear: ~(63+1)/2 ≈ 32 ops; average binary: ~5 ops
            boolean ok = sOps[0] >= fOps[0] * 4;
            System.out.printf("    [%s] slow=%,d fast=%,d ratio=%.1fx (need >=4x)%n",
                    ok ? "PASS" : "FAIL", sOps[0], fOps[0], (double) sOps[0] / fOps[0]);
            if (ok) passed++;
        }

        // --- Scenario 3: correctness check — both find the same class ---
        {
            Random rng = new Random(0);
            int maxSize = sizes[powerLargest];
            int mismatches = 0;
            for (int i = 0; i < 100_000; i++) {
                int q = 1 + rng.nextInt(maxSize);
                // find expected class with slow (reference)
                int res_s = POWER_SMALLEST;
                while (q > sizes[res_s] && res_s < powerLargest) res_s++;
                // find with fast
                int lo = POWER_SMALLEST, hi = powerLargest;
                while (lo < hi) {
                    int mid = lo + (hi - lo) / 2;
                    if (sizes[mid] < q) lo = mid + 1;
                    else hi = mid;
                }
                int res_f = lo;
                if (res_s != res_f) mismatches++;
            }
            total++;
            boolean ok = mismatches == 0;
            System.out.printf("  %-56s  [%s] mismatches=%d%n",
                    "memcached-0001 correctness (100k random queries)",
                    ok ? "PASS" : "FAIL", mismatches);
            if (ok) passed++;
        }

        System.out.println();
        System.out.printf("%d/%d PASS%n", passed, total);
        if (passed < total) System.exit(1);
    }
}