From bfd78c82e20957c08bb4ee7f891ae7400c5760fd Mon Sep 17 00:00:00 2001 From: Kevin McCoy Date: Sat, 1 Aug 2026 20:35:09 -0400 Subject: [PATCH] Add opt-in exclusive model-container ownership Take a non-blocking advisory POSIX lock before model planning and allocation when the host requests exclusivity. Share it across same-process contexts and release it on every open failure and the final close. Fail open when locking is unavailable, while returning WASTE_E_BUSY for actual contention. Expose the opt-in through the C API, CLI, and server, with focused lifecycle and fallback tests. --- CHANGELOG.md | 35 +++++++ Makefile | 5 +- cli/main.c | 10 +- docs/ENGINE.md | 24 +++++ docs/SERVE.md | 1 + serve/__main__.py | 10 +- serve/engine.py | 8 +- src/waste.c | 172 +++++++++++++++++++++++++++++- src/waste.h | 13 ++- tests/run.sh | 15 +++ tests/test_lock.c | 262 ++++++++++++++++++++++++++++++++++++++++++++++ 11 files changed, 545 insertions(+), 10 deletions(-) create mode 100644 tests/test_lock.c diff --git a/CHANGELOG.md b/CHANGELOG.md index 0608086..40cc214 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,6 +8,41 @@ measurement is the useful part. `docs/LEARNED.md` carries the full reasoning; this file carries what changed. Each entry names the section to read for the numbers behind it. +## 0.6.7 — 2026-08-10 + +The container format, arithmetic, routing, and default open behavior are +unchanged. A host can now ask cooperating WASTE processes for exclusive +ownership of one container before model-sized allocation begins. This is a +host-policy mechanism for a workstation daemon, not RAM accounting or a data +lock: containers remain read-only, concurrent opens remain the default, and an +unsupported advisory lock does not make a readable model unavailable. + +**Callers must recompile against this header.** `exclusive_open` was appended +to `waste_cfg`, and `WASTE_E_BUSY` was added to `waste_status`. The library is +still pre-1.0 and does not promise a stable ABI; `serve/engine.py`'s ctypes +mirror moved with the C header. + +### Added + +- **Opt-in single-process container ownership** + ([#29](https://github.com/sqliteai/waste/pull/29)). On POSIX hosts, + `waste_cfg.exclusive_open`, or `--exclusive-open` in the CLI and server, + takes a non-blocking advisory `flock` on the container directory. Multiple + contexts in one process share a device/inode-keyed reference; a cooperating + process that also requests exclusivity receives `WASTE_E_BUSY`. The last + close and every planning, budget, and partial-load failure release ownership; + descriptors are close-on-exec, and a forked child discards the copied + registry. Windows keeps its existing lifecycle behavior. + +### Fixed + +- **Advisory locking fails open when ownership cannot be established.** Only + actual `EWOULDBLOCK`/`EAGAIN` contention returns `WASTE_E_BUSY`. A directory + that is search-only, a filesystem without `flock`, or another non-contention + locking failure continues through the ordinary model-open path. This keeps + external FUSE, SMB, and NFS containers usable and leaves their real read + errors to the existing loader diagnostics. + ## 0.6.6 — 2026-08-05 The engine decodes exactly as 0.6.5 did and no container format moved. Two diff --git a/Makefile b/Makefile index 34a282c..1fea071 100644 --- a/Makefile +++ b/Makefile @@ -228,7 +228,7 @@ waste$(EXE): cli/main.o libwaste.a # the two failures tests/run.sh was written to catch, so a binary that # `test` builds and `clean` forgets defeats the check meant to notice it. TESTNAMES := test_kda test_container test_forward test_tokenizer test_k3parts \ - test_state test_vision test_image test_memory test_cpus sweep + test_state test_vision test_image test_memory test_cpus test_lock sweep TESTBINS := $(addsuffix $(EXE),$(TESTNAMES)) test: $(TESTBINS) @@ -274,6 +274,9 @@ test_memory$(EXE): tests/test_memory.o src/memory.o test_cpus$(EXE): tests/test_cpus.o $(CC) $(CFLAGS) -o $@ $^ $(LDLIBS) +test_lock$(EXE): tests/test_lock.o libwaste.a + $(CC) $(CFLAGS) -o $@ $^ $(LDLIBS) + %.o: %.c $(CC) $(CFLAGS) -c -o $@ $< diff --git a/cli/main.c b/cli/main.c index b43b980..78a8dd7 100644 --- a/cli/main.c +++ b/cli/main.c @@ -112,7 +112,7 @@ typedef struct { uint64_t budget; uint32_t ctx, max_tokens; float temperature, top_p; - int top_k, threads, quiet, learn, json, no_echo; + int top_k, threads, quiet, learn, json, no_echo, exclusive_open; int media_inlined; /* the media block is already in the prompt string, inside the user turn */ uint64_t seed; @@ -190,6 +190,7 @@ static int parse_opts(int argc, char **argv, int from, opts *o) else if (!strcmp(a, "--json")) o->json = 1; else if (!strcmp(a, "--raw")) o->raw = 1; else if (!strcmp(a, "--verify")) o->verify = 1; + else if (!strcmp(a, "--exclusive-open")) o->exclusive_open = 1; else if (!strcmp(a, "-")) { /* explicit stdin */ if (o->n_pos >= MAX_POS) { fprintf(stderr, "too many arguments\n"); return -1; } o->pos[o->n_pos++] = "-"; @@ -273,6 +274,7 @@ static waste_status open_model(const char *path, const opts *o, waste_ctx **ctx) * Kimi-Linear. Worth it for a container that was copied or downloaded * and has not been read since. */ cfg.verify_records = o->verify; + cfg.exclusive_open = o->exclusive_open; const waste_status st = waste_open(path, &cfg, ctx); /* Two statuses that say nothing useful on their own when --cpus is * what produced them, and it usually is: nothing else here can be @@ -283,6 +285,10 @@ static waste_status open_model(const char *path, const opts *o, waste_ctx **ctx) else if (o->cpus && st == WASTE_E_UNSUPPORTED) fprintf(stderr, "--cpus: this platform does not bind threads to " "CPUs (Linux and Windows only)\n"); + else if (o->exclusive_open && st == WASTE_E_BUSY) + fprintf(stderr, "--exclusive-open: another process owns this " + "container; stop it or retry without " + "--exclusive-open\n"); return st; } @@ -1220,6 +1226,7 @@ int main(int argc, char **argv) "options: --budget 8G --ctx N -n N --temp F --top-p F\n" " --top-k N --seed N --threads N --cpus LIST\n" " --stop STR --file F --json -q --learn --verify\n" + " --exclusive-open\n" " --stop ends generation when the text appears\n" " --json machine-readable output for eval, tokenize, plan,\n" " info and bench\n" @@ -1231,6 +1238,7 @@ int main(int argc, char **argv) " the cores differ: on a two-die Ryzen, six threads on one die\n" " measured 16-25%% faster than six split across both. Linux and\n" " Windows; the default is to leave placement to the OS\n" + " --exclusive-open asks for single-process container ownership\n" " --verify checks each expert record's checksum as it is read,\n" " for a container you have not read since copying it. Costs ~5%%\n" " on Kimi-Linear, ~1%% on K3; off otherwise\n", diff --git a/docs/ENGINE.md b/docs/ENGINE.md index 70efe47..139a0b2 100644 --- a/docs/ENGINE.md +++ b/docs/ENGINE.md @@ -23,6 +23,30 @@ state save/load, model introspection and aggregate stats. Deliberately *not* in the API: logging to stdout, signal handlers, config files, argument parsing. Those belong to the host — the CLI included. +### Optional container ownership + +On POSIX hosts, `waste_cfg.exclusive_open` asks `waste_open` to take a +non-blocking advisory lock on the container directory before memory planning or +model-sized allocation. A cooperating process that also requests exclusivity +for the same container receives `WASTE_E_BUSY`; it does not wait. Paths are +matched by device and inode, so aliases of one directory do not evade the +check. + +Contexts in one process remain independent as documented: they share a +reference-counted ownership entry, and the last `waste_close` releases it. +Failures during planning, budget validation, or partial model loading release +it as well. Lock descriptors are close-on-exec, and a forked child is treated +as a different process rather than inheriting the parent's registry. + +Concurrent opens remain the default. Containers are read-only, and process +ownership is host policy rather than a data-safety requirement; a workstation +daemon can opt in through `waste_cfg.exclusive_open`, or `--exclusive-open` on +the CLI and server. This is an advisory lock between cooperating WASTE +processes, not RAM accounting or a security boundary. If the directory cannot +be opened for locking or the filesystem does not support `flock`, the model +continues without ownership. Windows keeps its existing lifecycle behavior and +ignores the setting. + ## 2. CLI as a first-class client `cli/` links the library and adds only host concerns: argv parsing, a diff --git a/docs/SERVE.md b/docs/SERVE.md index 51b71c6..02e0bd8 100644 --- a/docs/SERVE.md +++ b/docs/SERVE.md @@ -292,6 +292,7 @@ python3 -m serve MODEL [options] --vision load the vision tower --verify check every expert record's crc32 as it is read --usage PATH learned hotlist (default /usage.waste) + --exclusive-open ask for POSIX single-process container ownership --max-tokens N default cap when a request does not set one (4096) --no-thinking answer without the think channel unless asked --allow-local-images diff --git a/serve/__main__.py b/serve/__main__.py index cb2c078..ebc9c01 100644 --- a/serve/__main__.py +++ b/serve/__main__.py @@ -23,7 +23,7 @@ from . import api # noqa: E402 from .engine import (CACHE_LFRU, CACHE_LRU, # noqa: E402 - WASTE_E_ARG, WASTE_E_UNSUPPORTED, + WASTE_E_ARG, WASTE_E_BUSY, WASTE_E_UNSUPPORTED, Engine, EngineError, build_info, physical_ram, plan_memory) from .server import serve # noqa: E402 @@ -121,6 +121,8 @@ def main(argv=None) -> int: "since. Costs ~5%% on Kimi-Linear, ~1%% on K3") g.add_argument("--usage", default=None, metavar="PATH", help="learned hotlist (default /usage.waste)") + g.add_argument("--exclusive-open", action="store_true", + help="ask for single-process ownership of this container") s = ap.add_argument_group("serving") s.add_argument("--max-tokens", type=bounded_int(1, (1 << 32) - 1), @@ -177,7 +179,8 @@ def main(argv=None) -> int: direct_io=not args.no_direct_io, vision=args.vision, verify_records=args.verify, - usage_path=args.usage) + usage_path=args.usage, + exclusive_open=args.exclusive_open) except EngineError as e: print(f"{e}", file=sys.stderr) # Two statuses that say nothing useful on their own when --cpus is @@ -187,6 +190,9 @@ def main(argv=None) -> int: elif args.cpus and e.status == WASTE_E_UNSUPPORTED: print("--cpus: this platform does not bind threads to CPUs " "(Linux and Windows only)", file=sys.stderr) + elif args.exclusive_open and e.status == WASTE_E_BUSY: + print("--exclusive-open: another process owns this container; " + "stop it or retry without --exclusive-open", file=sys.stderr) return 1 try: diff --git a/serve/engine.py b/serve/engine.py index e53987d..4fc6cdd 100644 --- a/serve/engine.py +++ b/serve/engine.py @@ -44,6 +44,7 @@ WASTE_E_ARG = -5 WASTE_E_UNSUPPORTED = -6 WASTE_E_CANCELLED = -7 +WASTE_E_BUSY = -8 # waste.h's waste_cache_policy. There is no third: a "pinned" policy was # listed there and never implemented, so it selected LFRU like everything @@ -87,7 +88,8 @@ class Cfg(C.Structure): ("use_direct_io", C.c_int), ("vision", C.c_int), ("verify_records", C.c_int), - ("usage_path", C.c_char_p)] + ("usage_path", C.c_char_p), + ("exclusive_open", C.c_int)] class GenParams(C.Structure): @@ -378,7 +380,8 @@ def __init__(self, model_path: str, *, direct_io: bool = True, vision: bool = False, verify_records: bool = False, - usage_path: Optional[str] = None): + usage_path: Optional[str] = None, + exclusive_open: bool = False): ram_budget_bytes = _bounded_int( "ram_budget_bytes", ram_budget_bytes, 0, (1 << 64) - 1) ctx_tokens = _bounded_int("ctx_tokens", ctx_tokens, 0, (1 << 32) - 1) @@ -408,6 +411,7 @@ def __init__(self, model_path: str, *, # and a temporary would be freed before waste_open reads it. self._usage = usage_path.encode() if usage_path else None cfg.usage_path = self._usage + cfg.exclusive_open = 1 if exclusive_open else 0 st = self.lib.waste_open(self.model_path.encode(), C.byref(cfg), C.byref(self._ctx)) diff --git a/src/waste.c b/src/waste.c index b59cbf6..8f91111 100644 --- a/src/waste.c +++ b/src/waste.c @@ -11,6 +11,7 @@ #include "waste.h" +#include #include #include #include @@ -21,6 +22,12 @@ #include #endif #include +#if !defined(_WIN32) +#include +#include +#include +#include +#endif #include "json.h" #include "memory.h" @@ -29,6 +36,8 @@ #include "tokenizer.h" #include "waste_backend.h" +typedef struct waste_model_lock waste_model_lock; + struct waste_ctx { waste_model m; waste_tok *tok; @@ -41,6 +50,7 @@ struct waste_ctx { char quant[64]; /* composed at open, reported by get_info */ char detail[128]; /* which record failed, for waste_error_detail */ waste_stats stats; + waste_model_lock *model_lock; /* Queued image embeddings, concatenated: img_each[] is how many rows * each queued image contributed, which is what expand needs to know @@ -51,6 +61,148 @@ struct waste_ctx { int img_n; }; +#if !defined(_WIN32) +struct waste_model_lock { + dev_t dev; + ino_t ino; + int fd; + unsigned refs; + pid_t owner; + waste_model_lock *next; +}; + +static pthread_mutex_t model_lock_mu = PTHREAD_MUTEX_INITIALIZER; +static pthread_once_t model_lock_once = PTHREAD_ONCE_INIT; +static waste_model_lock *model_locks; + +/* A forked child is a competing process, not another context in its parent. + * Close its inherited copies before it can consult the copied registry. The + * entries themselves remain allocated in the child: free is not async-signal + * safe, and inherited contexts ignore entries owned by a different pid. */ +static void model_lock_atfork_prepare(void) +{ + pthread_mutex_lock(&model_lock_mu); +} + +static void model_lock_atfork_parent(void) +{ + pthread_mutex_unlock(&model_lock_mu); +} + +static void model_lock_atfork_child(void) +{ + for (waste_model_lock *p = model_locks; p; p = p->next) close(p->fd); + model_locks = NULL; + pthread_mutex_unlock(&model_lock_mu); +} + +static void model_lock_init(void) +{ + /* pthread_atfork handlers cannot be unregistered. A host that loads this + * shared library dynamically must not dlclose it before a later fork, or + * the process would call these handlers after their code was unmapped. */ + (void)pthread_atfork(model_lock_atfork_prepare, model_lock_atfork_parent, + model_lock_atfork_child); +} + +/* One OS lock per container and process. Device/inode identity means aliases + * of the same directory share an entry. The registry supplies the reference + * semantics flock does not: closing one context must not release ownership + * while another context in this process still uses the container. + * + * Replacing the directory underneath an open context can change that identity + * and let the process contend with itself. The registry covers the normal, + * stable-container case; hosts that replace containers must close them first. */ +static waste_model_lock *model_lock_acquire(const char *path, int exclusive, + waste_status *status) +{ + *status = WASTE_OK; + if (!exclusive) return NULL; + + pthread_once(&model_lock_once, model_lock_init); + int flags = O_RDONLY; +#ifdef O_CLOEXEC + flags |= O_CLOEXEC; +#endif + const int fd = open(path, flags); + if (fd < 0) return NULL; +#ifndef O_CLOEXEC + const int fdflags = fcntl(fd, F_GETFD); + if (fdflags < 0 || fcntl(fd, F_SETFD, fdflags | FD_CLOEXEC)) { + close(fd); + return NULL; + } +#endif + struct stat st; + if (fstat(fd, &st)) { close(fd); return NULL; } + + pthread_mutex_lock(&model_lock_mu); + for (waste_model_lock *p = model_locks; p; p = p->next) { + if (p->dev == st.st_dev && p->ino == st.st_ino) { + p->refs++; + pthread_mutex_unlock(&model_lock_mu); + close(fd); + return p; + } + } + + int rc; + do rc = flock(fd, LOCK_EX | LOCK_NB); while (rc && errno == EINTR); + if (rc) { + const int busy = errno == EWOULDBLOCK || errno == EAGAIN; + pthread_mutex_unlock(&model_lock_mu); + close(fd); + if (busy) *status = WASTE_E_BUSY; + return NULL; + } + + waste_model_lock *p = (waste_model_lock *)calloc(1, sizeof *p); + if (!p) { + (void)flock(fd, LOCK_UN); + pthread_mutex_unlock(&model_lock_mu); + close(fd); + return NULL; + } + p->dev = st.st_dev; + p->ino = st.st_ino; + p->fd = fd; + p->refs = 1; + p->owner = getpid(); + p->next = model_locks; + model_locks = p; + pthread_mutex_unlock(&model_lock_mu); + return p; +} + +static void model_lock_release(waste_model_lock *entry) +{ + if (!entry || entry->owner != getpid()) return; + pthread_mutex_lock(&model_lock_mu); + if (--entry->refs == 0) { + waste_model_lock **pp = &model_locks; + while (*pp && *pp != entry) pp = &(*pp)->next; + if (*pp) *pp = entry->next; + (void)flock(entry->fd, LOCK_UN); + close(entry->fd); + free(entry); + } + pthread_mutex_unlock(&model_lock_mu); +} +#else +/* Keep non-POSIX lifecycle behavior unchanged. The public opt-in is ignored + * on hosts where this advisory ownership lock is not implemented. */ +struct waste_model_lock { int unused; }; +static waste_model_lock *model_lock_acquire(const char *path, int exclusive, + waste_status *status) +{ + (void)path; + (void)exclusive; + *status = WASTE_OK; + return NULL; +} +static void model_lock_release(waste_model_lock *entry) { (void)entry; } +#endif + /* What the container is actually stored as, composed once at open. It used @@ -87,6 +239,7 @@ const char *waste_strerror(waste_status s) case WASTE_E_ARG: return "invalid argument"; case WASTE_E_UNSUPPORTED: return "unsupported"; case WASTE_E_CANCELLED: return "cancelled by callback"; + case WASTE_E_BUSY: return "container is already open in another process"; } return "unknown error"; } @@ -408,8 +561,17 @@ waste_status waste_open(const char *model_path, const waste_cfg *cfg_in, c->cfg = cfg; snprintf(c->path, sizeof c->path, "%s", model_path); + waste_status lock_status = WASTE_OK; + c->model_lock = model_lock_acquire(model_path, cfg.exclusive_open, + &lock_status); + if (lock_status != WASTE_OK) { free(c); return lock_status; } + waste_status st = waste_plan_memory(model_path, cfg.ctx_tokens, &c->plan); - if (st != WASTE_OK) { free(c); return st; } + if (st != WASTE_OK) { + model_lock_release(c->model_lock); + free(c); + return st; + } /* Optional vision weights, decode buffers, tower activations and queued * embeddings are real memory, so all of them enter the floor. */ @@ -457,7 +619,11 @@ waste_status waste_open(const char *model_path, const waste_cfg *cfg_in, if (!cap || b <= cap) { budget = b; break; } } } - if (budget < c->plan.floor_bytes) { free(c); return WASTE_E_RAM_BUDGET; } + if (budget < c->plan.floor_bytes) { + model_lock_release(c->model_lock); + free(c); + return WASTE_E_RAM_BUDGET; + } /* A budget close to physical RAM backfires: the OS starts paging out * the engine's own expert cache, and a "hit" then costs a page fault @@ -495,6 +661,7 @@ waste_status waste_open(const char *model_path, const waste_cfg *cfg_in, * fail, and freeing only the context left all of it behind — * on K3 that is tens of gigabytes lost to one bad manifest. */ waste_model_free(&c->m); + model_lock_release(c->model_lock); free(c); return rc == -2 ? WASTE_E_FORMAT : WASTE_E_IO; } @@ -524,6 +691,7 @@ void waste_close(waste_ctx *c) waste_model_free(&c->m); waste_tok_free(c->tok); free(c->img); + model_lock_release(c->model_lock); free(c); } diff --git a/src/waste.h b/src/waste.h index 08c06bf..44b923d 100644 --- a/src/waste.h +++ b/src/waste.h @@ -48,8 +48,8 @@ extern "C" { */ #define WASTE_VERSION_MAJOR 0 #define WASTE_VERSION_MINOR 6 -#define WASTE_VERSION_PATCH 6 -#define WASTE_VERSION_STRING "0.6.6" +#define WASTE_VERSION_PATCH 7 +#define WASTE_VERSION_STRING "0.6.7" #define WASTE_VERSION_NUMBER (WASTE_VERSION_MAJOR * 10000 + \ WASTE_VERSION_MINOR * 100 + \ WASTE_VERSION_PATCH) @@ -70,6 +70,7 @@ typedef enum { WASTE_E_ARG = -5, WASTE_E_UNSUPPORTED = -6, /* arch/quant combination not built in */ WASTE_E_CANCELLED = -7, /* callback asked to stop */ + WASTE_E_BUSY = -8, /* another process owns this container */ } waste_status; /* Human-readable, static storage; never NULL. A coarse answer by design: @@ -222,6 +223,14 @@ typedef struct { * expert this container does not have are skipped, because it is one * of the few files the engine reads that nobody asked it to. */ const char *usage_path; + + /* Ask for single-process ownership of this container on POSIX hosts. + * Multiple contexts in that process share the ownership; a cooperating + * competing process that also requests exclusivity receives WASTE_E_BUSY + * before model-sized allocations begin. Off by default: containers are + * read-only, and whether concurrent loads are acceptable is host policy. + * Unsupported locks fail open, and non-POSIX hosts ignore this setting. */ + int exclusive_open; } waste_cfg; /* Removed in 0.6.0, having never done anything: `io_threads` (there is no diff --git a/tests/run.sh b/tests/run.sh index cf35765..dc948bc 100755 --- a/tests/run.sh +++ b/tests/run.sh @@ -53,6 +53,21 @@ else exit 1 fi +head_ "model-container ownership" +# This intentionally opens two contexts at once. Always give it a tiny +# dedicated container rather than duplicating a caller-supplied K3 load. +LOCK_MODEL="$TMP/lock-test.waste" +if ! python3 tools/make_test_container.py "$LOCK_MODEL" \ + >"$TMP/lock-container.log" 2>&1; then + sk "model-container ownership lock" \ + "could not build its synthetic container" +elif ./test_lock "$LOCK_MODEL" "$TMP" >"$TMP/lock.log" 2>&1; then + ok "opt-in process exclusion, references, fail-open and cleanup" +else + no "model-container ownership lock" + head -20 "$TMP/lock.log" +fi + # ---------------------------------------------------------------- unit ---- head_ "kernels vs the reference implementations" diff --git a/tests/test_lock.c b/tests/test_lock.c new file mode 100644 index 0000000..66b0444 --- /dev/null +++ b/tests/test_lock.c @@ -0,0 +1,262 @@ +/* SPDX-License-Identifier: Apache-2.0 */ +/* test_lock.c — container ownership is process-wide and leak-free. */ +#include "../src/waste.h" + +#include +#include +#include + +#if !defined(_WIN32) +#include +#include +#include +#include +#include +#include +#include + +static int failed; + +#define CHECK(expr, what) do { \ + if (!(expr)) { \ + fprintf(stderr, "FAIL line %d: %s\n", __LINE__, (what)); \ + failed = 1; \ + } \ +} while (0) + +static int copy_file(const char *src, const char *dst) +{ + const int in = open(src, O_RDONLY); + if (in < 0) return -1; + const int out = open(dst, O_WRONLY | O_CREAT | O_TRUNC, 0600); + if (out < 0) { close(in); return -1; } + char buf[16384]; + int rc = 0; + for (;;) { + ssize_t n; + do n = read(in, buf, sizeof buf); while (n < 0 && errno == EINTR); + if (n <= 0) { if (n < 0) rc = -1; break; } + ssize_t off = 0; + while (off < n) { + ssize_t put; + do put = write(out, buf + off, (size_t)(n - off)); + while (put < 0 && errno == EINTR); + if (put <= 0) { rc = -1; break; } + off += put; + } + if (rc) break; + } + if (close(out)) rc = -1; + close(in); + return rc; +} + +static int write_text(const char *path, const char *text) +{ + const int fd = open(path, O_WRONLY | O_CREAT | O_TRUNC, 0600); + if (fd < 0) return -1; + const size_t n = strlen(text); + const int rc = write(fd, text, n) == (ssize_t)n ? 0 : -1; + return close(fd) ? -1 : rc; +} + +/* A separately opened descriptor must not be able to take the directory's + * flock while the library owns it. */ +static int raw_lock_available(const char *model) +{ + const int fd = open(model, O_RDONLY); + if (fd < 0) return 0; + int rc; + do rc = flock(fd, LOCK_EX | LOCK_NB); while (rc && errno == EINTR); + if (!rc) (void)flock(fd, LOCK_UN); + close(fd); + return rc == 0; +} + +static int probe(const char *model, int exclusive, uint64_t budget, + waste_status expected) +{ + waste_cfg cfg; + waste_cfg_init(&cfg); + cfg.exclusive_open = exclusive; + cfg.ram_budget_bytes = budget; + waste_ctx *ctx = NULL; + const waste_status got = waste_open(model, &cfg, &ctx); + if (ctx) waste_close(ctx); + if (got != expected) { + fprintf(stderr, "probe: got %d (%s), expected %d (%s)\n", + got, waste_strerror(got), expected, waste_strerror(expected)); + return 1; + } + return 0; +} + +/* Exec makes this a genuinely separate library instance rather than relying + * on fork's copied registry. It also verifies the lock descriptor is closed + * across exec while the parent's descriptor continues to own the lock. */ +static int child_probe(const char *self, const char *model, int exclusive, + uint64_t budget, waste_status expected) +{ + const pid_t pid = fork(); + if (pid < 0) return 1; + if (pid == 0) { + char a[2], b[32], e[16]; + snprintf(a, sizeof a, "%d", exclusive); + snprintf(b, sizeof b, "%llu", (unsigned long long)budget); + snprintf(e, sizeof e, "%d", expected); + execlp(self, self, "--probe", model, a, b, e, (char *)NULL); + _exit(127); + } + int ws; + pid_t got; + do got = waitpid(pid, &ws, 0); while (got < 0 && errno == EINTR); + return got != pid || !WIFEXITED(ws) || WEXITSTATUS(ws) != 0; +} + +static void remove_variant(const char *dir, int trunk) +{ + char path[1024]; + snprintf(path, sizeof path, "%s/manifest.json", dir); + (void)unlink(path); + if (trunk) { + snprintf(path, sizeof path, "%s/trunk.bin", dir); + (void)unlink(path); + } + (void)rmdir(dir); +} + +int main(int argc, char **argv) +{ + if (argc == 6 && !strcmp(argv[1], "--probe")) { + const int exclusive = atoi(argv[3]); + const uint64_t budget = (uint64_t)strtoull(argv[4], NULL, 10); + const waste_status expected = (waste_status)strtol(argv[5], NULL, 10); + return probe(argv[2], exclusive, budget, expected); + } + if (argc != 3) { + fprintf(stderr, "usage: %s MODEL SCRATCH-DIR\n", argv[0]); + return 2; + } + const char *model = argv[1]; + const char *scratch = argv[2]; + + waste_cfg cfg; + waste_cfg_init(&cfg); + CHECK(cfg.exclusive_open == 0, "ownership must default off"); + CHECK(strstr(waste_strerror(WASTE_E_BUSY), "another process") != NULL, + "WASTE_E_BUSY must explain the contention"); + waste_memplan plan; + if (waste_plan_memory(model, cfg.ctx_tokens, &plan) != WASTE_OK) { + fprintf(stderr, "cannot plan lock-test container\n"); + return 1; + } + cfg.ram_budget_bytes = plan.floor_bytes; + + /* Ordinary opens do not take ownership. A second process therefore + * reaches budget validation instead of being rejected by container + * identity. */ + waste_ctx *ordinary = NULL; + CHECK(waste_open(model, &cfg, &ordinary) == WASTE_OK, "default open"); + CHECK(ordinary != NULL, "default context"); + CHECK(raw_lock_available(model), "ownership must be opt-in"); + if (ordinary) waste_close(ordinary); + + /* A search-only directory is still a readable container: known files can + * be opened through it, but opening the directory itself for flock fails. + * Exclusive ownership is advisory, so that lock failure must not turn a + * model the engine can read into an open failure. */ + struct stat model_mode; + const int have_mode = stat(model, &model_mode) == 0; + CHECK(have_mode, "read container permissions"); + if (have_mode) { + const int search_only = chmod(model, 0111) == 0; + CHECK(search_only, "make container search-only"); + if (search_only) { + CHECK(probe(model, 1, plan.floor_bytes, WASTE_OK) == 0, + "locking failure must proceed without ownership"); + } + CHECK(chmod(model, model_mode.st_mode & 0777) == 0, + "restore container permissions"); + } + + cfg.exclusive_open = 1; + + waste_ctx *a = NULL, *b = NULL; + CHECK(waste_open(model, &cfg, &a) == WASTE_OK, "first open"); + CHECK(a != NULL, "first context"); + if (!a) return 1; + CHECK(!raw_lock_available(model), "open context must own directory"); + + /* The second context shares the process entry instead of contending with + * its own flock. Keeping it open exercises reference-counted release. */ + CHECK(waste_open(model, &cfg, &b) == WASTE_OK, "same-process second open"); + CHECK(b != NULL, "second context"); + CHECK(child_probe(argv[0], model, 1, 1, WASTE_E_BUSY) == 0, + "competing process must receive WASTE_E_BUSY before budgeting"); + CHECK(child_probe(argv[0], model, 0, 1, WASTE_E_RAM_BUDGET) == 0, + "default open must bypass optional ownership"); + + waste_close(a); + a = NULL; + CHECK(!raw_lock_available(model), + "closing one context must retain the other context's ownership"); + CHECK(child_probe(argv[0], model, 1, 1, WASTE_E_BUSY) == 0, + "remaining same-process reference must exclude competitors"); + + waste_close(b); + b = NULL; + CHECK(raw_lock_available(model), "last close must release ownership"); + CHECK(child_probe(argv[0], model, 1, 1, WASTE_E_RAM_BUDGET) == 0, + "a new process must pass ownership after normal close"); + + /* Every return after acquisition must release the OS lock. */ + CHECK(probe(model, 1, 1, WASTE_E_RAM_BUDGET) == 0, + "budget failure status"); + CHECK(raw_lock_available(model), "budget failure must release ownership"); + + char bad_plan[1024], bad_load[1024], src[1024], dst[1024]; + snprintf(bad_plan, sizeof bad_plan, "%s/lock-bad-plan-%ld", scratch, + (long)getpid()); + CHECK(mkdir(bad_plan, 0700) == 0, "create malformed variant"); + snprintf(dst, sizeof dst, "%s/manifest.json", bad_plan); + CHECK(write_text(dst, "{") == 0, "write malformed manifest"); + waste_ctx *ctx = NULL; + const waste_status plan_st = waste_open(bad_plan, &cfg, &ctx); + CHECK(plan_st == WASTE_E_FORMAT && ctx == NULL, "planning failure status"); + CHECK(raw_lock_available(bad_plan), "planning failure must release ownership"); + remove_variant(bad_plan, 0); + + /* Copy enough for planning and trunk allocation, then omit codebooks so + * model loading fails after partial construction. */ + snprintf(bad_load, sizeof bad_load, "%s/lock-bad-load-%ld", scratch, + (long)getpid()); + CHECK(mkdir(bad_load, 0700) == 0, "create partial-load variant"); + snprintf(src, sizeof src, "%s/manifest.json", model); + snprintf(dst, sizeof dst, "%s/manifest.json", bad_load); + CHECK(copy_file(src, dst) == 0, "copy partial manifest"); + snprintf(src, sizeof src, "%s/trunk.bin", model); + snprintf(dst, sizeof dst, "%s/trunk.bin", bad_load); + CHECK(copy_file(src, dst) == 0, "copy partial trunk"); + ctx = NULL; + const waste_status load_st = waste_open(bad_load, &cfg, &ctx); + CHECK(load_st != WASTE_OK && load_st != WASTE_E_BUSY && ctx == NULL, + "partial model-load failure status"); + CHECK(raw_lock_available(bad_load), + "partial model-load failure must release ownership"); + remove_variant(bad_load, 1); + + if (failed) return 1; + puts("PASS model-container ownership lock"); + return 0; +} + +#else +int main(void) +{ + waste_cfg cfg; + waste_cfg_init(&cfg); + if (cfg.exclusive_open != 0) return 1; + puts("PASS model-container ownership lock (not used on this platform)"); + return 0; +} +#endif