Reserve run ownership before installing worker deadlines

CHANGELOG.md

@@ -22,10 +22,11 @@
 - Summarize ONNX protobuf graphs without a runtime, including versions, operators, initializer estimates, and ZipMap.
 - Cancel one model-bound inference attempt with immutable `Ortex.RunControl` options.
 - Bound caller wait with per-run timeouts and cancel managed inference when its caller exits.
+- Reserve run ownership before deadlines so duplicate calls cannot cancel another attempt.
 - Support checksum-verified precompiled CPU NIFs and source-build overrides.
 - Default to dynamic runtime loading and API 17, without runtime build headers.
 
 - Build and validate all seven artifacts, checksum maps, and a package in CI.
 - Stage attested draft releases only in RiverFinancial/ortex; no Hex publisher.
 
-Upgrade in a fresh VM; native resource layouts changed (including environment V2).
+Upgrade in a fresh VM; native resource representations changed (including environment and run-control V2).

docs/cancellation.md

@@ -3,8 +3,8 @@
 `Ortex.RunControl.new(model, outputs: [...], tag: "...")` prepares immutable
 options for one model-bound inference attempt. Pass it with `control: control` to
 `Ortex.run/3`. A different process can request termination with
-`Ortex.RunControl.cancel/1`, and `started?/1` reports whether the attempt has been
-claimed, including waiting for the session mutex. A claimed control cannot be
+`Ortex.RunControl.cancel/1`, and `started?/1` reports whether the native attempt
+has been entered, including waiting for the session mutex. A reserved control cannot be
 reused even if inference fails. Invalid options are rejected before an attempt.
 
 Termination calls ONNX Runtime's run-options flag without acquiring the session
@@ -49,6 +49,14 @@ ignored. Worker exceptions are re-raised with their original stack. Completed
 attempts do not trigger watcher cancellation, so rejecting a reused or wrong-model
 control cannot cancel its legitimate in-flight attempt.
 
+Ownership is atomically reserved in the caller before a watcher or deadline is
+installed. Native entry then consumes that reservation exactly once. This also
+rejects duplicates while the original worker has not yet reached its NIF: a
+zero-deadline or exiting duplicate cannot cancel the original attempt. The native
+control resource is V2; a fresh VM is required when upgrading from V1 controls.
+`started?/1` remains false for a reserved attempt until native entry, including
+entry that waits for the model mutex.
+
 Tests exercise actual loop deadlines, caller death, deadlines behind an occupied
 mutex, post-cancel inference, worker exception stacks, resource ownership, worker
 termination, response/deadline races, and lack of late messages. Telemetry emits

docs/integration.md

@@ -38,6 +38,7 @@ writes to another repository. The broad original PR is superseded by the stack.
 | 28 | Managed deadlines and caller-exit cancellation |
 | 29 | Dirty-scheduler measurement and policy |
 | 30 | Final lifecycle/managed-run soak and review ledger |
+| 31 | Atomic control ownership before worker/deadline creation |
 
 Each PR targets the preceding branch. Merge in order with merge commits to
 preserve ancestry. If using squash/rebase merges, rebase and retarget the remaining
@@ -59,7 +60,7 @@ ORTEX_BUILD=true ORT_DYLIB_PATH=/does-not-exist/libonnxruntime.so MIX_ENV=test \
   mix run --no-start scripts/check_binary_without_runtime.exs
 ```
 
-The assembled stack currently has 216 checks, with two optional ResNet checks
+The assembled stack currently has 218 checks, with two optional ResNet checks
 excluded when that external checkpoint is absent. Its fixture contains only
 synthetic data. Memory loading checks file/bytes score and descriptor parity,
 source-owner process exit, malformed-model recovery, and concurrent loading.
@@ -142,6 +143,12 @@ validation PR records exact integration run links without requiring a release ta
 Runtime cancellation is cooperative and may finish after the caller deadline.
 Accelerator runtime execution remains unvalidated without its libraries/hardware.
 
+The final ownership fix reserves a control before creating its watcher or deadline.
+Regression tests cover a zero-deadline duplicate before native entry, wrong-model
+reuse, and 16 concurrent public reservations with exactly one successful owner.
+After the fix, another 100,000 bounded calls per runtime retained checked scores,
+bounded sampled memory, and flat process counts. Run-control resources are V2.
+
 ## Deferred scope
 
 The stack covers the selected CPU/API 17–24 modernization, session/environment

lib/ortex/model.ex

@@ -93,6 +93,12 @@ defmodule Ortex.Model do
             )
 
         opts = %{opts | control: control}
+
+        case Ortex.Native.reserve_run(model.reference, control.reference) do
+          :ok -> :ok
+          {:error, message} -> raise message
+        end
+
         Ortex.Runner.run(control, opts.timeout, fn -> run_validated(model, inputs, opts) end)
       else
         run_validated(model, inputs, opts)

lib/ortex/native.ex

@@ -40,6 +40,7 @@ defmodule Ortex.Native do
   def new_run_control(_model, _options), do: :erlang.nif_error(:nif_not_loaded)
   def cancel_run(_control), do: :erlang.nif_error(:nif_not_loaded)
   def run_started(_control), do: :erlang.nif_error(:nif_not_loaded)
+  def reserve_run(_model, _control), do: :erlang.nif_error(:nif_not_loaded)
   def inference_scheduler, do: :erlang.nif_error(:nif_not_loaded)
   def run_controlled(_model, _inputs, _control), do: :erlang.nif_error(:nif_not_loaded)
   def run_named_controlled(_model, _inputs, _control), do: :erlang.nif_error(:nif_not_loaded)

lib/ortex/run_control.ex

@@ -47,6 +47,6 @@ defmodule Ortex.RunControl do
     end
   end
 
-  @doc "Returns whether the single attempt has been claimed, including waiting for the session mutex."
+  @doc "Returns whether the native attempt has been entered, including waiting for the session mutex."
   def started?(%__MODULE__{reference: reference}), do: Ortex.Native.run_started(reference)
 end

native/ortex/src/lib.rs

@@ -278,6 +278,17 @@ fn run_started(control: ResourceArc<run_control::RunControl>) -> bool {
     control.started()
 }
 
+#[rustler::nif]
+fn reserve_run(
+    model: ResourceArc<model::OrtexModel>,
+    control: ResourceArc<run_control::RunControl>,
+) -> NifResult<Atom> {
+    control
+        .reserve(&model)
+        .map_err(|e| rustler::Error::Term(Box::new(e.to_string())))?;
+    Ok(rustler::types::atom::ok())
+}
+
 #[cfg_attr(feature = "benchmark-dirty-cpu", rustler::nif(schedule = "DirtyCpu"))]
 #[cfg_attr(
     not(feature = "benchmark-dirty-cpu"),

native/ortex/src/run_control.rs

@@ -3,7 +3,7 @@ use crate::model::{self, InferenceOptions, OrtexModel};
 use ort::session::{HasSelectedOutputs, OutputSelector, RunOptions};
 use ort::Error;
 use rustler::{Resource, ResourceArc};
-use std::sync::atomic::{AtomicBool, Ordering};
+use std::sync::atomic::{AtomicU8, Ordering};
 
 pub struct PreparedRun {
     names: Vec<String>,
@@ -58,11 +58,11 @@ impl PreparedRun {
 pub struct RunControl {
     model: ResourceArc<OrtexModel>,
     prepared: PreparedRun,
-    claimed: AtomicBool,
+    state: AtomicU8,
 }
 // Encoding only copies the opaque resource handle; it cannot mutate run options.
 impl std::panic::RefUnwindSafe for RunControl {}
-#[rustler::resource_impl(name = "OrtexRunControlV1")]
+#[rustler::resource_impl(name = "OrtexRunControlV2")]
 impl Resource for RunControl {}
 
 impl RunControl {
@@ -71,15 +71,36 @@ impl RunControl {
         Ok(Self {
             model,
             prepared,
-            claimed: AtomicBool::new(false),
+            state: AtomicU8::new(0),
         })
     }
-    pub fn claim(&self, model: &OrtexModel) -> Result<(), Error> {
+    fn check_model(&self, model: &OrtexModel) -> Result<(), Error> {
         if !std::ptr::eq(&*self.model, model) {
             return Err(Error::new("run control belongs to another model"));
         }
-        self.claimed
-            .compare_exchange(false, true, Ordering::AcqRel, Ordering::Acquire)
+        Ok(())
+    }
+    // Reserve ownership before spawning a worker or installing its deadline.
+    // A rejected duplicate must never acquire authority to cancel another call.
+    pub fn reserve(&self, model: &OrtexModel) -> Result<(), Error> {
+        self.check_model(model)?;
+        self.state
+            .compare_exchange(0, 1, Ordering::AcqRel, Ordering::Acquire)
+            .map_err(|_| Error::new("run control has already been used"))?;
+        Ok(())
+    }
+    pub fn claim(&self, model: &OrtexModel) -> Result<(), Error> {
+        self.check_model(model)?;
+        // Direct internal NIF calls can enter from ready (0); public managed
+        // calls enter from reserved (1). Neither can execute twice (2).
+        self.state
+            .fetch_update(Ordering::AcqRel, Ordering::Acquire, |state| {
+                if state < 2 {
+                    Some(2)
+                } else {
+                    None
+                }
+            })
             .map_err(|_| Error::new("run control has already been used"))?;
         Ok(())
     }
@@ -90,6 +111,6 @@ impl RunControl {
         self.prepared.options.terminate()
     }
     pub fn started(&self) -> bool {
-        self.claimed.load(Ordering::Acquire)
+        self.state.load(Ordering::Acquire) == 2
     }
 }

test/managed_runs_test.exs

@@ -135,7 +135,10 @@ defmodule Ortex.ManagedRunsTest do
     end
 
     refute RunControl.started?(control)
-    assert_raise RuntimeError, ~r/terminate/, fn -> Ortex.run(model, input, control: control) end
+
+    assert_raise RuntimeError, ~r/already been used/, fn ->
+      Ortex.run(model, input, control: control)
+    end
 
     for timeout <- [-1, 4_294_967_296, 1.5, nil, :later] do
       assert_raise ArgumentError, fn -> Ortex.run(model, input, timeout: timeout) end
@@ -145,6 +148,66 @@ defmodule Ortex.ManagedRunsTest do
     assert {_, _} = Ortex.run(model, input, timeout: :infinity)
   end
 
+  test "a duplicate deadline cannot cancel a reserved attempt before its native entry" do
+    model = load("named_io")
+    control = RunControl.new(model, outputs: ["sum"])
+    :ok = Ortex.Native.reserve_run(model.reference, control.reference)
+    refute RunControl.started?(control)
+    input = {Nx.tensor([7.0, 5.0]), Nx.tensor([2.0, 1.0])}
+
+    assert_raise RuntimeError, ~r/already been used/, fn ->
+      Ortex.run(model, input, control: control, timeout: 0)
+    end
+
+    refs =
+      input
+      |> Tuple.to_list()
+      |> Enum.map(fn tensor ->
+        %Nx.Tensor{data: %Ortex.Backend{ref: ref}} = Nx.backend_transfer(tensor, Ortex.Backend)
+        ref
+      end)
+
+    [{:tensor, {ref, [2], :f, 32}}] =
+      Ortex.Native.run_controlled(model.reference, refs, control.reference)
+
+    assert Ortex.Native.to_binary(ref, 32, 2) |> Nx.from_binary(:f32) |> Nx.to_flat_list() == [
+             9.0,
+             6.0
+           ]
+  end
+
+  test "concurrent public reservations permit one successful owner" do
+    model = load("named_io")
+    control = RunControl.new(model, outputs: ["sum"])
+
+    results =
+      1..16
+      |> Task.async_stream(
+        fn _ ->
+          try do
+            {sum} =
+              Ortex.run(model, {Nx.tensor([7.0, 5.0]), Nx.tensor([2.0, 1.0])},
+                control: control,
+                timeout: 5000
+              )
+
+            {:success, Nx.to_flat_list(sum)}
+          rescue
+            error in RuntimeError -> {:rejected, error.message}
+          end
+        end,
+        max_concurrency: 16
+      )
+      |> Enum.to_list()
+
+    assert Enum.count(results, &(&1 == {:ok, {:success, [9.0, 6.0]}})) == 1
+
+    assert Enum.count(results, fn
+             {:ok, {:rejected, message}} -> message =~ "already been used"
+             _ -> false
+           end) == 15
+  end
+
   test "timed-out protocol workers terminate and leave no late reply or monitor messages" do
     model = load("named_io")
     caller = self()

test/run_control_test.exs

@@ -69,7 +69,7 @@ defmodule Ortex.RunControlTest do
     control = RunControl.new(model)
 
     assert_raise RuntimeError, ~r/another model/, fn ->
-      Ortex.run(other, inputs(), control: control)
+      Ortex.run(other, inputs(), control: control, timeout: 0)
     end
 
     assert_raise ArgumentError, fn ->