Impl streaming

Author: owjs3901

AGENTS.md

@@ -136,6 +136,7 @@ Feature flags:
 |---------|-----------|------|
 | `inprocess` | `vespera::inprocess` (= `vespera_inprocess`) | dispatch, register_app, envelopes |
 | `jni` | `vespera::jni` (= `vespera_jni`) + implies `inprocess` | RUNTIME, jni_app!, JNI symbol |
+| `mimalloc` | (with `jni`) mimalloc as the cdylib's `#[global_allocator]` | measured -15~19% on sync/direct dispatch vs Windows HeapAlloc |
 
 ## JNI ARCHITECTURE
 

Cargo.lock

@@ -22,6 +22,9 @@ default = [
 cron = ["dep:tokio-cron-scheduler"]
 inprocess = ["dep:vespera_inprocess"]
 jni = ["inprocess", "dep:vespera_jni"]
+# mimalloc as the cdylib's global allocator (see vespera_jni docs).
+# Weak dep syntax: only applies when the `jni` feature enables vespera_jni.
+mimalloc = ["vespera_jni?/mimalloc"]
 # Runtime validation: `#[derive(Schema)]` additionally emits
 # `impl garde::Validate` and the `Validated<T>` extractor is enabled.
 # The `garde` crate is bundled internally and never named by user code.

crates/vespera/Cargo.toml

@@ -10,6 +10,17 @@ repository.workspace = true
 vespera_inprocess = { workspace = true }
 jni = "0.22"
 tokio = { version = "1", features = ["rt-multi-thread"] }
+# Optional high-performance global allocator for the final cdylib.
+# Opt-in because #[global_allocator] is process-wide and must be the
+# embedding crate's decision.
+mimalloc = { version = "0.1", optional = true }
+
+[features]
+# Use mimalloc as the global allocator inside the JNI cdylib.  The
+# default OS allocator (Windows HeapAlloc in particular) is measurably
+# slower on the allocation-heavy dispatch paths (input Vec, response
+# collect, wire response, streaming chunks).
+mimalloc = ["dep:mimalloc"]
 
 [lints]
 workspace = true

crates/vespera_jni/Cargo.toml

@@ -1,5 +1,5 @@
 use std::{
-    cell::Cell,
+    cell::{Cell, RefCell},
     ffi::c_void,
     panic::{AssertUnwindSafe, catch_unwind, resume_unwind},
     ptr,
@@ -39,6 +39,104 @@ static RUNTIME_WORKER_THREADS: std::sync::OnceLock<Option<usize>> = std::sync::O
 
 thread_local! {
     static ASYNC_DAEMON_ENV: Cell<*mut jni::sys::JNIEnv> = const { Cell::new(ptr::null_mut()) };
+    static STREAMING_PULL_BUFFER: RefCell<Option<CachedStreamingChunkBuffer>> = const { RefCell::new(None) };
+    static STREAMING_PUSH_BUFFER: RefCell<Option<CachedStreamingChunkBuffer>> = const { RefCell::new(None) };
+}
+
+type StreamingChunkBuffer = Global<JByteArray<'static>>;
+
+#[derive(Clone, Copy)]
+enum StreamingBufferRole {
+    Pull,
+    Push,
+}
+
+impl StreamingBufferRole {
+    fn with_cache<R>(
+        self,
+        callback: impl FnOnce(&RefCell<Option<CachedStreamingChunkBuffer>>) -> R,
+    ) -> R {
+        match self {
+            Self::Pull => STREAMING_PULL_BUFFER.with(callback),
+            Self::Push => STREAMING_PUSH_BUFFER.with(callback),
+        }
+    }
+}
+
+struct CachedStreamingChunkBuffer {
+    size: usize,
+    array: StreamingChunkBuffer,
+    checked_out: bool,
+}
+
+// Released explicitly only after the streaming future returns normally.  If a
+// panic unwinds through a bidirectional dispatch while the request producer may
+// still be in `InputStream.read`, the cache stays checked out and future
+// dispatches allocate fresh buffers instead of aliasing the Java array.
+struct StreamingChunkBufferLease {
+    role: StreamingBufferRole,
+}
+
+impl StreamingChunkBufferLease {
+    const fn new(role: StreamingBufferRole) -> Self {
+        Self { role }
+    }
+
+    fn mark_reusable(self) {
+        self.role.with_cache(|cache| {
+            if let Some(cached) = cache.borrow_mut().as_mut() {
+                cached.checked_out = false;
+            }
+        });
+    }
+}
+
+fn new_streaming_chunk_buffer(
+    env: &mut jni::Env<'_>,
+    size: usize,
+) -> jni::errors::Result<StreamingChunkBuffer> {
+    let local = env.new_byte_array(size)?;
+    env.new_global_ref(&local)
+}
+
+fn checkout_streaming_chunk_buffer(
+    env: &mut jni::Env<'_>,
+    role: StreamingBufferRole,
+) -> jni::errors::Result<(StreamingChunkBuffer, Option<StreamingChunkBufferLease>)> {
+    let size = streaming_chunk_size();
+    role.with_cache(|cache| {
+        let mut slot = cache.borrow_mut();
+        let replace_cached = slot
+            .as_ref()
+            .is_none_or(|cached| cached.size != size && !cached.checked_out);
+
+        if replace_cached {
+            *slot = Some(CachedStreamingChunkBuffer {
+                size,
+                array: new_streaming_chunk_buffer(env, size)?,
+                checked_out: false,
+            });
+        }
+
+        let Some(cached) = slot.as_mut() else {
+            return Ok((new_streaming_chunk_buffer(env, size)?, None));
+        };
+
+        if cached.size != size || cached.checked_out {
+            return Ok((new_streaming_chunk_buffer(env, size)?, None));
+        }
+
+        let cached_array: &JByteArray<'static> = cached.array.as_ref();
+        let dispatch_array = env.new_global_ref(cached_array)?;
+        cached.checked_out = true;
+        Ok((dispatch_array, Some(StreamingChunkBufferLease::new(role))))
+    })
+}
+
+fn mark_streaming_buffer_reusable(lease: Option<StreamingChunkBufferLease>) {
+    if let Some(lease) = lease {
+        lease.mark_reusable();
+    }
 }
 
 fn attach_async_daemon_thread(jvm: &jni::JavaVM) -> jni::errors::Result<*mut jni::sys::JNIEnv> {
@@ -509,18 +607,23 @@ pub extern "system" fn Java_com_devfive_vespera_bridge_VesperaBridge_dispatchStr
             let stream_global: Global<JObject<'static>> = env.new_global_ref(&output_stream)?;
             let jvm = env.get_java_vm()?;
 
-            // One reusable Java chunk buffer for the whole stream.
-            let push_buf_local = env.new_byte_array(streaming_chunk_size())?;
-            let push_buf: Global<jni::objects::JByteArray<'static>> =
-                env.new_global_ref(&push_buf_local)?;
+            // One per-thread reusable Java chunk buffer for the whole stream.
+            let (push_buf, push_buf_lease) =
+                checkout_streaming_chunk_buffer(env, StreamingBufferRole::Push)?;
 
             let header_bytes = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
                 RUNTIME.block_on(vespera_inprocess::dispatch_streaming_async(
                     input,
                     make_push_closure(jvm, stream_global, push_buf),
                 ))
-            }))
-            .unwrap_or_else(|_| vespera_inprocess::error_wire(500, "panic in Rust engine"));
+            }));
+            let header_bytes = header_bytes.map_or_else(
+                |_| vespera_inprocess::error_wire(500, "panic in Rust engine"),
+                |header_bytes| {
+                    mark_streaming_buffer_reusable(push_buf_lease);
+                    header_bytes
+                },
+            );
 
             Ok(env.byte_array_from_slice(&header_bytes)?.into())
         })
@@ -576,15 +679,18 @@ pub extern "system" fn Java_com_devfive_vespera_bridge_VesperaBridge_dispatchFul
             let output_global: Global<JObject<'static>> = env.new_global_ref(&output_stream)?;
             let jvm = env.get_java_vm()?;
 
-            let chunk_size = streaming_chunk_size();
             // Pull and push run concurrently on different threads, so each
-            // direction owns its own global-ref'd buffer.
-            let pull_buf_local = env.new_byte_array(chunk_size)?;
-            let pull_buf: Global<jni::objects::JByteArray<'static>> =
-                env.new_global_ref(&pull_buf_local)?;
-            let push_buf_local = env.new_byte_array(chunk_size)?;
-            let push_buf: Global<jni::objects::JByteArray<'static>> =
-                env.new_global_ref(&push_buf_local)?;
+            // direction checks out its own per-thread cached buffer.
+            let (pull_buf, pull_buf_lease) =
+                checkout_streaming_chunk_buffer(env, StreamingBufferRole::Pull)?;
+            let (push_buf, push_buf_lease) =
+                match checkout_streaming_chunk_buffer(env, StreamingBufferRole::Push) {
+                    Ok(checked_out) => checked_out,
+                    Err(err) => {
+                        mark_streaming_buffer_reusable(pull_buf_lease);
+                        return Err(err);
+                    }
+                };
 
             // Closures capture clones of the JavaVM and Globals;
             // both types are Send+Sync.
@@ -604,8 +710,15 @@ pub extern "system" fn Java_com_devfive_vespera_bridge_VesperaBridge_dispatchFul
                     // Runs on the tokio worker driving the dispatch.
                     make_push_closure(push_jvm, push_global, push_buf),
                 ))
-            }))
-            .unwrap_or_else(|_| vespera_inprocess::error_wire(500, "panic in Rust engine"));
+            }));
+            let header_response = header_response.map_or_else(
+                |_| vespera_inprocess::error_wire(500, "panic in Rust engine"),
+                |header_response| {
+                    mark_streaming_buffer_reusable(pull_buf_lease);
+                    mark_streaming_buffer_reusable(push_buf_lease);
+                    header_response
+                },
+            );
 
             Ok(env.byte_array_from_slice(&header_response)?.into())
         })
@@ -649,10 +762,9 @@ pub extern "system" fn Java_com_devfive_vespera_bridge_VesperaBridge_dispatchStr
         let stream_global: Global<JObject<'static>> = env.new_global_ref(&output_stream)?;
         let jvm = env.get_java_vm()?;
 
-        // One reusable Java chunk buffer for the whole stream.
-        let push_buf_local = env.new_byte_array(streaming_chunk_size())?;
-        let push_buf: Global<jni::objects::JByteArray<'static>> =
-            env.new_global_ref(&push_buf_local)?;
+        // One per-thread reusable Java chunk buffer for the whole stream.
+        let (push_buf, push_buf_lease) =
+            checkout_streaming_chunk_buffer(env, StreamingBufferRole::Push)?;
 
         // Panic safety: catch_unwind absorbs Rust panics so the
         // JVM never sees an unwinding stack across the FFI
@@ -664,7 +776,7 @@ pub extern "system" fn Java_com_devfive_vespera_bridge_VesperaBridge_dispatchStr
         // rare — e.g. wire parse), the Java side will see a
         // dangling controller; document that follow-up callers
         // should set a timeout.
-        let _panic_result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
+        let panic_result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
             let header_for_cb = header_global;
             let jvm_for_cb = jvm.clone();
             let push = make_push_closure(jvm, stream_global, push_buf);
@@ -680,6 +792,9 @@ pub extern "system" fn Java_com_devfive_vespera_bridge_VesperaBridge_dispatchStr
                 push,
             ));
         }));
+        if panic_result.is_ok() {
+            mark_streaming_buffer_reusable(push_buf_lease);
+        }
 
         Ok(())
     });
@@ -718,14 +833,17 @@ pub extern "system" fn Java_com_devfive_vespera_bridge_VesperaBridge_dispatchFul
         let output_global: Global<JObject<'static>> = env.new_global_ref(&output_stream)?;
         let jvm = env.get_java_vm()?;
 
-        let chunk_size = streaming_chunk_size();
         // Pull and push run concurrently on different threads.
-        let pull_buf_local = env.new_byte_array(chunk_size)?;
-        let pull_buf: Global<jni::objects::JByteArray<'static>> =
-            env.new_global_ref(&pull_buf_local)?;
-        let push_buf_local = env.new_byte_array(chunk_size)?;
-        let push_buf: Global<jni::objects::JByteArray<'static>> =
-            env.new_global_ref(&push_buf_local)?;
+        let (pull_buf, pull_buf_lease) =
+            checkout_streaming_chunk_buffer(env, StreamingBufferRole::Pull)?;
+        let (push_buf, push_buf_lease) =
+            match checkout_streaming_chunk_buffer(env, StreamingBufferRole::Push) {
+                Ok(checked_out) => checked_out,
+                Err(err) => {
+                    mark_streaming_buffer_reusable(pull_buf_lease);
+                    return Err(err);
+                }
+            };
 
         let pull_jvm = jvm.clone();
         let pull_global = input_global;
@@ -737,7 +855,7 @@ pub extern "system" fn Java_com_devfive_vespera_bridge_VesperaBridge_dispatchFul
         // See dispatchStreamingWithHeader: panic absorbed silently,
         // recovery semantics depend on which side of the header
         // callback the panic landed.
-        let _panic_result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
+        let panic_result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
             RUNTIME.block_on(
                 vespera_inprocess::dispatch_bidirectional_streaming_with_header(
                     header_input,
@@ -753,6 +871,10 @@ pub extern "system" fn Java_com_devfive_vespera_bridge_VesperaBridge_dispatchFul
                 ),
             );
         }));
+        if panic_result.is_ok() {
+            mark_streaming_buffer_reusable(pull_buf_lease);
+            mark_streaming_buffer_reusable(push_buf_lease);
+        }
 
         Ok(())
     });

crates/vespera_jni/src/jni_impl.rs

@@ -17,6 +17,18 @@
 pub use jni;
 pub use vespera_inprocess;
 
+/// mimalloc as the process-wide allocator (feature `mimalloc`).
+///
+/// The JNI dispatch hot path allocates several times per call (input
+/// buffer, request body, response collection, wire response); the OS
+/// default allocator — Windows `HeapAlloc` in particular — is
+/// measurably slower than mimalloc on this pattern.  Opt-in because a
+/// `#[global_allocator]` is process-wide and belongs to the final
+/// cdylib's build decision.
+#[cfg(feature = "mimalloc")]
+#[global_allocator]
+static GLOBAL_ALLOC: mimalloc::MiMalloc = mimalloc::MiMalloc;
+
 /// Generate the `JNI_OnLoad` export that registers a single (default)
 /// app.  Backward-compatible sugar for the single-app case; new code
 /// targeting multiple apps should use [`jni_apps!`] directly.

crates/vespera_jni/src/lib.rs

@@ -12,7 +12,7 @@ name = "rust-jni-demo"
 path = "src/main.rs"
 
 [dependencies]
-vespera = { path = "../../crates/vespera", features = ["jni"] }
+vespera = { path = "../../crates/vespera", features = ["jni", "mimalloc"] }
 serde = { version = "1", features = ["derive"] }
 serde_json = "1"
 tokio = { version = "1", features = ["rt-multi-thread", "macros", "net"] }

examples/rust-jni-demo/Cargo.toml

@@ -202,3 +202,17 @@ Protocol: same 3 JVM invocations; run 1 discarded as cold; retained value is the
 | `async_completable_future` | 22,756 | 21,474 | **1,282 ns/op faster** (-5.6%) |
 
 The measured win is above the **100 ns/op** keep gate. Follow-up review found that the daemon-attached Tokio worker must explicitly clear pending Java exceptions after every completion callback because it no longer gets jni-rs scoped-detach cleanup. The implementation now clears pending exceptions after callback success, callback error, and callback unwind while preserving the callback return/error. A targeted regression guard, `AsyncDispatchExceptionHygieneTest.throwingFutureCompleteDoesNotPoisonNextAsyncCompletion`, first forces `CompletableFuture.complete()` to throw and then asserts a normal `dispatchAsync` still completes with status 200; it failed before the cleanup with a timeout and passes after the fix. A single post-fix sanity bench run measured `async_completable_future` at **16,107 ns/op** (informational only; not a replacement for the 3-JVM gate). Verification also passed `cargo clippy --workspace --all-targets -- -D warnings`, `cargo fmt --check`, `cargo test --workspace`, `cargo build -p rust-jni-demo --release`, and the full `:demo-app:test` Gradle suite (including `StreamingClosureStressTest` and the new hygiene guard).
+
+## Addendum (same day, later session): allocator + streaming buffer pooling
+
+Two further changes, paired same-session benches (GET /health, 100k iters, mimalloc build):
+
+| mode | default alloc | + mimalloc | + chunk-buffer pooling | total delta |
+|---|---:|---:|---:|---|
+| sync_dispatch_bytes | 2,870 | 2,314 | 2,322 | **-19%** |
+| direct_pooled | 2,376 | 2,017 | 2,000 | **-16%** |
+| response_streaming | 18,617* | 17,610 | **2,434** | **-87%** |
+| bidirectional_streaming | 37,543* | 32,326 | **2,605** | **-93%** |
+| async_completable_future | 22,038 | 19,468 | ~15,000 | **-32%** |
+
+\* with the 256 KiB chunk default: each streaming dispatch allocated+zeroed fresh 256 KiB Java arrays (bidi: two), costing ~10µs each — this addendum's TLS pooling (per-OS-thread cached Global<JByteArray>, fresh-alloc fallback when leased/reentrant) removes that per-dispatch cost entirely while keeping the 256 KiB throughput benefit for large transfers. mimalloc is opt-in via the vespera `mimalloc` cargo feature.