spring-ai/llm-gateway/src/main/java/com/ankurm/gateway/Budgets.java

package com.ankurm.gateway;

import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.function.ToLongFunction;

/**
 * Dollar caps per tenant, in microdollars. The point of the class is the order of operations:
 * {@link #reserve} takes the money BEFORE the provider is called, and {@link Reservation#settle}
 * corrects it afterwards. Checking the balance and recording the spend as two separate steps lets
 * concurrent requests all pass the check and all spend (see {@code BudgetTest}).
 */
public final class Budgets {

    /** Money held for one call. Settle with the real cost, or release if the call failed. */
    public final class Reservation {

        private final Account account;

        private final long held;

        private boolean done;

        private Reservation(Account account, long held) {
            this.account = account;
            this.held = held;
        }

        public synchronized void settle(long actualMicros) {
            if (!done) {
                done = true;
                account.settle(held, actualMicros);
            }
        }

        public synchronized void release() {
            if (!done) {
                done = true;
                account.settle(held, 0);
            }
        }
    }

    private static final class Account {

        private long spent;

        private long held;

        synchronized boolean reserve(long amount, long cap) {
            if (spent + held + amount > cap) {
                return false;
            }
            held += amount;
            return true;
        }

        synchronized void settle(long heldAmount, long actual) {
            held -= heldAmount;
            spent += actual;
        }

        synchronized long spent() {
            return spent;
        }
    }

    private final Map<String, Account> accounts = new ConcurrentHashMap<>();

    private final ToLongFunction<String> capMicros;

    public Budgets(ToLongFunction<String> capMicros) {
        this.capMicros = capMicros;
    }

    public Reservation reserve(String tenant, long estimateMicros) {
        Account a = accounts.computeIfAbsent(tenant, t -> new Account());
        long cap = capMicros.applyAsLong(tenant);
        if (!a.reserve(estimateMicros, cap)) {
            throw new BudgetExceededException("Budget for " + tenant + " would be exceeded: spent "
                    + a.spent() + " of " + cap + " microdollars, this call may cost up to " + estimateMicros);
        }
        return new Reservation(a, estimateMicros);
    }

    public long spent(String tenant) {
        Account a = accounts.get(tenant);
        return a == null ? 0 : a.spent();
    }
}