Add heap allocation audit for relayed connections

alloc_counter is a preloadable malloc counter (glibc, macOS) that writes a snapshot on SIGUSR2. measure_alloc.py warms ss-local and ss-server up, snapshots both, relays more connections and reports allocations per connection. It runs report-only as the alloc_audit_report ctest. interop.peers() now yields the started processes.
2026-10-27 03:03:43 +00:00 · 2026-10-14 12:51:47 +08:00
parent f90cfe9034
commit 641643ec48
3 changed files with 357 additions and 1 deletions
@@ -0,0 +1,217 @@
+/* SPDX-License-Identifier: GPL-3.0-or-later */
+/*
+ * alloc_counter: a preloadable shim that counts heap allocations.
+ *
+ * Load it into a program with LD_PRELOAD (glibc) or DYLD_INSERT_LIBRARIES
+ * (macOS). Every SIGUSR2 appends one line with the running totals to the
+ * file named by SS_ALLOC_AUDIT_OUT:
+ *
+ *     <pid> <allocations> <frees> <bytes>
+ *
+ * A driver takes a snapshot after warming a proxy up and another after more
+ * traffic; the difference is what the steady-state data path allocates.
+ * tests/alloc_audit/measure_alloc.py is that driver.
+ *
+ * The signal handler only uses async-signal-safe calls (open, write, close).
+ */
+#ifndef _GNU_SOURCE
+#define _GNU_SOURCE
+#endif
+#include <errno.h>
+#include <fcntl.h>
+#include <signal.h>
+#include <stddef.h>
+#include <stdint.h>
+#include <stdlib.h>
+#include <string.h>
+#include <unistd.h>
+
+static uint64_t allocations;
+static uint64_t frees;
+static uint64_t bytes;
+static const char *out_path;
+
+static void
+count_alloc(size_t size)
+{
+    __atomic_add_fetch(&allocations, 1, __ATOMIC_RELAXED);
+    __atomic_add_fetch(&bytes, (uint64_t)size, __ATOMIC_RELAXED);
+}
+
+static void
+count_free(void *ptr)
+{
+    if (ptr != NULL) {
+        __atomic_add_fetch(&frees, 1, __ATOMIC_RELAXED);
+    }
+}
+
+static char *
+format_u64(char *end, uint64_t value)
+{
+    do {
+        *--end = (char)('0' + value % 10);
+        value /= 10;
+    } while (value != 0);
+    return end;
+}
+
+static void
+snapshot(int signo)
+{
+    (void)signo;
+    int saved_errno = errno;
+    if (out_path != NULL) {
+        uint64_t fields[4] = {
+            (uint64_t)getpid(),
+            __atomic_load_n(&allocations, __ATOMIC_RELAXED),
+            __atomic_load_n(&frees, __ATOMIC_RELAXED),
+            __atomic_load_n(&bytes, __ATOMIC_RELAXED),
+        };
+        char line[128];
+        char *end = line + sizeof(line);
+        char *p   = end;
+        *--p = '\n';
+        for (int i = 3; i >= 0; i--) {
+            p = format_u64(p, fields[i]);
+            if (i != 0) {
+                *--p = ' ';
+            }
+        }
+        int fd = open(out_path, O_WRONLY | O_CREAT | O_APPEND, 0644);
+        if (fd >= 0) {
+            ssize_t ignored = write(fd, p, (size_t)(end - p));
+            (void)ignored;
+            close(fd);
+        }
+    }
+    errno = saved_errno;
+}
+
+__attribute__((constructor)) static void
+install(void)
+{
+    out_path = getenv("SS_ALLOC_AUDIT_OUT");
+    struct sigaction action;
+    memset(&action, 0, sizeof(action));
+    action.sa_handler = snapshot;
+    action.sa_flags   = SA_RESTART;
+    sigemptyset(&action.sa_mask);
+    sigaction(SIGUSR2, &action, NULL);
+}
+
+#if defined(__APPLE__)
+
+#include <malloc/malloc.h>
+
+static void *
+audit_malloc(size_t size)
+{
+    count_alloc(size);
+    return malloc(size);
+}
+
+static void *
+audit_calloc(size_t count, size_t size)
+{
+    count_alloc(count * size);
+    return calloc(count, size);
+}
+
+static void *
+audit_realloc(void *ptr, size_t size)
+{
+    count_alloc(size);
+    return realloc(ptr, size);
+}
+
+static int
+audit_posix_memalign(void **out, size_t align, size_t size)
+{
+    count_alloc(size);
+    return posix_memalign(out, align, size);
+}
+
+static void
+audit_free(void *ptr)
+{
+    count_free(ptr);
+    free(ptr);
+}
+
+#define SS_INTERPOSE(replacement, original)                                   \
+    __attribute__((used)) static const struct {                               \
+        const void *replacement_fn;                                           \
+        const void *original_fn;                                              \
+    } interpose_##original __attribute__((section("__DATA,__interpose"))) = { \
+        (const void *)(uintptr_t)&replacement, (const void *)(uintptr_t)&original }
+
+SS_INTERPOSE(audit_malloc, malloc);
+SS_INTERPOSE(audit_calloc, calloc);
+SS_INTERPOSE(audit_realloc, realloc);
+SS_INTERPOSE(audit_posix_memalign, posix_memalign);
+SS_INTERPOSE(audit_free, free);
+
+#elif defined(__GLIBC__)
+
+extern void *__libc_malloc(size_t);
+extern void *__libc_calloc(size_t, size_t);
+extern void *__libc_realloc(void *, size_t);
+extern void *__libc_memalign(size_t, size_t);
+extern void __libc_free(void *);
+
+void *
+malloc(size_t size)
+{
+    count_alloc(size);
+    return __libc_malloc(size);
+}
+
+void *
+calloc(size_t count, size_t size)
+{
+    count_alloc(count * size);
+    return __libc_calloc(count, size);
+}
+
+void *
+realloc(void *ptr, size_t size)
+{
+    count_alloc(size);
+    return __libc_realloc(ptr, size);
+}
+
+void *
+memalign(size_t align, size_t size)
+{
+    count_alloc(size);
+    return __libc_memalign(align, size);
+}
+
+void *
+aligned_alloc(size_t align, size_t size)
+{
+    return memalign(align, size);
+}
+
+int
+posix_memalign(void **out, size_t align, size_t size)
+{
+    void *p = memalign(align, size);
+    if (p == NULL) {
+        return ENOMEM;
+    }
+    *out = p;
+    return 0;
+}
+
+void
+free(void *ptr)
+{
+    count_free(ptr);
+    __libc_free(ptr);
+}
+
+#else
+#error "alloc_counter supports glibc and macOS only"
+#endif
@@ -0,0 +1,139 @@
+#!/usr/bin/env python3
+# SPDX-License-Identifier: GPL-3.0-or-later
+"""Count heap allocations per proxied TCP connection in ss-local and ss-server.
+
+Both peers run with the alloc_counter shim preloaded. After a warm-up the
+script snapshots each process's allocation counter, relays more connections,
+and snapshots again. The difference, divided by the number of connections,
+is what one connection costs in the steady state. The no-heap goal for the
+C++ core is zero.
+
+Report-only by default; --max-per-connection N turns it into a check.
+
+Example:
+    python3 tests/alloc_audit/measure_alloc.py --bin build/bin \\
+        --shim build/lib/libss_alloc_audit.dylib --method aes-256-gcm
+"""
+import argparse
+import base64
+import concurrent.futures
+import json
+import os
+from pathlib import Path
+import signal
+import sys
+import tempfile
+import threading
+import time
+
+sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
+from interop import TCPHandler, TCPOrigin, free_port, peers, tcp_case  # noqa: E402
+
+
+def read_snapshot(path, previous_lines, timeout=10):
+    deadline = time.monotonic() + timeout
+    while time.monotonic() < deadline:
+        lines = path.read_text().splitlines() if path.exists() else []
+        if len(lines) > previous_lines:
+            pid, allocations, frees, size = (int(v) for v in lines[-1].split())
+            return {"pid": pid, "allocations": allocations, "frees": frees, "bytes": size}, len(lines)
+        time.sleep(.02)
+    raise AssertionError(f"no allocation snapshot written to {path}")
+
+
+def snapshot_all(processes, outputs, counts):
+    results = []
+    for index, (process, path) in enumerate(zip(processes, outputs)):
+        os.kill(process.pid, signal.SIGUSR2)
+        result, counts[index] = read_snapshot(path, counts[index])
+        results.append(result)
+    return results
+
+
+def relay(local_port, origin_port, connections, concurrency, size):
+    with concurrent.futures.ThreadPoolExecutor(max_workers=concurrency) as pool:
+        futures = [pool.submit(tcp_case, local_port, origin_port, size) for _ in range(connections)]
+        for future in futures:
+            future.result(timeout=60)
+
+
+def main():
+    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
+    parser.add_argument("--bin", default="build/bin", help="directory with ss-local and ss-server")
+    parser.add_argument("--shim", required=True, help="path to the alloc_counter shared library")
+    parser.add_argument("--method", default="aes-256-gcm")
+    parser.add_argument("--warmup", type=int, default=50, help="connections before the first snapshot")
+    parser.add_argument("--connections", type=int, default=200, help="connections between snapshots")
+    parser.add_argument("--concurrency", type=int, default=10)
+    parser.add_argument("--size", type=int, default=16384, help="bytes echoed per connection")
+    parser.add_argument("--max-per-connection", type=float, help="fail if any peer exceeds this")
+    parser.add_argument("--json", action="store_true", help="print the result as JSON")
+    args = parser.parse_args()
+
+    binary_dir = Path(args.bin).resolve()
+    shim = str(Path(args.shim).resolve())
+    preload = "DYLD_INSERT_LIBRARIES" if sys.platform == "darwin" else "LD_PRELOAD"
+    key_size = 16 if args.method == "2022-blake3-aes-128-gcm" else 32
+    password = base64.b64encode(os.urandom(key_size)).decode("ascii")
+
+    with TCPOrigin(("127.0.0.1", 0), TCPHandler) as origin, \
+            tempfile.TemporaryDirectory(prefix="ss-alloc-") as temp:
+        threading.Thread(target=origin.serve_forever, daemon=True).start()
+        server_port, local_port = free_port(), free_port()
+        names = ["ss-server", "ss-local"]
+        outputs = [Path(temp) / f"{name}.counts" for name in names]
+        commands = [
+            [str(binary_dir / "ss-server"), "-s", "127.0.0.1", "-p", str(server_port),
+             "-k", password, "-m", args.method],
+            [str(binary_dir / "ss-local"), "-s", "127.0.0.1", "-p", str(server_port),
+             "-l", str(local_port), "-k", password, "-m", args.method],
+        ]
+        environments = []
+        for output in outputs:
+            env = os.environ.copy()
+            env[preload] = shim
+            env["SS_ALLOC_AUDIT_OUT"] = str(output)
+            environments.append(env)
+
+        # peers() takes one environment for all of its commands, so each
+        # peer gets its own context and its own output file.
+        with peers([commands[0]], [server_port], environments[0]) as server, \
+                peers([commands[1]], [local_port], environments[1]) as local:
+            processes = server + local
+            counts = [0, 0]
+            relay(local_port, origin.server_address[1], args.warmup, args.concurrency, args.size)
+            time.sleep(.5)
+            before = snapshot_all(processes, outputs, counts)
+            relay(local_port, origin.server_address[1], args.connections, args.concurrency, args.size)
+            time.sleep(.5)
+            after = snapshot_all(processes, outputs, counts)
+
+    report = {"method": args.method, "connections": args.connections, "size": args.size, "peers": {}}
+    failed = False
+    for name, b, a in zip(names, before, after):
+        delta = a["allocations"] - b["allocations"]
+        per = delta / args.connections
+        report["peers"][name] = {
+            "allocations": delta,
+            "frees": a["frees"] - b["frees"],
+            "bytes": a["bytes"] - b["bytes"],
+            "allocations_per_connection": round(per, 2),
+        }
+        if args.max_per_connection is not None and per > args.max_per_connection:
+            failed = True
+
+    if args.json:
+        print(json.dumps(report, indent=2))
+    else:
+        print(f"method {args.method}, {args.connections} connections of {args.size} bytes")
+        for name, peer in report["peers"].items():
+            print(f"  {name:9} {peer['allocations_per_connection']:8.2f} allocations/connection "
+                  f"({peer['allocations']} total, {peer['bytes']} bytes)")
+    if failed:
+        print(f"FAIL: more than {args.max_per_connection} allocations per connection")
+        return 1
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())
@@ -159,7 +159,7 @@ def peers(commands, ports, environment=None):
                process = subprocess.Popen(command, stdout=log, stderr=subprocess.STDOUT, env=environment)
                processes.append(process)
                wait_ready(process, port)
-            yield
+            yield processes
        except BaseException:
            for path in logs:
                print(path.name + ":\n" + path.read_text(errors="replace")[-6000:], flush=True)