gigatoken 0.2.2 → 0.4.0

data/ext/gigatoken/src/gvl.rs

@@ -17,7 +17,7 @@
 //! (`~/architect/src/github.com/socketry/io-event/ext/io/event/worker_pool.c:309-316`).
 //! Without such a scheduler (or with one lacking a worker pool), `rb_nogvl`
 //! degrades to exactly today's behavior: release the GVL, block this thread.
-//! See `docs/rb/async-design.md` and `docs/rb/async.md` for the full design
+//! See `docs/explanation/async-design.md` and `docs/how-to/run-under-async.md` for the full design
 //! and gotchas (worker pool is opt-in, defaults to one background worker).
 //!
 //! `func`/`data1` may now run on a different OS thread than the caller (the
@@ -36,15 +36,27 @@
 //! `VALUE` or thread-local state, so running one OS thread over instead of
 //! another changes nothing about its safety, and the types involved satisfy
 //! `Send`/`Sync` on their own merits.
+//!
+//! Interrupts. `rb_nogvl` ends by checking interrupts and raises any that
+//! are pending — a `Timeout`, `Thread#kill`, `Interrupt`, an `Async`
+//! timeout — by longjmping out of itself, over every Rust frame in between.
+//! Both entry points here therefore call it under `magnus::rb_sys::protect`
+//! (see [`run`]), so the raise comes back as an ordinary `Error` and every
+//! guard, `InputDocs` and Ruby String lock the caller holds is released by
+//! ordinary Rust unwinding on the way out. [`without_gvl_cancellable`] goes
+//! one further and supplies a real unblock function, so a pending interrupt
+//! cancels the work in flight instead of waiting for it.
 
 use std::any::Any;
 use std::ffi::c_void;
 use std::os::raw::c_int;
 use std::panic::{self, AssertUnwindSafe};
+use std::sync::atomic::{AtomicBool, Ordering};
 
+use magnus::Error;
 use rb_sys::rb_nogvl;
 
-/// `RB_NOGVL_OFFLOAD_SAFE` (`ruby/thread.h:84` in a current Ruby checkout;
+/// `RB_NOGVL_OFFLOAD_SAFE` (`ruby/thread.h:73` in Ruby 4.0.7's own headers;
 /// introduced by Ruby's `Fiber::Scheduler#blocking_operation_wait` support,
 /// first released in Ruby 3.4.0). Defined locally rather than taken from
 /// `rb_sys::` bindings: `rb-sys` bindgens its constants from the *building*
@@ -61,62 +73,184 @@ const RB_NOGVL_OFFLOAD_SAFE: c_int = 0x4;
 /// value, or a caught panic payload to re-raise once we're back on ordinary
 /// (non-`extern "C"`) Rust stack frames. Unwinding a panic directly across
 /// the `extern "C"` trampoline `rb_nogvl` calls into is undefined behavior;
-/// catching it here and resuming it from `without_gvl` below turns that into
-/// an ordinary Rust panic, which magnus's own `method!`/`function!` call
-/// trampolines already wrap in `catch_unwind` and convert into a fatal Ruby
-/// exception (`magnus::error::Error::from_panic`) — the same outcome any
-/// other panicking native method already gets, just carried safely across
-/// the extra C boundary this one call adds.
+/// catching it here and resuming it from [`Slot::take`] below turns that
+/// into an ordinary Rust panic, which magnus's own `method!`/`function!`
+/// call trampolines already wrap in `catch_unwind` and convert into a fatal
+/// Ruby exception (`magnus::error::Error::from_panic`) — the same outcome
+/// any other panicking native method already gets, just carried safely
+/// across the extra C boundary this one call adds.
 enum Outcome<R> {
     Value(R),
     Panic(Box<dyn Any + Send + 'static>),
 }
 
+/// The one thing `rb_nogvl`'s callback and [`run`] share: the callback takes
+/// the closure out of `input` and leaves its result in `output`. The rule is
+/// that nothing travels through `rb_nogvl`'s *return value* — the caller owns
+/// this slot, in an ordinary Rust frame.
+///
+/// That rule is what makes the raise `rb_nogvl` performs on its way out
+/// leak-free. The `protect` closure in [`run`] holds the single FFI call and
+/// owns nothing at all, so the longjmp skips no destructor that matters:
+/// whatever is in flight is either here — in a frame that unwinds normally
+/// afterwards — or in one of the caller's own frames, which unwind with it.
+/// It is also the only way to tell "the callback ran" from "it never did",
+/// which a fiber scheduler cancelling the offloaded operation before its
+/// worker pool picks it up makes a real case (io-event's
+/// `worker_pool_work_wait`, `ext/io/event/worker_pool.c`).
+struct Slot<F, R> {
+    input: Option<F>,
+    output: Option<Outcome<R>>,
+}
+
+impl<F, R> Slot<F, R> {
+    fn new(f: F) -> Self {
+        Self {
+            input: Some(f),
+            output: None,
+        }
+    }
+
+    /// The callback's value, or `None` if it never ran. A panic it caught
+    /// resumes here, on an ordinary Rust frame.
+    fn take(&mut self) -> Option<R> {
+        match self.output.take() {
+            Some(Outcome::Value(value)) => Some(value),
+            Some(Outcome::Panic(payload)) => panic::resume_unwind(payload),
+            None => None,
+        }
+    }
+}
+
 unsafe extern "C" fn call_without_gvl<F, R>(arg: *mut c_void) -> *mut c_void
 where
     F: FnOnce() -> R + Send,
     R: Send,
 {
-    // SAFETY: `arg` is the `*mut Option<F>` handed to `rb_nogvl` below, valid
-    // for the duration of that (synchronous) call, and this is the only
-    // place it's dereferenced.
-    let closure = unsafe { (*(arg as *mut Option<F>)).take() }
+    // SAFETY: `arg` is the `*mut Slot<F, R>` handed to `rb_nogvl` by `run`,
+    // valid for the duration of that (synchronous) call, and this is the
+    // only place it's dereferenced.
+    let slot = unsafe { &mut *(arg as *mut Slot<F, R>) };
+    let closure = slot
+        .input
+        .take()
         .expect("without_gvl callback invoked more than once");
-    let outcome = match panic::catch_unwind(AssertUnwindSafe(closure)) {
+    slot.output = Some(match panic::catch_unwind(AssertUnwindSafe(closure)) {
         Ok(value) => Outcome::Value(value),
         Err(payload) => Outcome::Panic(payload),
-    };
-    Box::into_raw(Box::new(outcome)) as *mut c_void
+    });
+    std::ptr::null_mut()
+}
+
+/// Run `slot`'s closure with the GVL released, under `rb_protect`.
+///
+/// `rb_nogvl` finishes by checking interrupts and raising any that are
+/// pending, by longjmping out of itself. `protect` catches that and returns
+/// it as an `Error` — including the `Tag::Fatal` a `Thread#kill` jumps with,
+/// which magnus resumes with `rb_jump_tag` once the caller's frames have
+/// unwound. The closure below is deliberately trivial: one FFI call and a
+/// `nil`, owning nothing the longjmp could strand.
+fn run<F, R>(
+    slot: &mut Slot<F, R>,
+    ubf: rb_sys::rb_unblock_function_t,
+    data2: *mut c_void,
+) -> Result<(), Error>
+where
+    F: FnOnce() -> R + Send,
+    R: Send,
+{
+    let arg = slot as *mut Slot<F, R> as *mut c_void;
+    let nil: rb_sys::VALUE = rb_sys::Qnil.into();
+    magnus::rb_sys::protect(|| {
+        // SAFETY: `arg` points at `slot`, which outlives this synchronous
+        // call; the callback is the only reader of it (see `Slot`).
+        unsafe {
+            rb_nogvl(
+                Some(call_without_gvl::<F, R>),
+                arg,
+                ubf,
+                data2,
+                RB_NOGVL_OFFLOAD_SAFE,
+            )
+        };
+        nil
+    })?;
+    Ok(())
+}
+
+/// `rb_nogvl`'s unblock function: Ruby calls it from another thread when an
+/// interrupt is pending for this one, and a fiber scheduler calls it to
+/// cancel an offloaded operation (`rb_fiber_scheduler_blocking_operation_cancel`
+/// "marks it as cancelled and calls the unblock function", Ruby 4.0.7's
+/// `ruby/fiber/scheduler.h:455-457` — which is how an `Async` timeout reaches
+/// an encode running on io-event's worker pool). It does exactly one thing:
+/// set the flag the core's encode loops poll (`src/batch.rs`), so the encode
+/// stops at the next document boundary instead of running to completion.
+unsafe extern "C" fn set_cancel(arg: *mut c_void) {
+    // SAFETY: `arg` is the `&AtomicBool` `without_gvl_cancellable` passed as
+    // `data2`, living in its frame for the whole `rb_nogvl` call — the only
+    // window in which Ruby may call this.
+    unsafe { &*(arg as *const AtomicBool) }.store(true, Ordering::Relaxed);
 }
 
 /// Run `f` with the GVL released: other Ruby threads may run while `f`
 /// executes, and — under a fiber scheduler with a worker pool — the calling
 /// fiber yields to the reactor while `f` runs on a background thread. `f`
 /// must not touch any Ruby object (`VALUE`) — only plain Rust data — per the
-/// Ruby C API's contract for this call. A panic inside `f` is caught and
-/// re-raised here rather than left to unwind across the C trampoline.
-pub fn without_gvl<F, R>(f: F) -> R
+/// Ruby C API's contract for this call. An interrupt that arrives meanwhile
+/// is delivered once `f` has finished, as an `Err`; use
+/// [`without_gvl_cancellable`] for work that can stop early.
+pub fn without_gvl<F, R>(f: F) -> Result<R, Error>
 where
     F: FnOnce() -> R + Send,
     R: Send,
 {
-    let mut slot = Some(f);
-    let arg = &mut slot as *mut Option<F> as *mut c_void;
-    let result = unsafe {
-        rb_nogvl(
-            Some(call_without_gvl::<F, R>),
-            arg,
-            None,
-            std::ptr::null_mut(),
-            RB_NOGVL_OFFLOAD_SAFE,
-        )
-    };
-    // SAFETY: `result` is the `Box::into_raw(Box::new(Outcome<R>))` pointer
-    // produced by the callback above, which always runs exactly once before
-    // `rb_nogvl` returns.
-    let outcome = *unsafe { Box::from_raw(result as *mut Outcome<R>) };
-    match outcome {
-        Outcome::Value(value) => value,
-        Outcome::Panic(payload) => panic::resume_unwind(payload),
+    let mut slot = Slot::new(f);
+    run(&mut slot, None, std::ptr::null_mut())?;
+    Ok(match slot.take() {
+        Some(value) => value,
+        // A fiber scheduler cancelled the offloaded operation before its
+        // worker pool ever started it, and nothing was raised (that would
+        // have come back as an `Err` above). The closure is still in the
+        // slot and still ours to run: do that here, GVL and all, rather
+        // than invent a result.
+        None => slot.input.take().expect("the callback left the closure")(),
+    })
+}
+
+/// [`without_gvl`] for work that can stop early: `attempt` builds the
+/// closure, which gets a cancellation token to poll and returns `None` if it
+/// saw the token set and cut its work short. An interrupt arriving during
+/// the run sets that token through [`set_cancel`], so the caller sees the
+/// interrupt instead of waiting out the whole batch.
+///
+/// `attempt` is a factory, not the closure itself, because a cancelled run
+/// sometimes has to be redone: Ruby calls the unblock function for every
+/// interrupt, including ones that never raise (a trap handler, a
+/// `Thread#wakeup`), and returning the half-encoded batch those produce
+/// would be a silently truncated result. The redo runs uninterruptibly —
+/// plain [`without_gvl`], which is exactly how this path behaved before
+/// cancellation existed — so a chatty signal handler can cost a batch one
+/// extra pass, never an unbounded number of them.
+pub fn without_gvl_cancellable<A, F, R>(mut attempt: A) -> Result<R, Error>
+where
+    A: FnMut() -> F,
+    F: FnOnce(&AtomicBool) -> Option<R> + Send,
+    R: Send,
+{
+    let cancel = AtomicBool::new(false);
+    let f = attempt();
+    let mut slot = Slot::new(|| f(&cancel));
+    run(
+        &mut slot,
+        Some(set_cancel),
+        &cancel as *const AtomicBool as *mut c_void,
+    )?;
+    if let Some(Some(value)) = slot.take() {
+        return Ok(value);
     }
+    let f = attempt();
+    let never_cancelled = AtomicBool::new(false);
+    without_gvl(|| f(&never_cancelled))
+        .map(|value| value.expect("nothing sets the token of an uncancellable run"))
 }

data/ext/gigatoken/src/lib.rs

@@ -1,3 +1,4 @@
+use gigatoken_rs::input::file_source::load_file;
 use gigatoken_rs::load_tokenizer::hf::{self, HfTokenizer};
 use gigatoken_rs::pretokenize::PretokenizerType;
 use magnus::{Error, Module, RString, Ruby, Value, function};
@@ -26,9 +27,9 @@ mod sentencepiece;
 mod sources;
 mod tokenizer;
 
-use error::raise;
+use error::{model_error, raise};
 use sentencepiece::SentencePieceTokenizer;
-use tokenizer::BPETokenizer;
+use tokenizer::{binary_string, BPETokenizer};
 
 // The gigatoken core crate exposes no version constant of its own, so this
 // is the ext crate's (gigatoken-rb's) version — see the builder report.
@@ -66,7 +67,20 @@ fn load_hf_json(ruby: &Ruby, data: RString) -> Result<Value, Error> {
     match hf::load_hf_slice(bytes) {
         Ok(HfTokenizer::Bpe(tokenizer)) => Ok(ruby.into_value(BPETokenizer::from_tokenizer(tokenizer))),
         Ok(HfTokenizer::SentencePiece(tokenizer)) => Ok(ruby.into_value(SentencePieceTokenizer::from_tokenizer(tokenizer))),
-        Err(e) => Err(raise(ruby, e.to_string())),
+        Err(e) => Err(model_error(ruby, e.to_string())),
+    }
+}
+
+/// One file's contents as a binary String, decompressed by extension the way
+/// the native file sources do it (`.gz`, `.zst`/`.zstd`, plain — see the core's
+/// `load_file`). The CLI's Ruby-side split reads through here so both sides of
+/// `gigatoken validate` see the same bytes with one decoder between them
+/// (`lib/gigatoken/cli/support.rb`); the encode paths never touch it, which is
+/// why it keeps the GVL.
+fn read_input(ruby: &Ruby, path: String) -> Result<RString, Error> {
+    match load_file(std::path::Path::new(&path)) {
+        Ok(file) => Ok(binary_string(ruby, file.as_bytes())),
+        Err(e) => Err(raise(ruby, format!("{path}: {e}"))),
     }
 }
 
@@ -77,6 +91,7 @@ fn init(ruby: &Ruby) -> Result<(), Error> {
     native.define_module_function("crate_version", function!(crate_version, 0))?;
     native.define_module_function("load_hf_json", function!(load_hf_json, 1))?;
     native.define_module_function("pretokenizer_names", function!(pretokenizer_names, 0))?;
+    native.define_module_function("read_input", function!(read_input, 1))?;
     native.define_module_function("set_max_cache_bytes", function!(set_max_cache_bytes, 1))?;
     native.define_module_function("get_max_cache_bytes", function!(get_max_cache_bytes, 0))?;
     sources::init(ruby, native)?;

data/ext/gigatoken/src/sentencepiece.rs

@@ -22,14 +22,16 @@ use magnus::{
     scan_args::{get_kwargs, scan_args},
 };
 
-use crate::error::raise;
+use crate::error::{input_error, raise};
 use crate::gvl::without_gvl;
 use crate::sources;
-use crate::tokenizer::{binary_string, packed_result, ragged_result};
+use crate::tokenizer::{
+    binary_string, packed_result, ragged_result, require_known_ids, snapshot_entry, snapshot_inputs,
+};
 
-/// Validate that `bytes` is UTF-8, raising `Gigatoken::Error` otherwise.
+/// Validate that `bytes` is UTF-8, raising `Gigatoken::InputError` otherwise.
 fn require_utf8<'a>(ruby: &Ruby, bytes: &'a [u8]) -> Result<&'a str, Error> {
-    std::str::from_utf8(bytes).map_err(|e| raise(ruby, format!("invalid UTF-8: {e}")))
+    std::str::from_utf8(bytes).map_err(|e| input_error(ruby, format!("invalid UTF-8: {e}")))
 }
 
 #[magnus::wrap(class = "Gigatoken::Native::SentencePieceTokenizer", free_immediately, size)]
@@ -40,7 +42,10 @@ pub struct SentencePieceTokenizer {
     tokenizer: SentencePieceBPE,
     // The one mutable piece. A `Mutex` rather than a `RefCell` so the wrapped
     // object is `Sync`: Ruby hands the same instance to every thread, and a
-    // `RefCell` shared that way is unsound (see `BPETokenizer`'s lock).
+    // `RefCell` shared that way is unsound. It is only ever taken with the
+    // GVL held — the batch paths build their own per-call encoders — so it
+    // cannot actually block, which is what `BPETokenizer`'s invariant asks
+    // of any lock on this side of the boundary.
     state: Mutex<EncodeState>,
 }
 
@@ -74,21 +79,27 @@ impl SentencePieceTokenizer {
     /// string is validated as UTF-8 and copied into an owned `String`
     /// before release: nothing Ruby-managed may be touched once the GVL is
     /// gone.
+    ///
+    /// The walk reads a [`snapshot_inputs`] copy of the caller's array, one
+    /// slot at a time — never a Rust borrow of the array's buffer held
+    /// across `try_convert`, which runs `to_str` and with it arbitrary Ruby
+    /// code, including code that re-homes the array out from under the
+    /// walk. Same contract as the BPE path: the result reflects the array
+    /// as it was at entry.
     fn encode_batch_ragged(ruby: &Ruby, rb_self: &Self, inputs: RArray) -> Result<(Vec<u32>, Vec<i64>), Error> {
-        // SAFETY: values are read (checked-converted to `RString`, then
-        // validated and copied into owned buffers below) before anything
-        // else runs.
-        let docs: Vec<String> = unsafe { inputs.as_slice() }
-            .iter()
-            .map(|&v| {
-                let s = RString::try_convert(v)?;
+        let snapshot = snapshot_inputs(inputs);
+        let docs: Vec<String> = (0..snapshot.len())
+            .map(|i| {
+                let s = RString::try_convert(snapshot_entry(snapshot, i))?;
+                // SAFETY: validated and copied into an owned buffer right
+                // here, with no Ruby call in between to invalidate it.
                 let bytes = unsafe { s.as_slice() };
                 require_utf8(ruby, bytes).map(str::to_owned)
             })
             .collect::<Result<_, _>>()?;
         let doc_refs: Vec<&str> = docs.iter().map(String::as_str).collect();
         let tokenizer: &SentencePieceBPE = &rb_self.tokenizer;
-        Ok(without_gvl(|| sp_encode_docs_ragged(tokenizer, &doc_refs)))
+        without_gvl(|| sp_encode_docs_ragged(tokenizer, &doc_refs))
     }
 
     fn encode_batch(ruby: &Ruby, rb_self: &Self, inputs: RArray) -> Result<RArray, Error> {
@@ -133,7 +144,7 @@ impl SentencePieceTokenizer {
         }
 
         let tokenizer: &SentencePieceBPE = &rb_self.tokenizer;
-        let encoded: std::io::Result<(Vec<u32>, Vec<i64>)> = without_gvl(|| {
+        let encoded = without_gvl(|| {
             sources::encode_files_ragged(&source, parallel, |files, format| {
                 for &region in files {
                     std::str::from_utf8(region).map_err(|e| {
@@ -146,7 +157,7 @@ impl SentencePieceTokenizer {
                     sp_encode_files_docs_serial(tokenizer, files, format)
                 })
             })
-        });
+        })?;
         encoded.map_err(|e| raise(ruby, e.to_string()))
     }
 
@@ -166,6 +177,7 @@ impl SentencePieceTokenizer {
 
     fn decode(ruby: &Ruby, rb_self: &Self, tokens: RArray) -> Result<RString, Error> {
         let ids: Vec<u32> = tokens.to_vec()?;
+        require_known_ids(ruby, &ids, rb_self.tokenizer.vocab_size())?;
         let ids: Vec<_> = ids.into_iter().map(Into::into).collect();
         let bytes = rb_self.tokenizer.decode(&ids);
         Ok(binary_string(ruby, &bytes))