Skip to main content

harmos_sdk/operations/
serving.rs

1/* ***
2 *
3 * Harmos
4 *
5 * SPDX-License-Identifier: Apache-2.0
6 *
7 * ***
8 */
9
10/* === Modules === */
11
12mod codec;
13mod owned;
14#[cfg(feature = "grpc")]
15pub(crate) mod protobuf;
16#[cfg(feature = "streaming")]
17mod streams;
18
19/* === Re-Exports === */
20
21pub use codec::WireCodecError;
22#[cfg(feature = "streaming")]
23pub use streams::{WireEndpoints, WireStreamLimits, WireStreams};
24
25/* === Imports === */
26
27// Standard Library Imports
28use std::collections::BTreeMap;
29use std::future::Future;
30use std::num::NonZeroUsize;
31use std::pin::Pin;
32use std::sync::{Arc, Mutex, Weak};
33use std::task::{Poll, Waker};
34// Third Party Imports
35use serde::{Deserialize, Serialize};
36use tracing::Instrument;
37// Local Imports
38use super::{Contract, ContractError, ExecutionContext, FailureKind, OperationDefinition, RegistrationBatch};
39use crate::QualifiedId;
40
41/* === Constants === */
42
43pub const WIRE_METHOD: &str = "harmos.operation.v2";
44
45pub const CHILD_METHOD: &str = "harmos.operation.child.v2";
46
47/* === Macros === */
48
49/* === Type Aliases === */
50
51type Answer<'a> = Pin<Box<dyn Future<Output = Result<Vec<u8>, ServingError>> + Send + 'a>>;
52type Finished<'a> = Pin<Box<dyn Future<Output = Result<(), ServingError>> + Send + 'a>>;
53type Finalization<S> = Pin<Box<dyn Future<Output = (S, Result<(), ServingError>)> + Send + 'static>>;
54type Prepare = Box<dyn Fn(Vec<u8>, bool) -> Result<Box<dyn Prepared>, ServingError> + Send + Sync>;
55
56/* === Structs === */
57
58struct Contextual<F: std::future::Future> {
59    future: Option<std::pin::Pin<Box<F>>>,
60    dispatcher: tracing::Dispatch,
61}
62
63#[derive(Debug, Serialize, Deserialize)]
64#[serde(deny_unknown_fields)]
65pub struct WireChildCall {
66    pub context: [u8; 32],
67
68    #[serde(default, skip_serializing_if = "Option::is_none")]
69    pub generation_context: Option<[u8; 32]>,
70
71    #[serde(default, skip_serializing_if = "Vec::is_empty")]
72    pub values: Vec<super::CallValue>,
73
74    pub operation: String,
75    pub fingerprint: [u8; 32],
76    pub generation: u64,
77    pub payload: Vec<u8>,
78
79    // Submitted children answer their instance id at admission and remain
80    // owned by the calling attempt until joined, released or detached.
81    #[serde(default)]
82    pub submit: bool,
83}
84
85// Drives one child the same attempt submitted; the host resolves the instance
86// only inside the attempt context that submitted it.
87#[derive(Debug, Serialize, Deserialize)]
88#[serde(deny_unknown_fields)]
89pub struct WireChildControl {
90    pub context: [u8; 32],
91    pub instance: u64,
92    pub action: WireChildAction,
93}
94
95pub struct WireTable {
96    owner: String,
97    handlers: BTreeMap<String, Route>,
98    retired: BTreeMap<String, u64>,
99    services: BTreeMap<String, super::registration::ServiceFactory>,
100    prepared: BTreeMap<u64, Arc<Slot>>,
101    capacity: Option<NonZeroUsize>,
102    closed: bool,
103    ordinary_closed: bool,
104    provided: Vec<std::any::TypeId>,
105
106    #[cfg(feature = "streaming")]
107    stream_limits: Option<WireStreamLimits>,
108}
109
110struct Slot {
111    operation: String,
112    lifecycle: bool,
113    resident: bool,
114
115    #[cfg(feature = "streaming")]
116    streams: Option<WireStreams>,
117
118    state: Mutex<SlotState>,
119}
120
121struct SlotState {
122    prepared: Option<Box<dyn Prepared>>,
123    finished: bool,
124    cancelled: bool,
125    waiters: Vec<Weak<Mutex<Option<Waker>>>>,
126}
127
128struct Loan {
129    slot: Arc<Slot>,
130    prepared: Option<Box<dyn Prepared>>,
131    finished: bool,
132}
133
134struct Route {
135    callees: &'static [super::WireDeclaration],
136    lifecycle: bool,
137    // Lifecycle actions and the observe control are built in: they serve the
138    // plugin's machinery and never appear in its declared catalog.
139    builtin: bool,
140    // A resident control runs until cancelled and is not charged against the
141    // ordinary invocation capacity.
142    resident: bool,
143    definition: OperationDefinition,
144    generation: Option<u64>,
145    closing: bool,
146    contract: Option<Contract>,
147    prepare: Prepare,
148}
149
150/* === Enums === */
151
152// Answers: `Join` the child's response payload, `Status` and `Cancel` an
153// encoded instance state, `Release` and `Detach` an empty payload.
154// Releasing cancels the child; detaching leaves it running without a handle.
155#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
156#[serde(rename_all = "snake_case")]
157pub enum WireChildAction {
158    Join,
159    Status,
160    Cancel,
161    Release,
162    Detach,
163}
164
165#[derive(Debug)]
166pub enum WireChild {
167    Call(WireChildCall),
168    Control(WireChildControl),
169}
170
171#[derive(Debug, Serialize, Deserialize)]
172#[serde(tag = "phase", rename_all = "snake_case", deny_unknown_fields)]
173pub enum WireCommand {
174    Prepare { invocation: u64, operation: String, fingerprint: [u8; 32], generation: u64, payload: Vec<u8> },
175    Initialize { invocation: u64, context: [u8; 32] },
176    Attempt { invocation: u64, context: [u8; 32], values: Vec<super::CallValue> },
177    Finalize { invocation: u64 },
178    StreamAttempt { invocation: u64, epoch: u64, context: [u8; 32], values: Vec<super::CallValue> },
179    StreamSend { invocation: u64, epoch: u64, payload: Vec<u8> },
180    StreamFinish { invocation: u64, epoch: u64 },
181    StreamReceive { invocation: u64, epoch: u64 },
182    PrepareGrpc { invocation: u64, operation: String, fingerprint: [u8; 32], generation: u64, payload: Vec<u8> },
183}
184
185#[derive(Debug, Serialize, Deserialize)]
186#[serde(tag = "status", rename_all = "snake_case", deny_unknown_fields)]
187pub enum WireAnswer {
188    Stream { epoch: u64, transferred: bool, payload: Option<Vec<u8>> },
189    Ready { payload: Vec<u8> },
190    Failed { failure: crate::Failure },
191}
192
193#[derive(Debug, thiserror::Error)]
194pub enum ServingError {
195    #[error("caller may not invoke operation `{operation}` without a host assignment")]
196    Unassigned { operation: &'static str },
197
198    #[error("operation endpoint types differ from the declaration")]
199    EndpointTypes,
200
201    #[error("carrier invocation does not support stream commands")]
202    StreamUnsupported,
203
204    #[cfg(feature = "streaming")]
205    #[error("streaming catalog requires explicit provider stream limits")]
206    StreamLimits,
207
208    #[cfg(feature = "streaming")]
209    #[error("carrier stream failed: {0}")]
210    Stream(#[source] super::streaming::Error),
211
212    #[cfg(feature = "streaming")]
213    #[error("carrier stream direction already has an outstanding request")]
214    StreamBusy,
215
216    #[cfg(feature = "streaming")]
217    #[error("carrier stream direction or attempt is not available")]
218    StreamPhase,
219
220    #[error(transparent)]
221    Lifecycle(#[from] harmos_domain::model::operations::lifecycle::Error),
222
223    #[error(transparent)]
224    Contract(#[from] ContractError),
225
226    #[error(transparent)]
227    Declaration(#[from] crate::Error),
228
229    #[error("operation `{operation}` declares resource `{resource}` that its registration does not provide")]
230    Resource { operation: String, resource: &'static str },
231
232    #[error("operation `{operation}` has no bound carrier handler")]
233    Unknown { operation: String },
234
235    #[error("operation `{operation}` requires streaming endpoints unsupported by this carrier binding")]
236    Streaming { operation: String },
237
238    #[error("operation `{operation}` has a mismatched fingerprint or stale generation")]
239    Identity { operation: String },
240
241    #[error("carrier invocation `{invocation}` is already prepared")]
242    Duplicate { invocation: u64 },
243
244    #[error("carrier invocation `{invocation}` has not been prepared")]
245    Missing { invocation: u64 },
246
247    #[error("carrier operation has not completed initialization")]
248    Uninitialized,
249
250    #[error("carrier operation initialization already started and cannot be replayed")]
251    InitializationStarted,
252
253    #[error("carrier invocation capacity is exhausted")]
254    Capacity,
255
256    #[error("carrier ordinary invocations have not finished settlement")]
257    Unsettled,
258
259    #[error("carrier operation admission is closing")]
260    Closing,
261
262    #[error("carrier handler definitions differ from the selected operation catalog")]
263    Catalog,
264
265    #[error("operation payload codec failed")]
266    Codec(#[source] crate::MessageError),
267
268    #[error("carrier handler panicked")]
269    Panicked,
270
271    // The handler's failure, written once at its origin.
272    #[error(transparent)]
273    Handler(crate::Failure),
274
275    #[error("native handler failed: {0}")]
276    Native(#[source] std::sync::Arc<dyn std::error::Error + Send + Sync>),
277}
278
279/* === Traits === */
280
281trait InContext: std::future::Future + Sized {
282    fn in_context(self, dispatcher: tracing::Dispatch) -> Contextual<Self> {
283        Contextual { future: Some(Box::pin(self)), dispatcher }
284    }
285
286    fn in_current_context(self) -> Contextual<Self> {
287        self.in_context(tracing::dispatcher::get_default(Clone::clone))
288    }
289}
290
291trait Prepared: Send {
292    #[cfg(feature = "streaming")]
293    fn streams(&self) -> Option<WireStreams> {
294        None
295    }
296
297    // The grant the prepared request needs, answered to the host with Prepare.
298    fn grant(&self) -> Option<crate::plugin::Grant> {
299        None
300    }
301    fn bind<'a>(&'a mut self, context: &'a ExecutionContext) -> Finished<'a>;
302    fn attempt<'a>(&'a mut self, context: &'a ExecutionContext) -> Answer<'a>;
303    fn finalize(&mut self) -> Finished<'_>;
304}
305
306/* === Implementations === */
307
308impl<F: std::future::Future> InContext for F {}
309
310impl<F: std::future::Future> std::future::Future for Contextual<F> {
311    type Output = F::Output;
312
313    fn poll(self: std::pin::Pin<&mut Self>, context: &mut std::task::Context<'_>) -> std::task::Poll<Self::Output> {
314        let owner = self.get_mut();
315        let _dispatcher = tracing::dispatcher::set_default(&owner.dispatcher);
316        owner.future.as_mut().expect("context future remains until drop").as_mut().poll(context)
317    }
318}
319
320impl<F: std::future::Future> Drop for Contextual<F> {
321    fn drop(&mut self) {
322        tracing::dispatcher::with_default(&self.dispatcher, || drop(self.future.take()));
323    }
324}
325
326impl WireTable {
327    pub fn new(owner: &str) -> Result<Self, ServingError> {
328        if !harmos_domain::identity::segment(owner) {
329            return Err(ServingError::Catalog);
330        }
331        Ok(Self {
332            owner: owner.to_owned(),
333            services: BTreeMap::new(),
334            handlers: BTreeMap::new(),
335            retired: BTreeMap::new(),
336            prepared: BTreeMap::new(),
337            capacity: None,
338            closed: false,
339            ordinary_closed: false,
340            provided: Vec::new(),
341            #[cfg(feature = "streaming")]
342            stream_limits: None,
343        })
344    }
345
346    // The Carrier Supplies These Resources To Every Attempt's Execution Context.
347    pub fn provide(&mut self, resource: std::any::TypeId) {
348        self.provided.push(resource);
349    }
350
351    pub fn provided(&self) -> &[std::any::TypeId] {
352        &self.provided
353    }
354
355    pub fn capacity(&mut self, capacity: NonZeroUsize) {
356        self.capacity = Some(capacity);
357    }
358
359    #[cfg(feature = "streaming")]
360    pub fn stream_limits(&mut self, limits: WireStreamLimits) -> Result<(), ServingError> {
361        limits.validate()?;
362        self.stream_limits = Some(limits);
363        Ok(())
364    }
365
366    pub fn register(&mut self, batch: RegistrationBatch) -> Result<(), ServingError> {
367        if self.closed || self.ordinary_closed {
368            return Err(ServingError::Closing);
369        }
370        let mut staged = Self::new(&self.owner)?;
371        staged.capacity = self.capacity;
372        staged.provided.clone_from(&self.provided);
373        #[cfg(feature = "streaming")]
374        {
375            staged.stream_limits = self.stream_limits;
376        }
377        for factory in batch.operations {
378            factory.install(&mut staged)?;
379        }
380        for factory in batch.services {
381            factory.policy.validate().map_err(|_| ServingError::Catalog)?;
382            let id = QualifiedId::qualify(&self.owner, &factory.local_id).map_err(|_| ServingError::Catalog)?;
383            let id = id.as_str().to_owned();
384            if staged.handlers.contains_key(&id) || staged.services.insert(id, factory).is_some() {
385                return Err(ServingError::Catalog);
386            }
387        }
388        if staged.handlers.keys().chain(staged.services.keys()).any(|id| self.handlers.contains_key(id) || self.services.contains_key(id)) {
389            return Err(ServingError::Catalog);
390        }
391        self.handlers.append(&mut staged.handlers);
392        self.services.append(&mut staged.services);
393        Ok(())
394    }
395
396    pub async fn unregister(&mut self, id: &QualifiedId) -> Result<(), ServingError> {
397        let key = id.as_str();
398        if !self.handlers.contains_key(key) && !self.services.contains_key(key) {
399            return Err(ServingError::Unknown { operation: key.to_owned() });
400        }
401        if let Some(handler) = self.handlers.get_mut(key) {
402            handler.closing = true;
403        }
404        let slots: Vec<_> =
405            self.prepared.iter().filter(|(_, slot)| slot.operation == key).map(|(invocation, slot)| (*invocation, slot.clone())).collect();
406        for (_, slot) in &slots {
407            slot.cancel();
408        }
409        for (_, slot) in &slots {
410            slot.finalize().await?;
411        }
412        for (invocation, _) in slots {
413            self.prepared.remove(&invocation);
414        }
415        if let Some(handler) = self.handlers.remove(key)
416            && let Some(generation) = handler.generation
417        {
418            self.retired.insert(key.to_owned(), generation);
419        }
420        self.services.remove(key);
421        Ok(())
422    }
423
424    pub fn definitions(&self) -> Vec<OperationDefinition> {
425        self.handlers.values().filter(|handler| !handler.builtin).map(|handler| handler.definition.clone()).collect()
426    }
427
428    pub fn validate(&self, definitions: &[OperationDefinition]) -> Result<(), ServingError> {
429        if definitions.len() != self.handlers.values().filter(|handler| !handler.builtin).count() {
430            return Err(ServingError::Catalog);
431        }
432        for definition in definitions {
433            if self
434                .handlers
435                .get(QualifiedId::qualify(&self.owner, definition.id.as_str()).map_err(|_| ServingError::Catalog)?.as_str())
436                .is_none_or(|bound| bound.builtin || bound.definition != *definition)
437            {
438                return Err(ServingError::Catalog);
439            }
440        }
441        let unique: std::collections::BTreeSet<_> = definitions.iter().map(|definition| &definition.id).collect();
442        if unique.len() != definitions.len() {
443            return Err(ServingError::Catalog);
444        }
445        Ok(())
446    }
447
448    pub fn validate_callees(&self, assignments: &[u8]) -> Result<(), ServingError> {
449        let assignments: Vec<super::AssignedOperation> =
450            if assignments.is_empty() { Vec::new() } else { postcard::from_bytes(assignments).map_err(|_| ServingError::Catalog)? };
451        for route in self.handlers.values() {
452            for callee in route.callees {
453                let id = callee.operation.id;
454                let fingerprint = (callee.operation.fingerprint)().map_err(|_| ServingError::Catalog)?;
455                let local = QualifiedId::qualify(&self.owner, id).map_err(|_| ServingError::Catalog)?;
456                let actual = self
457                    .handlers
458                    .get(local.as_str())
459                    .map(|entry| entry.definition.fingerprint.0)
460                    .or_else(|| {
461                        assignments
462                            .iter()
463                            .find(|entry| entry.operation.split_once('.').is_some_and(|(_, name)| name == id))
464                            .map(|entry| entry.fingerprint)
465                    })
466                    .ok_or(ServingError::Unassigned { operation: id })?;
467                if actual != fingerprint.0 {
468                    return Err(ServingError::Identity { operation: id.to_owned() });
469                }
470            }
471        }
472        Ok(())
473    }
474
475    pub fn validate_contracts(&self, contracts: &[Contract]) -> Result<(), ServingError> {
476        if contracts.len() != self.handlers.values().filter(|handler| !handler.builtin).count() {
477            return Err(ServingError::Catalog);
478        }
479        let mut unique = std::collections::BTreeSet::new();
480        for contract in contracts {
481            let id = &contract.operation().id;
482            if !unique.insert(id)
483                || self
484                    .handlers
485                    .get(QualifiedId::qualify(&self.owner, id.as_str()).map_err(|_| ServingError::Catalog)?.as_str())
486                    .is_none_or(|bound| bound.builtin || bound.contract.as_ref() != Some(contract))
487            {
488                return Err(ServingError::Catalog);
489            }
490        }
491        Ok(())
492    }
493
494    // Answers the request's grant as the Prepare payload.
495    pub fn prepare(
496        &mut self, invocation: u64, operation: &str, fingerprint: [u8; 32], generation: u64, payload: Vec<u8>,
497    ) -> Result<Vec<u8>, ServingError> {
498        self.prepare_codec(invocation, operation, fingerprint, generation, payload, false)
499    }
500
501    fn prepare_codec(
502        &mut self, invocation: u64, operation: &str, fingerprint: [u8; 32], generation: u64, payload: Vec<u8>, grpc: bool,
503    ) -> Result<Vec<u8>, ServingError> {
504        if self.closed {
505            return Err(ServingError::Closing);
506        }
507        self.prepared.retain(|_, slot| !slot.is_finished());
508        if self.prepared.contains_key(&invocation) {
509            return Err(ServingError::Duplicate { invocation });
510        }
511        let handler = self.handlers.get_mut(operation).ok_or_else(|| ServingError::Unknown { operation: operation.to_owned() })?;
512        if handler.definition.fingerprint.0 != fingerprint
513            || generation == 0
514            || handler.generation.is_some_and(|bound| bound != generation)
515            || self.retired.get(operation).is_some_and(|retired| generation <= *retired)
516        {
517            return Err(ServingError::Identity { operation: operation.to_owned() });
518        }
519        // The authenticated host assigns a generation with its first preparation.
520        // Keep it through withdrawal so delayed preparations cannot bind a replacement.
521        handler.generation = Some(generation);
522        if handler.closing || (self.ordinary_closed && !handler.lifecycle) {
523            return Err(ServingError::Closing);
524        }
525        let occupied = self.prepared.values().filter(|slot| slot.lifecycle == handler.lifecycle && !slot.resident).count();
526        // One accounted reserved lifecycle slot remains available even when
527        // ordinary admission is full. Failed SDK finalization retains that slot.
528        let capacity = if handler.lifecycle { 1 } else { self.capacity.ok_or(ServingError::Capacity)?.get() };
529        if occupied >= capacity && !handler.resident {
530            return Err(ServingError::Capacity);
531        }
532        let prepared = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| (handler.prepare)(payload, grpc)))
533            .map_err(|_| ServingError::Panicked)??;
534        let answer = encode_prepared_grant(prepared.grant().as_ref());
535        let slot = Slot {
536            operation: operation.to_owned(),
537            lifecycle: handler.lifecycle,
538            resident: handler.resident,
539            #[cfg(feature = "streaming")]
540            streams: prepared.streams(),
541            state: Mutex::new(SlotState {
542                prepared: Some(prepared),
543                finished: false,
544                cancelled: false,
545                waiters: Vec::new(),
546            }),
547        };
548        self.prepared.insert(invocation, Arc::new(slot));
549        Ok(answer)
550    }
551
552    #[cfg(feature = "streaming")]
553    pub fn streams(&self, invocation: u64) -> Result<WireStreams, ServingError> {
554        let slot = self.prepared.get(&invocation).ok_or(ServingError::Missing { invocation })?;
555        if self.closed || (self.ordinary_closed && !slot.lifecycle) || slot.is_finished() {
556            return Err(ServingError::Closing);
557        }
558        slot.streams.clone().ok_or(ServingError::StreamPhase)
559    }
560
561    #[tracing::instrument(target = "harmos::streams", name = "streams.initialize", level = "trace", skip_all, fields(invocation))]
562    pub async fn bind(&mut self, invocation: u64, context: &ExecutionContext) -> Result<(), ServingError> {
563        if self.closed {
564            return Err(ServingError::Closing);
565        }
566        let slot = self.prepared.get(&invocation).ok_or(ServingError::Missing { invocation })?;
567        if self.ordinary_closed && !slot.lifecycle {
568            return Err(ServingError::Closing);
569        }
570        slot.bind(invocation, context).await
571    }
572
573    #[tracing::instrument(target = "harmos::streams", name = "streams.attempt", level = "trace", skip_all, fields(invocation))]
574    pub async fn attempt(&mut self, invocation: u64, context: &ExecutionContext) -> Result<Vec<u8>, ServingError> {
575        if self.closed {
576            return Err(ServingError::Closing);
577        }
578        let slot = self.prepared.get(&invocation).ok_or(ServingError::Missing { invocation })?;
579        if self.ordinary_closed && !slot.lifecycle {
580            return Err(ServingError::Closing);
581        }
582        slot.attempt(invocation, context, None).await
583    }
584
585    #[tracing::instrument(target = "harmos::streams", name = "streams.finalize", level = "trace", skip_all, fields(invocation))]
586    pub async fn finalize(&mut self, invocation: u64) -> Result<(), ServingError> {
587        if let Some(slot) = self.prepared.get(&invocation) {
588            slot.finalize().await?;
589        }
590        self.prepared.remove(&invocation);
591        Ok(())
592    }
593
594    fn traffic(&self, command: WireCommand) -> Pin<Box<dyn Future<Output = WireAnswer> + Send + 'static>> {
595        #[cfg(not(feature = "streaming"))]
596        {
597            let _ = command;
598            Box::pin(async { wire_answer(Err(ServingError::StreamUnsupported)) })
599        }
600        #[cfg(feature = "streaming")]
601        {
602            let invocation = match &command {
603                WireCommand::StreamSend { invocation, .. }
604                | WireCommand::StreamFinish { invocation, .. }
605                | WireCommand::StreamReceive { invocation, .. } => *invocation,
606                _ => unreachable!("only stream traffic bypasses the phase loan"),
607            };
608            let streams = self.streams(invocation);
609            Box::pin(async move {
610                let streams = match streams {
611                    Ok(streams) => streams,
612                    Err(error) => return wire_answer(Err(error)),
613                };
614                match command {
615                    WireCommand::StreamSend { epoch, payload, .. } => {
616                        wire_answer(guarded(async { streams.send(epoch, &payload).await.map(|()| Vec::new()) }).await)
617                    }
618                    WireCommand::StreamFinish { epoch, .. } => wire_answer(streams.finish_input(epoch).map(|()| Vec::new())),
619                    WireCommand::StreamReceive { epoch, .. } => match guarded(streams.receive(epoch)).await {
620                        Ok(payload) => WireAnswer::Stream {
621                            epoch,
622                            transferred: streams.transferred(),
623                            payload,
624                        },
625                        Err(error) => wire_answer(Err(error)),
626                    },
627                    _ => unreachable!("only stream traffic bypasses the phase loan"),
628                }
629            })
630        }
631    }
632
633    pub fn begin(
634        &mut self, command: WireCommand, context: Arc<ExecutionContext>,
635    ) -> Pin<Box<dyn Future<Output = WireAnswer> + Send + 'static>> {
636        let (invocation, phase) = command_identity(&command);
637        let span = tracing::debug_span!(target: "harmos::carrier", "carrier.receive", invocation, phase, outcome = tracing::field::Empty);
638        if command.is_stream_traffic() {
639            return Box::pin(self.traffic(command).instrument(span).in_current_context());
640        }
641        let slot = match &command {
642            WireCommand::Initialize { invocation, .. }
643            | WireCommand::Attempt { invocation, .. }
644            | WireCommand::StreamAttempt { invocation, .. }
645            | WireCommand::Finalize { invocation } => self.prepared.get(invocation).cloned(),
646            WireCommand::Prepare { .. } | WireCommand::PrepareGrpc { .. } => None,
647            WireCommand::StreamSend { .. } | WireCommand::StreamFinish { .. } | WireCommand::StreamReceive { .. } => {
648                unreachable!("stream traffic already dispatched")
649            }
650        };
651        let command = match command {
652            WireCommand::PrepareGrpc {
653                invocation,
654                operation,
655                fingerprint,
656                generation,
657                payload,
658            } => {
659                let result = span.in_scope(|| self.prepare_codec(invocation, &operation, fingerprint, generation, payload, true));
660                return Box::pin(async move { wire_answer(result) }.instrument(span).in_current_context());
661            }
662            WireCommand::Prepare {
663                invocation,
664                operation,
665                fingerprint,
666                generation,
667                payload,
668            } => {
669                let result = span.in_scope(|| self.prepare(invocation, &operation, fingerprint, generation, payload));
670                return Box::pin(async move { wire_answer(result) }.instrument(span).in_current_context());
671            }
672            command => command,
673        };
674        let refused = self.closed || (self.ordinary_closed && slot.as_ref().is_some_and(|slot| !slot.lifecycle));
675        if refused && matches!(command, WireCommand::Initialize { .. } | WireCommand::Attempt { .. } | WireCommand::StreamAttempt { .. }) {
676            return Box::pin(async { wire_answer(Err(ServingError::Closing)) }.instrument(span).in_current_context());
677        }
678        Box::pin(
679            async move {
680                let result = match command {
681                    WireCommand::Initialize { invocation, .. } => match slot {
682                        Some(slot) => slot.bind(invocation, &context).await.map(|()| Vec::new()),
683                        None => Err(ServingError::Missing { invocation }),
684                    },
685                    WireCommand::Attempt { invocation, .. } => match slot {
686                        Some(slot) => slot.attempt(invocation, &context, None).await,
687                        None => Err(ServingError::Missing { invocation }),
688                    },
689                    WireCommand::StreamAttempt { invocation, epoch, .. } => match slot {
690                        Some(slot) => slot.attempt(invocation, &context, Some(epoch)).await,
691                        None => Err(ServingError::Missing { invocation }),
692                    },
693                    WireCommand::StreamSend { .. } | WireCommand::StreamFinish { .. } | WireCommand::StreamReceive { .. } => {
694                        unreachable!("stream traffic already dispatched")
695                    }
696                    WireCommand::Finalize { .. } => match slot {
697                        Some(slot) => slot.finalize().await.map(|()| Vec::new()),
698                        None => Ok(Vec::new()),
699                    },
700                    WireCommand::Prepare { .. } | WireCommand::PrepareGrpc { .. } => {
701                        unreachable!("preparation is synchronous")
702                    }
703                };
704                wire_answer(result)
705            }
706            .instrument(span)
707            .in_current_context(),
708        )
709    }
710
711    pub fn close_ordinary(&mut self) -> Pin<Box<dyn Future<Output = Result<(), ServingError>> + Send + 'static>> {
712        self.ordinary_closed = true;
713        let slots: Vec<_> = self.prepared.values().filter(|slot| !slot.lifecycle).cloned().collect();
714        for slot in &slots {
715            slot.cancel();
716        }
717        Box::pin(async move {
718            for slot in slots {
719                slot.finalize().await?;
720            }
721            Ok(())
722        })
723    }
724
725    pub fn release_ordinary(&mut self) -> Result<impl Send + use<>, ServingError> {
726        if !self.ordinary_closed || self.prepared.values().any(|slot| !slot.lifecycle && !slot.is_finished()) {
727            return Err(ServingError::Unsettled);
728        }
729        // Hand Off Retired Owners So The Caller Can Drop Them Outside Its Lock
730        let retired_slots = self.prepared.extract_if(.., |_, slot| !slot.lifecycle).collect::<Vec<_>>();
731        let retired_handlers = self.handlers.extract_if(.., |_, handler| !handler.lifecycle).collect::<Vec<_>>();
732        let retired_services = std::mem::take(&mut self.services);
733        Ok((retired_slots, retired_handlers, retired_services))
734    }
735
736    pub async fn settle_ordinary(&mut self) -> Result<(), ServingError> {
737        self.close_ordinary().await?;
738        let retired = self.release_ordinary()?;
739        drop(retired);
740        Ok(())
741    }
742
743    pub fn close(&mut self) -> Pin<Box<dyn Future<Output = Result<(), ServingError>> + Send + 'static>> {
744        self.closed = true;
745        let slots: Vec<_> = self.prepared.values().cloned().collect();
746        for slot in &slots {
747            slot.cancel();
748        }
749        Box::pin(async move {
750            for slot in slots {
751                slot.finalize().await?;
752            }
753            Ok(())
754        })
755    }
756
757    pub fn release(&mut self) -> Result<impl Send + use<>, ServingError> {
758        if !self.closed || self.prepared.values().any(|slot| !slot.is_finished()) {
759            return Err(ServingError::Unsettled);
760        }
761        Ok((std::mem::take(&mut self.prepared), std::mem::take(&mut self.handlers), std::mem::take(&mut self.services)))
762    }
763
764    pub async fn settle(&mut self) -> Result<(), ServingError> {
765        self.close().await?;
766        let retired = self.release()?;
767        drop(retired);
768        Ok(())
769    }
770
771    pub fn is_settled(&self) -> bool {
772        self.closed && self.prepared.values().all(|slot| slot.is_finished())
773    }
774
775    pub fn is_closed(&self) -> bool {
776        self.closed
777    }
778
779    #[tracing::instrument(target = "harmos::carrier", name = "carrier.answer", level = "debug", skip_all, fields(invocation = command_identity(&command).0, phase = command_identity(&command).1))]
780    pub async fn answer(&mut self, command: WireCommand, context: &ExecutionContext) -> WireAnswer {
781        if command.is_stream_traffic() {
782            return self.traffic(command).await;
783        }
784        if let WireCommand::StreamAttempt { invocation, epoch, .. } = command {
785            if self.closed || (self.ordinary_closed && self.prepared.get(&invocation).is_some_and(|slot| !slot.lifecycle)) {
786                return wire_answer(Err(ServingError::Closing));
787            }
788            let Some(slot) = self.prepared.get(&invocation) else {
789                return wire_answer(Err(ServingError::Missing { invocation }));
790            };
791            return wire_answer(slot.attempt(invocation, context, Some(epoch)).await);
792        }
793        let result = match command {
794            WireCommand::PrepareGrpc {
795                invocation,
796                operation,
797                fingerprint,
798                generation,
799                payload,
800            } => self.prepare_codec(invocation, &operation, fingerprint, generation, payload, true),
801            WireCommand::Prepare {
802                invocation,
803                operation,
804                fingerprint,
805                generation,
806                payload,
807            } => self.prepare(invocation, &operation, fingerprint, generation, payload),
808            WireCommand::Initialize { invocation, .. } => self.bind(invocation, context).await.map(|()| Vec::new()),
809            WireCommand::Attempt { invocation, .. } => self.attempt(invocation, context).await,
810            WireCommand::StreamAttempt { .. }
811            | WireCommand::StreamSend { .. }
812            | WireCommand::StreamFinish { .. }
813            | WireCommand::StreamReceive { .. } => {
814                unreachable!("stream command already dispatched")
815            }
816            WireCommand::Finalize { invocation } => self.finalize(invocation).await.map(|()| Vec::new()),
817        };
818        wire_answer(result)
819    }
820}
821
822impl Slot {
823    fn cancel(&self) {
824        let waiters = {
825            let mut state = self.state.lock().expect("invocation state lock");
826            state.cancelled = true;
827            std::mem::take(&mut state.waiters)
828        };
829        #[cfg(feature = "streaming")]
830        if let Some(streams) = &self.streams {
831            streams.close();
832        }
833        for waiter in waiters {
834            if let Some(waiter) = waiter.upgrade() {
835                let wake = waiter.lock().expect("invocation waiter lock").take();
836                if let Some(wake) = wake {
837                    wake.wake();
838                }
839            }
840        }
841    }
842
843    async fn active<T>(&self, future: impl Future<Output = Result<T, ServingError>>) -> Result<T, ServingError> {
844        let mut future = std::pin::pin!(future);
845        let waiter = Arc::new(Mutex::new(None));
846        std::future::poll_fn(|context| {
847            {
848                let mut state = self.state.lock().expect("invocation state lock");
849                if state.cancelled {
850                    return Poll::Ready(Err(ServingError::Closing));
851                }
852                *waiter.lock().expect("invocation waiter lock") = Some(context.waker().clone());
853                state.waiters.retain(|entry| entry.strong_count() != 0);
854                if !state.waiters.iter().any(|entry| entry.ptr_eq(&Arc::downgrade(&waiter))) {
855                    state.waiters.push(Arc::downgrade(&waiter));
856                }
857            }
858            future.as_mut().poll(context)
859        })
860        .await
861    }
862
863    fn is_finished(&self) -> bool {
864        self.state.lock().expect("invocation state lock").finished
865    }
866
867    async fn borrow(self: &Arc<Self>) -> Option<Loan> {
868        let waiter = Arc::new(Mutex::new(None));
869        std::future::poll_fn(|context| {
870            let mut state = self.state.lock().expect("invocation state lock");
871            if state.finished {
872                return Poll::Ready(None);
873            }
874            if let Some(prepared) = state.prepared.take() {
875                return Poll::Ready(Some(Loan {
876                    slot: Arc::clone(self),
877                    prepared: Some(prepared),
878                    finished: false,
879                }));
880            }
881            *waiter.lock().expect("invocation waiter lock") = Some(context.waker().clone());
882            state.waiters.retain(|entry| entry.strong_count() != 0);
883            if !state.waiters.iter().any(|entry| entry.ptr_eq(&Arc::downgrade(&waiter))) {
884                state.waiters.push(Arc::downgrade(&waiter));
885            }
886            Poll::Pending
887        })
888        .await
889    }
890}
891
892impl Slot {
893    #[tracing::instrument(target = "harmos::streams", name = "streams.initialize", level = "trace", skip_all, fields(invocation))]
894    async fn bind(self: &Arc<Self>, invocation: u64, context: &ExecutionContext) -> Result<(), ServingError> {
895        let mut loan = self.borrow().await.ok_or(ServingError::Missing { invocation })?;
896        self.active(loan.prepared.as_mut().expect("loaned invocation state").bind(context)).await
897    }
898
899    #[tracing::instrument(target = "harmos::streams", name = "streams.attempt", level = "trace", skip_all, fields(invocation))]
900    async fn attempt(self: &Arc<Self>, invocation: u64, context: &ExecutionContext, epoch: Option<u64>) -> Result<Vec<u8>, ServingError> {
901        let mut loan = self.borrow().await.ok_or(ServingError::Missing { invocation })?;
902        #[cfg(feature = "streaming")]
903        let attempt = match (&self.streams, epoch) {
904            (Some(streams), Some(epoch)) => Some(streams.start(epoch)?),
905            (None, None) => None,
906            _ => return Err(ServingError::StreamUnsupported),
907        };
908        #[cfg(not(feature = "streaming"))]
909        if epoch.is_some() {
910            return Err(ServingError::StreamUnsupported);
911        }
912        let result = self.active(loan.prepared.as_mut().expect("loaned invocation state").attempt(context)).await;
913        #[cfg(feature = "streaming")]
914        if let Some(attempt) = attempt {
915            attempt.end();
916        }
917        result
918    }
919
920    #[tracing::instrument(target = "harmos::carrier", name = "carrier.finalize", level = "debug", skip_all)]
921    async fn finalize(self: &Arc<Self>) -> Result<(), ServingError> {
922        #[cfg(feature = "streaming")]
923        if let Some(streams) = &self.streams {
924            streams.close();
925        }
926        let Some(mut loan) = self.borrow().await else {
927            return Ok(());
928        };
929        loan.prepared.as_mut().expect("loaned invocation state").finalize().await?;
930        // Drop Physical State Before Releasing Its Capacity
931        loan.finished = true;
932        let prepared = loan.prepared.take();
933        drop(prepared);
934        Ok(())
935    }
936}
937
938impl Drop for Loan {
939    fn drop(&mut self) {
940        let waiters = {
941            let mut state = self.slot.state.lock().expect("invocation state lock");
942            state.prepared = self.prepared.take();
943            state.finished = self.finished;
944            std::mem::take(&mut state.waiters)
945        };
946        for waiter in waiters {
947            if let Some(waiter) = waiter.upgrade() {
948                let wake = waiter.lock().expect("invocation waiter lock").take();
949                if let Some(wake) = wake {
950                    wake.wake();
951                }
952            }
953        }
954    }
955}
956
957impl WireCommand {
958    pub const LIMIT: usize = codec::LIMIT;
959    fn is_stream_traffic(&self) -> bool {
960        matches!(self, Self::StreamSend { .. } | Self::StreamFinish { .. } | Self::StreamReceive { .. })
961    }
962}
963
964impl ServingError {
965    pub fn classification(&self) -> FailureKind {
966        match self {
967            Self::Handler(failure) => failure.code(),
968            Self::Panicked => FailureKind::Ambiguous,
969            Self::Codec(_) | Self::Declaration(_) | Self::Lifecycle(_) | Self::Contract(_) | Self::Catalog | Self::Streaming { .. } => {
970                FailureKind::Invalid
971            }
972            _ => FailureKind::Permanent,
973        }
974    }
975}
976
977/* === Public Functions === */
978
979// A Prepare answer is empty when the request needs no grant, else the grant as JSON.
980pub fn encode_prepared_grant(grant: Option<&crate::plugin::Grant>) -> Vec<u8> {
981    grant.map_or_else(Vec::new, |grant| serde_json::to_vec(grant).expect("a grant serializes to JSON"))
982}
983
984pub fn decode_prepared_grant(payload: &[u8]) -> Result<Option<crate::plugin::Grant>, serde_json::Error> {
985    if payload.is_empty() {
986        return Ok(None);
987    }
988    serde_json::from_slice(payload).map(Some)
989}
990
991/* === Private Functions === */
992
993fn command_identity(command: &WireCommand) -> (u64, &'static str) {
994    match command {
995        WireCommand::Prepare { invocation, .. } => (*invocation, "prepare"),
996        WireCommand::PrepareGrpc { invocation, .. } => (*invocation, "prepare_grpc"),
997        WireCommand::Initialize { invocation, .. } => (*invocation, "initialize"),
998        WireCommand::Attempt { invocation, .. } => (*invocation, "attempt"),
999        WireCommand::Finalize { invocation, .. } => (*invocation, "finalize"),
1000        WireCommand::StreamAttempt { invocation, .. } => (*invocation, "stream_attempt"),
1001        WireCommand::StreamSend { invocation, .. } => (*invocation, "stream_send"),
1002        WireCommand::StreamFinish { invocation, .. } => (*invocation, "stream_finish"),
1003        WireCommand::StreamReceive { invocation, .. } => (*invocation, "stream_receive"),
1004    }
1005}
1006
1007fn wire_answer(result: Result<Vec<u8>, ServingError>) -> WireAnswer {
1008    tracing::Span::current().record("outcome", if result.is_ok() { "ready" } else { "failed" });
1009    match result {
1010        Ok(payload) => WireAnswer::Ready { payload },
1011        Err(ServingError::Handler(failure)) => WireAnswer::Failed { failure },
1012        Err(error) => WireAnswer::Failed {
1013            failure: crate::Failure::originate(error.classification(), &error, crate::failure::Origin::Plugin),
1014        },
1015    }
1016}
1017
1018async fn guarded<T>(future: impl Future<Output = Result<T, ServingError>>) -> Result<T, ServingError> {
1019    let mut future = std::pin::pin!(future);
1020    std::future::poll_fn(|context| match std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| future.as_mut().poll(context))) {
1021        Ok(result) => result,
1022        Err(_) => std::task::Poll::Ready(Err(ServingError::Panicked)),
1023    })
1024    .await
1025}
1026
1027/* === Tests === */
1028
1029#[cfg(test)]
1030mod tests {
1031    use super::*;
1032
1033    #[test]
1034    fn withdrawal_cancels_a_pending_active_future() {
1035        let slot = Slot {
1036            operation: "host.read".to_owned(),
1037            lifecycle: false,
1038            resident: false,
1039            #[cfg(feature = "streaming")]
1040            streams: None,
1041            state: Mutex::new(SlotState {
1042                prepared: None,
1043                finished: false,
1044                cancelled: false,
1045                waiters: Vec::new(),
1046            }),
1047        };
1048        let mut active = Box::pin(slot.active(std::future::pending::<Result<(), ServingError>>()));
1049        let mut context = std::task::Context::from_waker(Waker::noop());
1050        assert!(active.as_mut().poll(&mut context).is_pending());
1051        slot.cancel();
1052        assert!(matches!(active.as_mut().poll(&mut context), Poll::Ready(Err(ServingError::Closing))));
1053    }
1054}