1use std::{any::Any, cell::Cell, fmt, marker::PhantomData, rc::Rc};
2
3use futures::future::LocalBoxFuture;
4
5use super::{
6 DiagnosticEvent, DiagnosticOutcome, DiagnosticSource, InvocationContext, NativeAppRuntime,
7 NativeStreamEndpointBinding, RequestPermit, RuntimeFailure, await_with_generation_context,
8 diagnostics::diagnostic_operation, schedule_module_supervision_after_failure,
9};
10
11pub trait StreamCapability: 'static {
13 type OpenRequest: 'static;
15 type Message: 'static;
17 type DomainError: 'static;
19 const ID: &'static str;
21 const DESCRIPTOR_VERSION: &'static str;
23}
24
25#[derive(Clone, Debug, PartialEq)]
27pub enum StreamEvent<M, E> {
28 Message(M),
30 PeerHalfClosed,
32 Terminal(Result<(), E>),
34}
35
36pub enum NativeStreamItem {
38 Message(Box<dyn Any>),
40 PeerHalfClosed,
42 Terminal(Result<(), Box<dyn Any>>),
44}
45
46impl fmt::Debug for NativeStreamItem {
47 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
48 match self {
49 Self::Message(_) => formatter.write_str("Message(<erased>)"),
50 Self::PeerHalfClosed => formatter.write_str("PeerHalfClosed"),
51 Self::Terminal(Ok(())) => formatter.write_str("Terminal(Ok(()))"),
52 Self::Terminal(Err(_)) => formatter.write_str("Terminal(Err(<erased>))"),
53 }
54 }
55}
56
57pub trait NativeStreamSession: fmt::Debug {
59 fn send(&self, message: Box<dyn Any>) -> LocalBoxFuture<'static, Result<(), RuntimeFailure>>;
61 fn receive(&self) -> LocalBoxFuture<'static, Result<NativeStreamItem, RuntimeFailure>>;
63 fn close_send(&self) -> LocalBoxFuture<'static, Result<(), RuntimeFailure>>;
65 fn cancel(&self);
67}
68
69pub type NativeStreamOpenFuture = LocalBoxFuture<
71 'static,
72 Result<Result<Box<dyn NativeStreamSession>, Box<dyn Any>>, RuntimeFailure>,
73>;
74
75pub trait NativeStreamEndpoint: fmt::Debug {
77 fn capability_id(&self) -> &'static str;
79 fn descriptor_version(&self) -> &'static str;
81 fn operations(&self) -> &'static [&'static str];
83 fn open(
85 &self,
86 operation: &str,
87 request: Box<dyn Any>,
88 context: InvocationContext,
89 ) -> NativeStreamOpenFuture;
90}
91
92#[derive(Debug)]
94pub struct NativeStreamHandle<C: StreamCapability> {
95 endpoints: Vec<NativeStreamEndpointBinding>,
96 runtime: Rc<NativeAppRuntime>,
97 caller_instance: String,
98 allow_before_ready: bool,
99 capability: PhantomData<fn() -> C>,
100}
101
102impl<C: StreamCapability> NativeStreamHandle<C> {
103 pub(crate) fn from_endpoints(
104 endpoints: &[NativeStreamEndpointBinding],
105 runtime: Rc<NativeAppRuntime>,
106 caller_instance: &str,
107 allow_before_ready: bool,
108 ) -> Self {
109 Self {
110 endpoints: endpoints.to_vec(),
111 runtime,
112 caller_instance: caller_instance.to_owned(),
113 allow_before_ready,
114 capability: PhantomData,
115 }
116 }
117
118 pub fn binding_count(&self) -> usize {
120 self.endpoints.len()
121 }
122
123 pub async fn open(
125 &self,
126 operation: &str,
127 request: C::OpenRequest,
128 ) -> Result<Result<NativeStream<C>, C::DomainError>, RuntimeFailure> {
129 let context = self.next_context();
130 self.open_with_context(operation, context, request).await
131 }
132
133 pub async fn open_with_context(
135 &self,
136 operation: &str,
137 context: InvocationContext,
138 request: C::OpenRequest,
139 ) -> Result<Result<NativeStream<C>, C::DomainError>, RuntimeFailure> {
140 let context = context
141 .for_caller(&self.caller_instance)
142 .for_target(C::ID, operation);
143 let started_at = (self.runtime.driver.now)();
144 let operation_name = self
145 .endpoints
146 .first()
147 .and_then(|endpoint| diagnostic_operation(endpoint.state.operations, operation));
148 let request_id = context.request_id();
149 self.runtime
150 .diagnostics
151 .emit(DiagnosticSource::Invocation, started_at, |_| {
152 DiagnosticEvent::InvocationStarted {
153 request_id,
154 caller_instance: Some(self.caller_instance.clone()),
155 provider_instance: self
156 .endpoints
157 .first()
158 .map(|endpoint| endpoint.module_instance.clone()),
159 capability: C::ID,
160 operation: operation_name,
161 }
162 });
163 let result = self
164 .open_with_context_inner(operation, context, request)
165 .await;
166 let outcome = match &result {
167 Ok(Ok(_)) => DiagnosticOutcome::Succeeded,
168 Ok(Err(_)) => DiagnosticOutcome::DomainError,
169 Err(error) => DiagnosticOutcome::RuntimeFailure(error.into()),
170 };
171 self.runtime.diagnostics.emit(
172 DiagnosticSource::Invocation,
173 (self.runtime.driver.now)(),
174 |_| DiagnosticEvent::InvocationCompleted {
175 request_id,
176 caller_instance: Some(self.caller_instance.clone()),
177 provider_instance: self
178 .endpoints
179 .first()
180 .map(|endpoint| endpoint.module_instance.clone()),
181 capability: C::ID,
182 operation: operation_name,
183 outcome,
184 elapsed: (self.runtime.driver.now)().saturating_sub(started_at),
185 },
186 );
187 if let Err(error) = &result {
188 self.runtime.diagnostics.emit_runtime_failure(
189 (self.runtime.driver.now)(),
190 self.endpoints
191 .first()
192 .map(|endpoint| endpoint.module_instance.as_str()),
193 error,
194 );
195 }
196 result
197 }
198
199 async fn open_with_context_inner(
200 &self,
201 operation: &str,
202 context: InvocationContext,
203 request: C::OpenRequest,
204 ) -> Result<Result<NativeStream<C>, C::DomainError>, RuntimeFailure> {
205 if self.runtime.shutdown_started.get()
206 || (!self.allow_before_ready && self.runtime.admission.is_closed())
207 {
208 return Err(RuntimeFailure::AdmissionClosed);
209 }
210 let endpoint = match self.endpoints.as_slice() {
211 [] => return Err(RuntimeFailure::Unavailable { capability: C::ID }),
212 [endpoint] => endpoint,
213 endpoints => {
214 return Err(RuntimeFailure::AmbiguousBinding {
215 capability: C::ID,
216 providers: endpoints.len(),
217 });
218 }
219 };
220 let snapshot = endpoint
221 .state
222 .snapshot()
223 .ok_or(RuntimeFailure::Unavailable { capability: C::ID })?;
224 let admission = endpoint
225 .admission(operation)
226 .ok_or_else(|| RuntimeFailure::UnknownOperation {
227 capability: C::ID,
228 operation: operation.to_owned(),
229 })?
230 .clone();
231 let permit = admission
232 .acquire(C::ID, operation, &context, &self.runtime.driver)
233 .await?;
234 if !endpoint.state.is_current(snapshot.generation) {
235 return Err(RuntimeFailure::Unavailable { capability: C::ID });
236 }
237 let generation_cancellation = snapshot.cancellation.clone();
238 let outcome = await_with_generation_context(
239 &self.runtime.driver,
240 &context,
241 snapshot.cancellation,
242 C::ID,
243 snapshot
244 .endpoint
245 .open(operation, Box::new(request), context.clone()),
246 )
247 .await
248 .map_err(|error| {
249 schedule_module_supervision_after_failure(
250 &self.runtime,
251 &endpoint.module_instance,
252 error,
253 )
254 })?
255 .map_err(|error| {
256 schedule_module_supervision_after_failure(
257 &self.runtime,
258 &endpoint.module_instance,
259 error,
260 )
261 })?;
262 match outcome {
263 Ok(session) => Ok(Ok(NativeStream::new(
264 session,
265 self.runtime.clone(),
266 generation_cancellation,
267 endpoint.module_instance.clone(),
268 context,
269 permit,
270 ))),
271 Err(error) => Ok(Err(error
272 .downcast::<C::DomainError>()
273 .map(|error| *error)
274 .map_err(|_| RuntimeFailure::ProtocolViolation { capability: C::ID })?)),
275 }
276 }
277
278 fn next_context(&self) -> InvocationContext {
279 InvocationContext::new(
280 self.next_request_id(),
281 None,
282 super::CancellationToken::new(),
283 )
284 .with_caller_instance(self.caller_instance.clone())
285 }
286
287 fn next_request_id(&self) -> super::RequestId {
288 let request_id = self.runtime.request_ids.get();
289 self.runtime.request_ids.set(request_id.saturating_add(1));
290 request_id
291 }
292}
293
294#[derive(Debug)]
296pub struct NativeStream<C: StreamCapability> {
297 inner: Rc<dyn NativeStreamSession>,
298 runtime: Rc<NativeAppRuntime>,
299 generation_cancellation: super::CancellationToken,
300 module_instance: String,
301 context: InvocationContext,
302 _permit: RequestPermit,
303 local_half_closed: Cell<bool>,
304 peer_half_closed: Cell<bool>,
305 terminal_seen: Cell<bool>,
306 cancelled: Cell<bool>,
307 capability: PhantomData<fn() -> C>,
308}
309
310impl<C: StreamCapability> NativeStream<C> {
311 fn new(
312 session: Box<dyn NativeStreamSession>,
313 runtime: Rc<NativeAppRuntime>,
314 generation_cancellation: super::CancellationToken,
315 module_instance: String,
316 context: InvocationContext,
317 permit: RequestPermit,
318 ) -> Self {
319 Self {
320 inner: Rc::from(session),
321 runtime,
322 generation_cancellation,
323 module_instance,
324 context,
325 _permit: permit,
326 local_half_closed: Cell::new(false),
327 peer_half_closed: Cell::new(false),
328 terminal_seen: Cell::new(false),
329 cancelled: Cell::new(false),
330 capability: PhantomData,
331 }
332 }
333
334 pub async fn send(&self, message: C::Message) -> Result<(), RuntimeFailure> {
336 if let Some(error) = self.cancelled_outcome() {
337 return Err(error);
338 }
339 if self.local_half_closed.get() || self.terminal_seen.get() {
340 return Err(Self::protocol_violation());
341 }
342 let inner = self.inner.clone();
343 await_with_generation_context(
344 &self.runtime.driver,
345 &self.context,
346 self.generation_cancellation.clone(),
347 C::ID,
348 inner.send(Box::new(message)),
349 )
350 .await
351 .map_err(|error| self.finish_with_error(error))?
352 .map_err(|error| self.finish_with_error(error))
353 }
354
355 pub async fn receive(&self) -> Result<StreamEvent<C::Message, C::DomainError>, RuntimeFailure> {
357 if let Some(error) = self.cancelled_outcome() {
358 return Err(error);
359 }
360 if self.terminal_seen.get() {
361 return Err(Self::protocol_violation());
362 }
363 let inner = self.inner.clone();
364 let item = await_with_generation_context(
365 &self.runtime.driver,
366 &self.context,
367 self.generation_cancellation.clone(),
368 C::ID,
369 inner.receive(),
370 )
371 .await
372 .map_err(|error| self.finish_with_error(error))?
373 .map_err(|error| self.finish_with_error(error))?;
374 match item {
375 super::NativeStreamItem::Message(message) => {
376 if self.peer_half_closed.get() {
377 return Err(self.finish_with_error(Self::protocol_violation()));
378 }
379 message
380 .downcast::<C::Message>()
381 .map(|message| StreamEvent::Message(*message))
382 .map_err(|_| self.finish_with_error(Self::protocol_violation()))
383 }
384 super::NativeStreamItem::PeerHalfClosed => {
385 if self.peer_half_closed.replace(true) {
386 return Err(self.finish_with_error(Self::protocol_violation()));
387 }
388 Ok(StreamEvent::PeerHalfClosed)
389 }
390 super::NativeStreamItem::Terminal(outcome) => {
391 if self.terminal_seen.replace(true) {
392 return Err(self.finish_with_error(Self::protocol_violation()));
393 }
394 let outcome = match outcome {
395 Ok(()) => Ok(()),
396 Err(error) => Err(error
397 .downcast::<C::DomainError>()
398 .map(|error| *error)
399 .map_err(|_| self.finish_with_error(Self::protocol_violation()))?),
400 };
401 Ok(StreamEvent::Terminal(outcome))
402 }
403 }
404 }
405
406 pub async fn close_send(&self) -> Result<(), RuntimeFailure> {
408 if let Some(error) = self.cancelled_outcome() {
409 return Err(error);
410 }
411 if self.terminal_seen.get() || self.local_half_closed.replace(true) {
412 return Err(Self::protocol_violation());
413 }
414 let inner = self.inner.clone();
415 let result = await_with_generation_context(
416 &self.runtime.driver,
417 &self.context,
418 self.generation_cancellation.clone(),
419 C::ID,
420 inner.close_send(),
421 )
422 .await
423 .map_err(|error| self.finish_with_error(error))?
424 .map_err(|error| self.finish_with_error(error));
425 let resource_exhausted = result
426 .as_ref()
427 .err()
428 .is_some_and(|error| matches!(error, RuntimeFailure::ResourceExhausted { .. }));
429 if resource_exhausted {
430 self.local_half_closed.set(false);
431 }
432 result
433 }
434
435 pub fn cancel(&self) {
437 if !self.terminal_seen.get() && !self.cancelled.replace(true) {
438 self.context.cancellation().cancel();
439 self.inner.cancel();
440 }
441 }
442
443 pub const fn request_id(&self) -> super::RequestId {
445 self.context.request_id()
446 }
447
448 fn protocol_violation() -> RuntimeFailure {
449 RuntimeFailure::ProtocolViolation { capability: C::ID }
450 }
451
452 fn cancelled_outcome(&self) -> Option<RuntimeFailure> {
453 if !self.cancelled.get() {
454 return None;
455 }
456 if self.terminal_seen.replace(true) {
457 Some(Self::protocol_violation())
458 } else {
459 Some(RuntimeFailure::Cancelled {
460 request_id: self.context.request_id(),
461 })
462 }
463 }
464
465 fn schedule_failure(&self, error: RuntimeFailure) -> RuntimeFailure {
466 schedule_module_supervision_after_failure(&self.runtime, &self.module_instance, error)
467 }
468
469 fn finish_with_error(&self, error: RuntimeFailure) -> RuntimeFailure {
470 let error = self.schedule_failure(error);
471 self.runtime.diagnostics.emit_runtime_failure(
472 (self.runtime.driver.now)(),
473 Some(&self.module_instance),
474 &error,
475 );
476 if !matches!(error, RuntimeFailure::ResourceExhausted { .. }) {
477 self.terminal_seen.set(true);
478 if !self.cancelled.replace(true) {
479 self.context.cancellation().cancel();
480 self.inner.cancel();
481 }
482 }
483 error
484 }
485}
486
487impl<C: StreamCapability> Drop for NativeStream<C> {
488 fn drop(&mut self) {
489 if !self.cancelled.replace(true) && !self.terminal_seen.get() {
490 self.inner.cancel();
491 }
492 }
493}
494
495pub type StreamSession<C> = NativeStream<C>;