From 4842388b20939b7ddc2cb6810be1aab1919bb372 Mon Sep 17 00:00:00 2001 From: Mohammad Dashti Date: Tue, 13 Oct 2026 22:45:14 -0700 Subject: [PATCH] Armed a stream from a Nexus read that waits and finds nothing. The Nexus read takes arm and answers armed, so a Nexus caller sleeps until its operation's progress says that the stream moved. --- .../nexus/temporal_streams.nexusrpc.yaml | 13 +++ crates/streams/src/nexus/mod.rs | 8 +- crates/streams/src/nexus/wire.rs | 4 + crates/streams/tests/nexus_contract.rs | 2 + crates/streams/tests/nexus_service.rs | 108 +++++++++++++++++- 5 files changed, 133 insertions(+), 2 deletions(-) diff --git a/crates/streams/nexus/temporal_streams.nexusrpc.yaml b/crates/streams/nexus/temporal_streams.nexusrpc.yaml index f34570d8c3..f160afec35 100644 --- a/crates/streams/nexus/temporal_streams.nexusrpc.yaml +++ b/crates/streams/nexus/temporal_streams.nexusrpc.yaml @@ -188,6 +188,13 @@ $defs: producer attempts the reader already saw, so its superseded and stale marks are exact across calls and across handler Workers. A state that doesn't parse is refused as a cursor mistake. + arm: + type: boolean + description: >- + For a reader that sleeps until its caller hears of progress. When + the call waits its whole wait_ms and collects nothing, the endpoint + arms the stream, so the next append notifies the stream's callers. + A reader with no callback leaves it out. required: [stream] additionalProperties: false @@ -279,6 +286,12 @@ $defs: also when the answer was cut short by its byte budget. It holds at most 64 KiB, which is about 87 KiB as base64, so it can ride in a Workflow's History. A longer state is refused. + armed: + type: boolean + description: >- + The call asked to arm, collected nothing and armed the stream, so + the caller waits for progress. Never true with records or with done. + An empty answer without it falls back to waiting for progress too. required: [records, next_token, done] additionalProperties: true diff --git a/crates/streams/src/nexus/mod.rs b/crates/streams/src/nexus/mod.rs index a17e29a289..00c40a88ae 100644 --- a/crates/streams/src/nexus/mod.rs +++ b/crates/streams/src/nexus/mod.rs @@ -307,6 +307,7 @@ impl NexusCalls { next_token: latest.cursor, done: false, state: None, + armed: false, }), Err(error) => failed(failure::of_stream(&error)), }; @@ -340,7 +341,7 @@ impl NexusCalls { wait: wait.try_into().ok(), max_records: max_records as u32, state: state.clone(), - arm: false, + arm: input.arm.unwrap_or(false), }; let mut read = match self.streams.read_carried(request).await { Ok(read) => read, @@ -613,6 +614,7 @@ impl Budget { next_token: read.response.cursor.clone(), done: false, state: Some(Base64::default()), + armed: false, }); let state = crossing_size(read.state.encoded_len().min(MAX_STATE_BYTES)); Self { @@ -687,6 +689,8 @@ fn finish_read( wires.push(wire); kept.extend(stored); } + // An answer the budget cut holds records, so only a whole empty answer stays armed. + let armed = response.armed && !cut && wires.is_empty(); let (next_token, done) = if cut { // A supersession belongs to the record after it, so it goes into the next answer too. if wires.last().is_some_and(|wire| wire.superseded.is_some()) { @@ -707,6 +711,7 @@ fn finish_read( next_token, done, state: Some(Base64(state.encode_to_vec())), + armed, }) } @@ -780,6 +785,7 @@ fn read_failure(error: StreamError, state: Vec) -> NexusStep { next_token: cursor, done: false, state: (!state.is_empty()).then_some(Base64(state)), + armed: false, }), _ => failed(failure::of_stream(&error)), } diff --git a/crates/streams/src/nexus/wire.rs b/crates/streams/src/nexus/wire.rs index 5c9ba45b6b..6cef74ee41 100644 --- a/crates/streams/src/nexus/wire.rs +++ b/crates/streams/src/nexus/wire.rs @@ -96,6 +96,8 @@ pub(crate) struct ReadInput { /// Left as base64, so an oversized state is refused before it is decoded. #[serde(default)] pub(crate) state: Option, + #[serde(default)] + pub(crate) arm: Option, } #[derive(Debug, Clone, PartialEq, Eq, Serialize)] @@ -105,6 +107,8 @@ pub(crate) struct ReadOutput { pub(crate) done: bool, #[serde(skip_serializing_if = "Option::is_none")] pub(crate) state: Option, + #[serde(skip_serializing_if = "std::ops::Not::not")] + pub(crate) armed: bool, } #[derive(Debug, Clone, PartialEq, Eq, Serialize)] diff --git a/crates/streams/tests/nexus_contract.rs b/crates/streams/tests/nexus_contract.rs index ce8790c061..300e9b6c98 100644 --- a/crates/streams/tests/nexus_contract.rs +++ b/crates/streams/tests/nexus_contract.rs @@ -58,6 +58,8 @@ fn full_input(contract: &Value, name: &str, workflow_id: &str) -> Value { ("ReadInput", "wait_ms") => json!(0), ("ReadInput", "latest_only") => json!(false), ("ReadInput", "state") => Value::Null, + // So the answer carries `armed`, which the output check holds to the contract. + ("ReadInput", "arm") => json!(true), (model, member) => panic!("the contract added {model}.{member}; teach this test"), }; out.insert(member, value); diff --git a/crates/streams/tests/nexus_service.rs b/crates/streams/tests/nexus_service.rs index 38e9581c54..c4f7d80692 100644 --- a/crates/streams/tests/nexus_service.rs +++ b/crates/streams/tests/nexus_service.rs @@ -22,7 +22,7 @@ use temporalio_common::protos::temporal::api::{ use temporalio_streams::{ MemoryStore, OwnerClient, OwnerDescription, OwnerError, StreamStore, Streams, StreamsOptions, nexus::{NexusCalls, NexusOptions}, - proto::{ChainId, NexusFailure, StreamRecord, StreamRecordKind}, + proto::{ChainId, NexusFailure, StoreAppendRequest, StreamRecord, StreamRecordKind}, }; static IDS: AtomicUsize = AtomicUsize::new(0); @@ -633,6 +633,108 @@ mod cases { assert!(small.get("error").is_none()); assert_eq!(texts(&answer), ["small"]); } + fn arming(wf: &str, after: &str) -> Value { + let mut input = read_input(wf, after); + input["arm"] = json!(true); + input + } + + /// Appends one record to the stream through the store, answering the arm it reported. + async fn arm_reported(store: &Arc, wf: &str, sequence: i64) -> Option { + store + .append(StoreAppendRequest { + chain: Some(ChainId { + namespace: NAMESPACE.to_string(), + workflow_id: wf.to_string(), + first_run_id: "run-1".to_string(), + }), + topic: TOPIC.to_string(), + producer_id: "store".to_string(), + attempt: 1, + sequence, + digest: vec![sequence as u8], + records: vec![ + StreamRecord { + topic: TOPIC.to_string(), + ..Default::default() + } + .encode_to_vec(), + ], + }) + .await + .unwrap() + .arm + } + + pub(crate) async fn an_empty_read_that_asks_to_arm_answers_armed(store: Arc) { + let calls = service(store.clone(), NexusOptions::default()); + let wf = workflow(); + let answer = read(&calls, arming(&wf, "")).await.unwrap(); + assert!(records(&answer).is_empty()); + assert_eq!(answer["armed"], true, "{answer}"); + assert_eq!(answer["done"], false); + assert!( + arm_reported(&store, &wf, 1) + .await + .is_some_and(|arm| arm > 0) + ); + } + + pub(crate) async fn an_answer_with_records_never_says_armed(store: Arc) { + let calls = service(store.clone(), NexusOptions::default()); + let wf = workflow(); + write(&calls, &wf, 1, &["one"]).await; + let answer = read(&calls, arming(&wf, "")).await.unwrap(); + assert_eq!(texts(&answer), ["one"]); + assert!( + answer.get("armed").is_none_or(|armed| armed == false), + "{answer}" + ); + assert_eq!(arm_reported(&store, &wf, 1).await, Some(0)); + } + + pub(crate) async fn a_done_answer_never_arms(store: Arc) { + let wf = workflow(); + write( + &service(store.clone(), NexusOptions::default()), + &wf, + 1, + &["one"], + ) + .await; + let ended = NexusCalls::new( + Arc::new(Streams::new( + store, + Arc::new(Ended), + StreamsOptions::default(), + )), + NexusOptions::default(), + ); + let mut after = String::new(); + for _ in 0..10 { + let mut input = arming(&wf, &after); + input["wait_ms"] = json!(5000); + let answer = read(&ended, input).await.unwrap(); + if answer["done"] == true { + assert!( + answer.get("armed").is_none_or(|armed| armed == false), + "{answer}" + ); + return; + } + after = token(&answer).to_string(); + } + panic!("the read never said done"); + } + + pub(crate) async fn a_read_without_arm_never_arms(store: Arc) { + let calls = service(store.clone(), NexusOptions::default()); + let wf = workflow(); + let answer = read_after(&calls, &wf, "").await; + assert!(records(&answer).is_empty()); + assert!(answer.get("armed").is_none(), "{answer}"); + assert_eq!(arm_reported(&store, &wf, 1).await, Some(0)); + } } macro_rules! lanes { @@ -671,6 +773,10 @@ macro_rules! all_cases { a_record_above_the_budget_still_crosses_alone, an_append_above_the_record_limit_is_refused, a_stored_record_above_the_limit_crosses_as_an_error, + an_empty_read_that_asks_to_arm_answers_armed, + an_answer_with_records_never_says_armed, + a_done_answer_never_arms, + a_read_without_arm_never_arms, ); }; }