The tenant that leaked across requests

import java.util.*;
import java.util.concurrent.*;

/**
 * Solution: the tenant that leaked across requests.
 *
 * The gateway puts the caller's tenant into a ThreadLocal at the start of a
 * request so that every layer below can read it without threading it through
 * forty signatures. A filter clears it at the end.
 *
 * Two tickets, and only the first was taken seriously:
 *
 *   1. SECURITY: a customer of tenant-b saw a report belonging to tenant-a.
 *      Intermittent, only under load, never reproducible locally. The report
 *      code is correct and reads the tenant from the context, as designed.
 *
 *   2. MEMORY: the service's heap floor creeps up over days. The retained set
 *      is dominated by objects reachable from Thread instances.
 *
 * They are the same bug. Find it, fix it, and then answer the design
 * question: what would remove this class of bug rather than this instance.
 *
 * WHAT WAS WRONG
 *   The two tickets were one bug. remove() sat at the end of the happy path
 *   instead of in a finally, so any request that threw left its tenant and
 *   its 512 KB buffer attached to a pooled thread.
 *
 *   The memory ticket is the obvious consequence: pooled threads never die,
 *   so the value is retained for the life of the application, once per
 *   thread that ever failed.
 *
 *   The security ticket is the same fact seen from the other side. The next
 *   request to land on that thread inherits the stale context. A request
 *   that SETS a tenant hides this by overwriting it, which is why it looked
 *   intermittent — only a request with no tenant of its own, like the health
 *   check here, actually reads the leftover.
 *
 * WHAT WOULD REMOVE THE CLASS OF BUG
 *   Virtual threads are not pooled: the thread really does die after its
 *   task, so the entry goes with it. ScopedValue (Java 21+) is the intended
 *   replacement for this pattern outright — the value is bound for the
 *   duration of a call and unbound when it returns, with no remove() to
 *   forget. Where neither is available, the filter that sets the context
 *   should be the only code that can, and it should use try/finally once
 *   rather than trusting every handler to.
 */
public class Solution {

    /** Set per request. Read by every layer below. */
    static final ThreadLocal<String> TENANT = new ThreadLocal<>();

    /** Stands in for whatever the request actually retained. */
    static final ThreadLocal<byte[]> REQUEST_BUFFER = new ThreadLocal<>();

    record Report(String forTenant, String servedTenant) {
        boolean leaked() {
            return !forTenant.equals(servedTenant);
        }
    }

    /**
     * A layer far below the filter. It does the right thing: it does not
     * take the tenant as a parameter, it reads it from the context.
     */
    static String currentTenant() {
        String tenant = TENANT.get();
        return tenant == null ? "<none>" : tenant;
    }

    /**
     * FIX: remove() in a finally.
     *
     * The clear has to run on the path that threw, because that is the only
     * path where it was ever going to matter. A ThreadLocal is cleaned up
     * when its thread dies, and a pooled thread is built not to die — so
     * "the next request will overwrite it" is not a cleanup strategy. It is
     * also not true for a request that has no tenant of its own.
     */
    static Report handleRequest(String tenant, boolean willFail) {
        TENANT.set(tenant);
        REQUEST_BUFFER.set(new byte[512 * 1024]);
        try {
            if (willFail) {
                throw new IllegalStateException("downstream timeout for " + tenant);
            }
            return new Report(tenant, currentTenant());
        } finally {
            TENANT.remove();
            REQUEST_BUFFER.remove();
        }
    }

    /**
     * An internal health check. It has no tenant — that is correct and
     * intended — and it reads the context only to label its audit line.
     *
     * This is the request that gets hurt, and it does nothing wrong.
     */
    static Report handleInternalRequest() {
        return new Report("<none>", currentTenant());
    }

    public static void main(String[] args) throws Exception {
        // One thread, so the interleaving is deterministic. In production it
        // is fifty, which only changes how often you hit it.
        ExecutorService pool = Executors.newFixedThreadPool(1);

        List<Report> served = new ArrayList<>();

        System.out.println("── four requests, one pooled thread ──");

        // Request 1: tenant-a, and it fails downstream.
        try {
            pool.submit(() -> handleRequest("tenant-a", true)).get();
        } catch (ExecutionException expected) {
            System.out.println("  1  tenant-a  failed: " + expected.getCause().getMessage());
        }

        // Request 2: the health check. No tenant of its own. Reads the
        // context exactly as every other layer does.
        Report internal = pool.submit(Solution::handleInternalRequest).get();
        served.add(internal);
        System.out.println("  2  internal  attributed to: " + internal.servedTenant()
                + (internal.leaked() ? "   <-- it has no tenant" : ""));

        // Request 3: tenant-c, completes normally.
        Report third = pool.submit(() -> handleRequest("tenant-c", false)).get();
        served.add(third);
        System.out.println("  3  tenant-c  served: " + third.servedTenant());

        // Request 4: tenant-d, fails too.
        try {
            pool.submit(() -> handleRequest("tenant-d", true)).get();
        } catch (ExecutionException expected) {
            System.out.println("  4  tenant-d  failed: " + expected.getCause().getMessage());
        }

        // What is still attached to the pooled thread once everything is done?
        String strandedTenant = pool.submit(TENANT::get).get();
        byte[] strandedBuffer = pool.submit(REQUEST_BUFFER::get).get();
        pool.shutdown();

        System.out.println();
        boolean noLeakAcrossRequests = served.stream().noneMatch(Report::leaked);
        boolean threadLeftClean = strandedTenant == null && strandedBuffer == null;

        System.out.println("stranded on the thread : tenant=" + strandedTenant
                + ", buffer=" + (strandedBuffer == null ? "none"
                        : (strandedBuffer.length / 1024) + " KB"));
        System.out.println();
        System.out.println("no request saw another's tenant : " + noLeakAcrossRequests);
        System.out.println("thread left clean for the pool  : " + threadLeftClean);
        System.out.println(noLeakAcrossRequests && threadLeftClean ? "PASS" : "FAIL");
    }
}