From b5cc28cb5d93532ce9067ffe9e133d89feacf03c Mon Sep 17 00:00:00 2001
From: zlexdev <zlexd.chqcode@gmail.com>
Date: Tue, 20 Oct 2026 00:48:33 +0300
Subject: [PATCH] Let callers set the ONNX Runtime threads of the sessions

Trustmark::new builds both sessions with 8 intra-op threads. A host
that loads one model per worker thread then runs workers x 2 x 8
threads on fewer cores. Trustmark::with_options takes SessionOptions
(intra-op threads, spinning); new() keeps its behaviour.
---
 rust/CHANGELOG.md |  5 ++++
 rust/src/lib.rs   | 68 +++++++++++++++++++++++++++++++++++++++++------
 2 files changed, 65 insertions(+), 8 deletions(-)

diff --git a/rust/CHANGELOG.md b/rust/CHANGELOG.md
index f84912d..7f8fffc 100644
--- a/rust/CHANGELOG.md
+++ b/rust/CHANGELOG.md
@@ -6,6 +6,11 @@ crate is versioned independently of the Python package in this repository.
 Entries for releases before 0.3.0 were reconstructed from the commit history,
 since those versions were published manually and without tags.
 
+## Unreleased
+
+- Add `Trustmark::with_options` and `SessionOptions` to set the intra-op thread count and
+  spinning of the ONNX Runtime sessions. `Trustmark::new` keeps its behaviour (8 threads).
+
 ## 0.4.0
 
 - Downgrade `ort` to 2.0.0-rc.10 to restore support for Intel macOS.
diff --git a/rust/src/lib.rs b/rust/src/lib.rs
index ce66dcb..63a5eb2 100644
--- a/rust/src/lib.rs
+++ b/rust/src/lib.rs
@@ -91,6 +91,29 @@ impl From<bits::Error> for Error {
 pub use bits::Version;
 pub use model::Variant;
 
+/// ONNX Runtime threading for the encoder and decoder sessions.
+///
+/// The default keeps the behaviour of [`Trustmark::new`]: 8 intra-op threads per session and the
+/// runtime's default spinning. A host that runs several models side by side (one per worker
+/// thread, for example) sets `intra_threads` to its own budget and usually turns spinning off,
+/// so idle sessions do not burn CPU.
+#[derive(Copy, Clone, Debug, PartialEq, Eq)]
+pub struct SessionOptions {
+    /// Intra-op threads of each session (the encoder and the decoder each get this many).
+    pub intra_threads: usize,
+    /// Whether intra-op threads spin while waiting for work; `None` keeps the runtime default.
+    pub intra_op_spinning: Option<bool>,
+}
+
+impl Default for SessionOptions {
+    fn default() -> Self {
+        Self {
+            intra_threads: 8,
+            intra_op_spinning: None,
+        }
+    }
+}
+
 impl Trustmark {
     /// Load a Trustmark model.
     pub fn new<P: AsRef<Path>>(
@@ -98,14 +121,27 @@ impl Trustmark {
         variant: Variant,
         version: Version,
     ) -> Result<Self, Error> {
-        let encoder = Session::builder()?
-            .with_optimization_level(GraphOptimizationLevel::Level3)?
-            .with_intra_threads(8)?
-            .commit_from_file(models.as_ref().join(variant.encoder_filename()))?;
-        let decoder = Session::builder()?
-            .with_optimization_level(GraphOptimizationLevel::Level3)?
-            .with_intra_threads(8)?
-            .commit_from_file(models.as_ref().join(variant.decoder_filename()))?;
+        Self::with_options(models, variant, version, SessionOptions::default())
+    }
+
+    /// Load a Trustmark model with explicit ONNX Runtime threading.
+    pub fn with_options<P: AsRef<Path>>(
+        models: P,
+        variant: Variant,
+        version: Version,
+        options: SessionOptions,
+    ) -> Result<Self, Error> {
+        let session = |file: String| -> Result<Session, Error> {
+            let mut builder = Session::builder()?
+                .with_optimization_level(GraphOptimizationLevel::Level3)?
+                .with_intra_threads(options.intra_threads)?;
+            if let Some(spinning) = options.intra_op_spinning {
+                builder = builder.with_intra_op_spinning(spinning)?;
+            }
+            Ok(builder.commit_from_file(models.as_ref().join(file))?)
+        };
+        let encoder = session(variant.encoder_filename())?;
+        let decoder = session(variant.decoder_filename())?;
         Ok(Self {
             encoder: Mutex::new(encoder),
             decoder: Mutex::new(decoder),
@@ -201,6 +237,22 @@ mod tests {
         Trustmark::new("./models", Variant::Q, Version::Bch5).unwrap();
     }
 
+    #[test]
+    fn roundtrip_with_one_thread_and_no_spinning() {
+        let options = SessionOptions {
+            intra_threads: 1,
+            intra_op_spinning: Some(false),
+        };
+        let tm = Trustmark::with_options("./models", Variant::Q, Version::Bch5, options).unwrap();
+        let input = image::open("../images/ghost.png").unwrap();
+        let watermark = "1011011110011000111111000000011111011111011100000110110110111".to_owned();
+        let encoded = tm.encode(watermark.clone(), input, 0.95).unwrap();
+        let decoded = tm
+            .decode(DynamicImage::ImageRgba8(encoded.to_rgba8()))
+            .unwrap();
+        assert_eq!(watermark, decoded);
+    }
+
     fn roundtrip(path: impl AsRef<Path>) {
         let tm = Trustmark::new("./models", Variant::Q, Version::Bch5).unwrap();
         let input = image::open(path.as_ref()).unwrap();