%% Copyright (c) 2026 Benoit Chesneau. Licensed under the MIT License.
%% See the LICENSE file at the project root.
%%
-module(erllama_cache_writer).
-moduledoc """
File-tier save orchestrator with a leak-proof ETS counting
semaphore.
Modelled on `py_semaphore` from erlang-python (Discord pattern):
one ETS table holding `{running, N}` and `{max, M}` counters;
atomic `update_counter` for fast inspection from any process.
Unlike the bare ETS pattern, this module also runs a small
gen_server that owns the holders map and monitors every active
acquirer. If a holder dies between acquire and release (SIGKILL,
process crash, etc.), the gen_server's `'DOWN'` handler releases
the slot. No leaks.
Hot path:
`current/0` and `max_concurrent/0` read ETS directly
(no server hop).
`acquire/1` and `release/0` go through this gen_server, which
serialises monitor bookkeeping. Per-call overhead is roughly
one gen_server hop (~1 us); for save paths measured in
seconds this is invisible.
Save pipeline run by `save/4` in the caller's process:
acquire (with backoff up to AcquireTimeoutMs)
meta_srv:reserve_save -> Token
disk_srv:save (returns once the file is linked + validated)
meta_srv:announce_saved -> available
release (try/after)
A `try/after` releases on normal exits and exceptions; the
gen_server's monitor catches everything else.
""".
-behaviour(gen_server).
-include("erllama_cache.hrl").
-export([
start_link/0,
start_link/1,
save/4,
save/5,
acquire/1,
release/0,
current/0,
max_concurrent/0,
set_max_concurrent/1
]).
-export([init/1, handle_call/3, handle_cast/2, handle_info/2]).
-define(SERVER, ?MODULE).
-define(TABLE, erllama_cache_writer_sem).
-define(COUNTER_KEY, running).
-define(MAX_KEY, max).
-define(BACKOFF_MS, 5).
-define(MAX_BACKOFF_MS, 50).
-define(DEFAULT_TIMEOUT_MS, 30000).
-define(SLOT_PD_KEY, {?MODULE, slot}).
-record(state, {
%% MonRef -> CallerPid; one entry per active acquire
holders = #{} :: #{reference() => pid()}
}).
-type state() :: #state{}.
%% =============================================================================
%% Public API: lifecycle
%% =============================================================================
-spec start_link() -> {ok, pid()} | {error, term()}.
start_link() ->
start_link(default_max_concurrent()).
-spec start_link(pos_integer()) -> {ok, pid()} | {error, term()}.
start_link(Max) when is_integer(Max), Max > 0 ->
gen_server:start_link({local, ?SERVER}, ?MODULE, [Max], []).
-spec current() -> non_neg_integer().
current() ->
case ets:lookup(?TABLE, ?COUNTER_KEY) of
[{_, N}] -> N;
[] -> 0
end.
-spec max_concurrent() -> pos_integer().
max_concurrent() ->
case ets:lookup(?TABLE, ?MAX_KEY) of
[{_, Max}] -> Max;
[] -> default_max_concurrent()
end.
-spec set_max_concurrent(pos_integer()) -> ok.
set_max_concurrent(Max) when is_integer(Max), Max > 0 ->
ets:insert(?TABLE, {?MAX_KEY, Max}),
ok.
%% =============================================================================
%% Public API: semaphore
%% =============================================================================
-doc """
Try to acquire a slot. Blocks with exponential backoff up to
Timeout milliseconds. On success the holder ref is stashed in the
process dictionary so `release/0` finds it.
""".
-spec acquire(timeout()) -> ok | {error, max_concurrent}.
acquire(infinity) ->
acquire_loop(infinity, 0, ?BACKOFF_MS);
acquire(Timeout) when is_integer(Timeout), Timeout >= 0 ->
acquire_loop(Timeout, erlang:monotonic_time(millisecond), ?BACKOFF_MS).
-spec release() -> ok.
release() ->
case erase(?SLOT_PD_KEY) of
Ref when is_reference(Ref) ->
gen_server:call(?SERVER, {release, Ref});
undefined ->
ok
end.
%% =============================================================================
%% Public API: save pipeline
%% =============================================================================
-spec save(
atom(),
disk | ram_file,
erllama_cache_kvc:build_meta(),
binary()
) -> {ok, erllama_cache:cache_key()} | {error, term()}.
save(TierSrv, Tier, BuildMeta, Payload) ->
save(TierSrv, Tier, BuildMeta, Payload, ?DEFAULT_TIMEOUT_MS).
-spec save(
atom(),
disk | ram_file,
erllama_cache_kvc:build_meta(),
binary(),
timeout()
) -> {ok, erllama_cache:cache_key()} | {error, term()}.
save(TierSrv, Tier, BuildMeta, Payload, AcquireTimeout) when
Tier =:= disk; Tier =:= ram_file
->
case acquire(AcquireTimeout) of
ok ->
try
do_save(TierSrv, Tier, BuildMeta, Payload)
after
release()
end;
{error, _} = E ->
E
end.
%% =============================================================================
%% gen_server callbacks
%% =============================================================================
-spec init([pos_integer()]) -> {ok, state()}.
init([Max]) ->
case ets:whereis(?TABLE) of
undefined ->
_ = ets:new(?TABLE, [
named_table,
public,
{write_concurrency, true},
{read_concurrency, true}
]),
ets:insert(?TABLE, [{?COUNTER_KEY, 0}, {?MAX_KEY, Max}]);
_Tid ->
ets:insert(?TABLE, {?COUNTER_KEY, 0}),
ets:insert(?TABLE, {?MAX_KEY, Max})
end,
{ok, #state{}}.
handle_call({try_acquire, Pid}, _From, S) ->
Max = max_concurrent(),
case ets:update_counter(?TABLE, ?COUNTER_KEY, {2, 1}) of
N when N =< Max ->
Ref = erlang:monitor(process, Pid),
{reply, {ok, Ref}, S#state{holders = (S#state.holders)#{Ref => Pid}}};
_ ->
decr(),
{reply, {error, busy}, S}
end;
handle_call({release, Ref}, _From, S) ->
case maps:take(Ref, S#state.holders) of
{_Pid, Holders1} ->
erlang:demonitor(Ref, [flush]),
decr(),
{reply, ok, S#state{holders = Holders1}};
error ->
%% Stale release (already cleaned up by 'DOWN').
{reply, ok, S}
end;
handle_call(_Msg, _From, S) ->
{reply, {error, unknown_call}, S}.
-spec handle_cast(term(), state()) -> {noreply, state()}.
handle_cast(_Msg, S) ->
{noreply, S}.
handle_info({'DOWN', Ref, process, _DownPid, _Reason}, S) ->
case maps:take(Ref, S#state.holders) of
{_HolderPid, Holders1} ->
decr(),
{noreply, S#state{holders = Holders1}};
error ->
{noreply, S}
end;
handle_info(_Msg, S) ->
{noreply, S}.
%% =============================================================================
%% Internal: acquire loop
%% =============================================================================
acquire_loop(Timeout, Start, Backoff) ->
case gen_server:call(?SERVER, {try_acquire, self()}) of
{ok, Ref} ->
put(?SLOT_PD_KEY, Ref),
ok;
{error, busy} ->
case check_timeout(Timeout, Start) of
continue ->
Jitter = rand:uniform(Backoff div 2 + 1),
timer:sleep(Backoff + Jitter),
acquire_loop(Timeout, Start, min(Backoff * 2, ?MAX_BACKOFF_MS));
timeout ->
{error, max_concurrent}
end
end.
check_timeout(infinity, _) ->
continue;
check_timeout(Timeout, Start) ->
case erlang:monotonic_time(millisecond) - Start >= Timeout of
true -> timeout;
false -> continue
end.
decr() ->
%% Saturating decrement of the counter; clamps at 0.
_ = ets:update_counter(?TABLE, ?COUNTER_KEY, {2, -1, 0, 0}),
ok.
-spec default_max_concurrent() -> pos_integer().
default_max_concurrent() ->
case application:get_env(erllama, writer_max_concurrent) of
{ok, N} when is_integer(N), N > 0 -> N;
_ -> erlang:system_info(schedulers) * 2 + 1
end.
%% =============================================================================
%% Internal: pipeline
%% =============================================================================
do_save(TierSrv, Tier, BuildMeta, Payload) ->
case derive_key(BuildMeta) of
{ok, Key} ->
case erllama_cache_meta_srv:reserve_save(Key, Tier, self()) of
{ok, Token} ->
do_save_reserved(TierSrv, Key, Token, BuildMeta, Payload);
{error, already_present} ->
erllama_cache_counters:incr(?C_DUPLICATE_DROPPED),
{error, already_present};
{error, _} = E ->
E
end;
{error, _} = E ->
E
end.
derive_key(BuildMeta) ->
try
{ok,
erllama_cache_key:make(#{
fingerprint => maps:get(fingerprint, BuildMeta),
quant_type => maps:get(quant_type, BuildMeta),
ctx_params_hash => maps:get(ctx_params_hash, BuildMeta),
tokens => maps:get(tokens, BuildMeta)
})}
catch
error:{badkey, K} -> {error, {missing_key, K}};
error:Reason -> {error, Reason}
end.
%% Stage-aware reservation flow:
%% reserve_save (already done) → write_tmp → check_reservation
%% → publish (link) → mark_published → announce_saved
%%
%% check_reservation right before link is the cancel-token check that
%% prevents a slow writer (whose reservation TTL expired and got
%% recycled to another writer) from publishing under stale authority.
%% mark_published advances the meta server's stage to post_link so a
%% writer that crashes between link and announce_saved triggers the
%% validate-and-adopt path rather than the placeholder-row delete.
do_save_reserved(TierSrv, Key, Token, BuildMeta, Payload) ->
Root = erllama_cache_disk_srv:dir(TierSrv),
case erllama_cache_disk_srv:write_tmp(Root, BuildMeta, Payload) of
{ok, WriteOut} ->
after_write_tmp(Key, Token, BuildMeta, WriteOut);
{error, _} = E ->
_ = erllama_cache_meta_srv:cancel_reservation(Key, Token),
E
end.
after_write_tmp(Key, Token, BuildMeta, WriteOut) ->
case erllama_cache_meta_srv:check_reservation(Key, Token) of
ok ->
after_check_reservation(Key, Token, BuildMeta, WriteOut);
{error, expired} = E ->
erllama_cache_disk_srv:abort_tmp(WriteOut),
E
end.
after_check_reservation(Key, Token, BuildMeta, WriteOut) ->
case erllama_cache_disk_srv:publish(WriteOut) of
{ok, _SameKey, Header, Size} ->
#{final := FinalPath} = WriteOut,
after_publish(Key, Token, BuildMeta, FinalPath, Header, Size);
{error, _} = E ->
erllama_cache_disk_srv:abort_tmp(WriteOut),
_ = erllama_cache_meta_srv:cancel_reservation(Key, Token),
E
end.
after_publish(Key, Token, BuildMeta, FinalPath, Header, Size) ->
case erllama_cache_meta_srv:mark_published(Key, Token, FinalPath) of
ok ->
TokensBin = erllama_cache_key:encode_tokens(maps:get(tokens, BuildMeta)),
case
erllama_cache_meta_srv:announce_saved(
Key, Token, Size, Header, TokensBin
)
of
ok ->
bump_save_counter(maps:get(save_reason, BuildMeta, unknown)),
{ok, Key};
{error, _} = E ->
%% Token expired between mark_published and announce_saved
%% (writer was slow; TTL sweep ran and validated-and-adopted
%% the linked file under a fresh reservation). The data is
%% on disk and indexed; the writer's caller just sees the
%% stale `{error, expired}` and should treat it as a
%% non-fatal save-skip. The save_counter bump for this
%% reason intentionally does not fire.
E
end;
{error, _} = E ->
%% Token was recycled between check_reservation and
%% mark_published. The file is published; the meta sweep
%% will validate-and-adopt it on its next pass. Same
%% non-fatal outcome as above.
E
end.
bump_save_counter(cold) -> erllama_cache_counters:incr(?C_SAVES_COLD);
bump_save_counter(continued) -> erllama_cache_counters:incr(?C_SAVES_CONTINUED);
bump_save_counter(finish) -> erllama_cache_counters:incr(?C_SAVES_FINISH);
bump_save_counter(evict) -> erllama_cache_counters:incr(?C_SAVES_EVICT);
bump_save_counter(shutdown) -> erllama_cache_counters:incr(?C_SAVES_SHUTDOWN);
bump_save_counter(_) -> ok.
