Fix async issue

Author: owjs3901

Cargo.lock

@@ -20,6 +20,10 @@ tokio = { version = "1", features = ["rt"] }
 [dev-dependencies]
 criterion = { version = "0.8", features = ["html_reports"] }
 tokio = { version = "1", features = ["rt-multi-thread", "macros"] }
+# `FutureExt::catch_unwind` for the `async_spawn_pattern` bench, which
+# A/Bs the vespera_jni `dispatchAsync` spawn-mechanism change (inner
+# `tokio::spawn` vs in-place `catch_unwind`).
+futures-util = { version = "0.3", default-features = false, features = ["std"] }
 
 [[bench]]
 name = "dispatch"

crates/vespera_inprocess/Cargo.toml

@@ -23,6 +23,7 @@
 
 use std::collections::HashMap;
 use std::ops::ControlFlow;
+use std::panic::AssertUnwindSafe;
 use std::sync::Mutex;
 
 use axum::{
@@ -32,6 +33,7 @@ use axum::{
     routing::{get, post},
 };
 use criterion::{BenchmarkId, Criterion, Throughput, criterion_group, criterion_main};
+use futures_util::FutureExt;
 use serde::{Deserialize, Serialize};
 use tokio::runtime::Runtime;
 use vespera_inprocess::{
@@ -485,6 +487,61 @@ fn bench_streaming_path(c: &mut Criterion) {
     drop(runtime);
 }
 
+/// #2 isolation: the `vespera_jni::dispatchAsync` spawn mechanism.
+///
+/// Both variants run the dispatch task on a shared multi-thread runtime
+/// (the outer `tokio::spawn`, common to both) and differ only in how a
+/// panic in the dispatch future is isolated:
+///
+/// - `double_spawn_pre`: a **second** `tokio::spawn` (panic → `JoinError`),
+///   the pre-#2 shape — one extra task allocation + scheduler hop.
+/// - `single_spawn_catch_unwind_post`: `FutureExt::catch_unwind` in place,
+///   the post-#2 shape — same panic → fallback, no second task.
+///
+/// The inner future is trivial so the spawn/catch_unwind overhead is the
+/// dominant cost and the delta isolates exactly what #2 removes per async
+/// dispatch (independent of the dispatch payload size).
+fn bench_async_spawn_pattern(c: &mut Criterion) {
+    let runtime = tokio::runtime::Builder::new_multi_thread()
+        .worker_threads(4)
+        .enable_all()
+        .build()
+        .expect("multi-thread runtime");
+    let mut group = c.benchmark_group("async_spawn_pattern");
+
+    group.bench_function("double_spawn_pre", |b| {
+        b.iter(|| {
+            runtime.block_on(async {
+                tokio::spawn(async move {
+                    tokio::spawn(async { vec![0u8; 64] })
+                        .await
+                        .unwrap_or_else(|_| vec![1u8; 16])
+                })
+                .await
+                .unwrap()
+            })
+        });
+    });
+
+    group.bench_function("single_spawn_catch_unwind_post", |b| {
+        b.iter(|| {
+            runtime.block_on(async {
+                tokio::spawn(async move {
+                    AssertUnwindSafe(async { vec![0u8; 64] })
+                        .catch_unwind()
+                        .await
+                        .unwrap_or_else(|_| vec![1u8; 16])
+                })
+                .await
+                .unwrap()
+            })
+        });
+    });
+
+    group.finish();
+    drop(runtime);
+}
+
 criterion_group!(
     benches,
     bench_router_path,
@@ -493,6 +550,7 @@ criterion_group!(
     bench_resolve_path,
     bench_contended_path,
     bench_headers_path,
-    bench_streaming_path
+    bench_streaming_path,
+    bench_async_spawn_pattern
 );
 criterion_main!(benches);

crates/vespera_inprocess/benches/dispatch.rs

@@ -247,14 +247,29 @@ struct WireHeaders<'a>(&'a http::HeaderMap);
 impl Serialize for WireHeaders<'_> {
     fn serialize<S: serde::Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
         use serde::ser::SerializeMap;
-        // `HeaderMap::keys` yields each distinct name exactly once;
-        // pre-size to the exact distinct-key count so the collect never
-        // reallocates.
-        let mut names: Vec<&str> = Vec::with_capacity(self.0.keys_len());
-        names.extend(self.0.keys().map(http::HeaderName::as_str));
+        // `HeaderMap::keys` yields each distinct name exactly once.  The
+        // overwhelmingly common response carries only a handful of header
+        // names, so sort them in a stack buffer and skip the per-response
+        // heap `Vec`; header sets larger than the stack cap fall back to a
+        // heap `Vec`.  Output is byte-identical either way (same sorted
+        // order over the same names), as locked by tests/wire_contract.rs.
+        const STACK_CAP: usize = 32;
+        let key_count = self.0.keys_len();
+        let mut stack_names: [&str; STACK_CAP] = [""; STACK_CAP];
+        let mut heap_names: Vec<&str>;
+        let names: &mut [&str] = if key_count <= STACK_CAP {
+            for (slot, name) in stack_names.iter_mut().zip(self.0.keys()) {
+                *slot = name.as_str();
+            }
+            &mut stack_names[..key_count]
+        } else {
+            heap_names = Vec::with_capacity(key_count);
+            heap_names.extend(self.0.keys().map(http::HeaderName::as_str));
+            &mut heap_names[..]
+        };
         names.sort_unstable();
         let mut map = serializer.serialize_map(Some(names.len()))?;
-        for name in names {
+        for &name in names.iter() {
             let mut values = self.0.get_all(name).iter();
             let first = values
                 .next()

crates/vespera_inprocess/src/wire.rs

@@ -10,6 +10,11 @@ repository.workspace = true
 vespera_inprocess = { workspace = true }
 jni = "0.22"
 tokio = { version = "1", features = ["rt-multi-thread"] }
+# `FutureExt::catch_unwind` for the async dispatch panic-isolation path
+# (replaces a redundant second `tokio::spawn`).  Already in the workspace
+# dependency tree via tokio/axum/tower, so this adds no new crate to the
+# build — only `std` is needed for the `catch_unwind` combinator.
+futures-util = { version = "0.3", default-features = false, features = ["std"] }
 # 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.

crates/vespera_jni/Cargo.toml

@@ -1,5 +1,6 @@
 use std::{cell::RefCell, future::Future, sync::LazyLock};
 
+use futures_util::FutureExt;
 use jni::EnvUnowned;
 use jni::errors::ThrowRuntimeExAndDefault;
 use jni::objects::{Global, JByteArray, JByteBuffer, JClass, JObject};
@@ -568,16 +569,27 @@ pub extern "system" fn Java_com_devfive_vespera_bridge_VesperaBridge_dispatchAsy
             }
         };
 
-        // The inner task converts Rust panics into JoinError, preserving
-        // always-complete semantics for the Java future.  Scheduling
-        // itself is wrapped in `catch_unwind` so a failure to build or
-        // schedule on the shared runtime completes the future (with a
-        // 500) instead of leaving the Java caller hanging.
+        // A panic in the dispatch future is caught **in place** with
+        // `FutureExt::catch_unwind` instead of isolating it in a second
+        // `tokio::spawn` task — same panic → 500 wire fallback (preserving
+        // always-complete semantics for the Java future), but one fewer
+        // task allocation + scheduler hop per async dispatch.  The inner
+        // spawn never bought parallelism here (the outer task awaited it
+        // immediately), so it was pure overhead.  `AssertUnwindSafe` is
+        // sound: a panic drops the half-run dispatch and we return a fresh
+        // `error_wire`; the registered `Router` is `Arc`-shared and is not
+        // left observably inconsistent.  The outer `catch_unwind` still
+        // guards `RUNTIME.spawn` itself so a scheduling failure completes
+        // the future (with a 500) instead of leaving the Java caller
+        // hanging.
         let scheduled = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
             RUNTIME.spawn(async move {
-                let response = tokio::spawn(vespera_inprocess::dispatch_from_bytes_async(input))
-                    .await
-                    .unwrap_or_else(|_| vespera_inprocess::error_wire(500, "panic in Rust engine"));
+                let response = std::panic::AssertUnwindSafe(
+                    vespera_inprocess::dispatch_from_bytes_async(input),
+                )
+                .catch_unwind()
+                .await
+                .unwrap_or_else(|_| vespera_inprocess::error_wire(500, "panic in Rust engine"));
 
                 let _ = with_cached_daemon_env(&jvm, |env| -> jni::errors::Result<()> {
                     complete_future(env, &future_for_task, &response)