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