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 .for_caller(&self.caller_instance)
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(C::ID, operation, &context, &self.runtime.driver)
230 .await?;
231 if !endpoint.state.is_current(snapshot.generation) {
232 return Err(RuntimeFailure::Unavailable { capability: C::ID });
233 }
234 let generation_cancellation = snapshot.cancellation.clone();
235 let outcome = await_with_generation_context(
236 &self.runtime.driver,
237 &context,
238 snapshot.cancellation,
239 C::ID,
240 snapshot
241 .endpoint
242 .open(operation, Box::new(request), context.clone()),
243 )
244 .await
245 .map_err(|error| {
246 schedule_module_supervision_after_failure(
247 &self.runtime,
248 &endpoint.module_instance,
249 error,
250 )
251 })?
252 .map_err(|error| {
253 schedule_module_supervision_after_failure(
254 &self.runtime,
255 &endpoint.module_instance,
256 error,
257 )
258 })?;
259 match outcome {
260 Ok(session) => Ok(Ok(NativeStream::new(
261 session,
262 self.runtime.clone(),
263 generation_cancellation,
264 endpoint.module_instance.clone(),
265 context,
266 permit,
267 ))),
268 Err(error) => Ok(Err(error
269 .downcast::<C::DomainError>()
270 .map(|error| *error)
271 .map_err(|_| RuntimeFailure::ProtocolViolation { capability: C::ID })?)),
272 }
273 }
274
275 fn next_context(&self) -> InvocationContext {
276 InvocationContext::new(
277 self.next_request_id(),
278 None,
279 super::CancellationToken::new(),
280 )
281 .with_caller_instance(self.caller_instance.clone())
282 }
283
284 fn next_request_id(&self) -> super::RequestId {
285 let request_id = self.runtime.request_ids.get();
286 self.runtime.request_ids.set(request_id.saturating_add(1));
287 request_id
288 }
289}
290
291#[derive(Debug)]
293pub struct NativeStream<C: StreamCapability> {
294 inner: Rc<dyn NativeStreamSession>,
295 runtime: Rc<NativeAppRuntime>,
296 generation_cancellation: super::CancellationToken,
297 module_instance: String,
298 context: InvocationContext,
299 _permit: RequestPermit,
300 local_half_closed: Cell<bool>,
301 peer_half_closed: Cell<bool>,
302 terminal_seen: Cell<bool>,
303 cancelled: Cell<bool>,
304 capability: PhantomData<fn() -> C>,
305}
306
307impl<C: StreamCapability> NativeStream<C> {
308 fn new(
309 session: Box<dyn NativeStreamSession>,
310 runtime: Rc<NativeAppRuntime>,
311 generation_cancellation: super::CancellationToken,
312 module_instance: String,
313 context: InvocationContext,
314 permit: RequestPermit,
315 ) -> Self {
316 Self {
317 inner: Rc::from(session),
318 runtime,
319 generation_cancellation,
320 module_instance,
321 context,
322 _permit: permit,
323 local_half_closed: Cell::new(false),
324 peer_half_closed: Cell::new(false),
325 terminal_seen: Cell::new(false),
326 cancelled: Cell::new(false),
327 capability: PhantomData,
328 }
329 }
330
331 pub async fn send(&self, message: C::Message) -> Result<(), RuntimeFailure> {
333 if let Some(error) = self.cancelled_outcome() {
334 return Err(error);
335 }
336 if self.local_half_closed.get() || self.terminal_seen.get() {
337 return Err(self.protocol_violation());
338 }
339 let inner = self.inner.clone();
340 await_with_generation_context(
341 &self.runtime.driver,
342 &self.context,
343 self.generation_cancellation.clone(),
344 C::ID,
345 inner.send(Box::new(message)),
346 )
347 .await
348 .map_err(|error| self.finish_with_error(error))?
349 .map_err(|error| self.finish_with_error(error))
350 }
351
352 pub async fn receive(&self) -> Result<StreamEvent<C::Message, C::DomainError>, RuntimeFailure> {
354 if let Some(error) = self.cancelled_outcome() {
355 return Err(error);
356 }
357 if self.terminal_seen.get() {
358 return Err(self.protocol_violation());
359 }
360 let inner = self.inner.clone();
361 let item = await_with_generation_context(
362 &self.runtime.driver,
363 &self.context,
364 self.generation_cancellation.clone(),
365 C::ID,
366 inner.receive(),
367 )
368 .await
369 .map_err(|error| self.finish_with_error(error))?
370 .map_err(|error| self.finish_with_error(error))?;
371 match item {
372 super::NativeStreamItem::Message(message) => {
373 if self.peer_half_closed.get() {
374 return Err(self.finish_with_error(self.protocol_violation()));
375 }
376 message
377 .downcast::<C::Message>()
378 .map(|message| StreamEvent::Message(*message))
379 .map_err(|_| self.finish_with_error(self.protocol_violation()))
380 }
381 super::NativeStreamItem::PeerHalfClosed => {
382 if self.peer_half_closed.replace(true) {
383 return Err(self.finish_with_error(self.protocol_violation()));
384 }
385 Ok(StreamEvent::PeerHalfClosed)
386 }
387 super::NativeStreamItem::Terminal(outcome) => {
388 if self.terminal_seen.replace(true) {
389 return Err(self.finish_with_error(self.protocol_violation()));
390 }
391 let outcome = match outcome {
392 Ok(()) => Ok(()),
393 Err(error) => Err(error
394 .downcast::<C::DomainError>()
395 .map(|error| *error)
396 .map_err(|_| self.finish_with_error(self.protocol_violation()))?),
397 };
398 Ok(StreamEvent::Terminal(outcome))
399 }
400 }
401 }
402
403 pub async fn close_send(&self) -> Result<(), RuntimeFailure> {
405 if let Some(error) = self.cancelled_outcome() {
406 return Err(error);
407 }
408 if self.terminal_seen.get() || self.local_half_closed.replace(true) {
409 return Err(self.protocol_violation());
410 }
411 let inner = self.inner.clone();
412 let result = await_with_generation_context(
413 &self.runtime.driver,
414 &self.context,
415 self.generation_cancellation.clone(),
416 C::ID,
417 inner.close_send(),
418 )
419 .await
420 .map_err(|error| self.finish_with_error(error))?
421 .map_err(|error| self.finish_with_error(error));
422 let resource_exhausted = result
423 .as_ref()
424 .err()
425 .is_some_and(|error| matches!(error, RuntimeFailure::ResourceExhausted { .. }));
426 if resource_exhausted {
427 self.local_half_closed.set(false);
428 }
429 result
430 }
431
432 pub fn cancel(&self) {
434 if !self.terminal_seen.get() && !self.cancelled.replace(true) {
435 self.context.cancellation().cancel();
436 self.inner.cancel();
437 }
438 }
439
440 pub const fn request_id(&self) -> super::RequestId {
442 self.context.request_id()
443 }
444
445 fn protocol_violation(&self) -> RuntimeFailure {
446 RuntimeFailure::ProtocolViolation { capability: C::ID }
447 }
448
449 fn cancelled_outcome(&self) -> Option<RuntimeFailure> {
450 if !self.cancelled.get() {
451 return None;
452 }
453 if self.terminal_seen.replace(true) {
454 Some(self.protocol_violation())
455 } else {
456 Some(RuntimeFailure::Cancelled {
457 request_id: self.context.request_id(),
458 })
459 }
460 }
461
462 fn schedule_failure(&self, error: RuntimeFailure) -> RuntimeFailure {
463 schedule_module_supervision_after_failure(&self.runtime, &self.module_instance, error)
464 }
465
466 fn finish_with_error(&self, error: RuntimeFailure) -> RuntimeFailure {
467 let error = self.schedule_failure(error);
468 self.runtime.diagnostics.emit_runtime_failure(
469 (self.runtime.driver.now)(),
470 Some(&self.module_instance),
471 &error,
472 );
473 if !matches!(error, RuntimeFailure::ResourceExhausted { .. }) {
474 self.terminal_seen.set(true);
475 if !self.cancelled.replace(true) {
476 self.context.cancellation().cancel();
477 self.inner.cancel();
478 }
479 }
480 error
481 }
482}
483
484impl<C: StreamCapability> Drop for NativeStream<C> {
485 fn drop(&mut self) {
486 if !self.cancelled.replace(true) && !self.terminal_seen.get() {
487 self.inner.cancel();
488 }
489 }
490}
491
492pub type StreamSession<C> = NativeStream<C>;