buntrail_engine/remote/
transport.rs

1use super::*;
2use buntrail_protocol::{
3    frame::{Decoder, Frame},
4    types::{Context, Header, Message},
5};
6use std::{
7    collections::{HashMap, VecDeque},
8    io::{Read, Write},
9    os::unix::net::UnixStream,
10    sync::{
11        Condvar, Mutex, Weak,
12        mpsc::{self, SyncSender},
13    },
14    task::{Poll, Waker},
15    time::Duration,
16};
17struct Inbox {
18    frames: VecDeque<Frame>,
19    bytes: usize,
20    error: Option<CoreError>,
21    waker: Option<Waker>,
22}
23struct Mailbox {
24    cancelled: AtomicBool,
25    finished: AtomicBool,
26    unsent_refusal: AtomicBool,
27    state: Mutex<Inbox>,
28    ready: Condvar,
29}
30impl Mailbox {
31    fn new() -> Arc<Self> {
32        Arc::new(Self {
33            cancelled: AtomicBool::new(false),
34            finished: AtomicBool::new(false),
35            unsent_refusal: AtomicBool::new(false),
36            state: Mutex::new(Inbox {
37                frames: VecDeque::new(),
38                bytes: 0,
39                error: None,
40                waker: None,
41            }),
42            ready: Condvar::new(),
43        })
44    }
45    fn fail(&self, e: CoreError) {
46        let mut s = self.state.lock().unwrap();
47        s.error = Some(e);
48        if let Some(w) = s.waker.take() {
49            w.wake()
50        }
51        self.ready.notify_all();
52    }
53    fn refuse_unsent(&self, error: CoreError) {
54        // Only the worker can establish that this owner was removed without
55        // entering the socket output. A deadline alone cannot prove that.
56        self.unsent_refusal.store(true, Ordering::SeqCst);
57        self.fail(error);
58    }
59    fn push(&self, f: Frame) -> Result<(), CoreError> {
60        let mut s = self.state.lock().unwrap();
61        if s.bytes + f.payload.len() > 1024 * 1024 || s.frames.len() >= 10000 {
62            return Err(CoreError::new(
63                ErrorCode::Busy,
64                "Remote receive window exceeded",
65            ));
66        }
67        s.bytes += f.payload.len();
68        s.frames.push_back(f);
69        if let Some(w) = s.waker.take() {
70            w.wake()
71        }
72        self.ready.notify_all();
73        Ok(())
74    }
75    fn take(s: &mut Inbox) -> Option<Result<Frame, CoreError>> {
76        if let Some(f) = s.frames.pop_front() {
77            s.bytes -= f.payload.len();
78            Some(Ok(f))
79        } else {
80            s.error.clone().map(Err)
81        }
82    }
83}
84type Credits = Arc<Mutex<HashMap<u64, (Header, usize, Arc<Scope>)>>>;
85struct SendFrame {
86    frame: Frame,
87    mailbox: Option<Arc<Mailbox>>,
88    scope: Arc<Scope>,
89}
90#[derive(Clone)]
91pub(super) struct Transport {
92    sender: SyncSender<SendFrame>,
93    credits: Credits,
94    control: SyncSender<SendFrame>,
95    pub context: Context,
96    scope: Arc<Scope>,
97    scopes: Arc<Mutex<Vec<Weak<Scope>>>>,
98    registrations: Arc<Mutex<()>>,
99    pub capabilities: Arc<Vec<String>>,
100    retired: Arc<Retirement>,
101    pub runtime: RuntimeGeneration,
102    pub connected: Arc<AtomicBool>,
103    ids: Arc<std::sync::atomic::AtomicU64>,
104    pub(super) authentication_ready: Arc<AtomicBool>,
105    pub(super) physical: Option<Arc<crate::environment::PhysicalOwner>>,
106    pub(super) connection_permit: Option<Arc<crate::environment::ConnectionPermit>>,
107    requests: Arc<std::sync::atomic::AtomicUsize>,
108}
109struct Scope {
110    context: Context,
111    lease: WorkspaceLease,
112    active: AtomicBool,
113    release_confirmed: Mutex<Option<buntrail_protocol::types::ReleaseIntent>>,
114}
115#[derive(Default)]
116struct Retirement {
117    done: AtomicBool,
118    wake: Mutex<Option<Waker>>,
119}
120struct Retire(Arc<Retirement>, Arc<AtomicBool>);
121impl Drop for Retire {
122    fn drop(&mut self) {
123        self.1.store(false, Ordering::SeqCst);
124        self.0.done.store(true, Ordering::SeqCst);
125        if let Some(wake) = self.0.wake.lock().unwrap().take() {
126            wake.wake();
127        }
128    }
129}
130pub struct Shutdown(Arc<Retirement>);
131impl std::future::Future for Shutdown {
132    type Output = ();
133    fn poll(self: Pin<&mut Self>, cx: &mut std::task::Context<'_>) -> Poll<()> {
134        let mut wake = self.0.wake.lock().unwrap();
135        if self.0.done.load(Ordering::SeqCst) {
136            Poll::Ready(())
137        } else {
138            *wake = Some(cx.waker().clone());
139            Poll::Pending
140        }
141    }
142}
143pub(super) struct Request {
144    pub header: Header,
145    mailbox: Arc<Mailbox>,
146    transport: Transport,
147    finished: bool,
148}
149impl Drop for Request {
150    fn drop(&mut self) {
151        self.transport.requests.fetch_sub(1, Ordering::SeqCst);
152        if !self.finished {
153            self.mailbox.cancelled.store(true, Ordering::SeqCst);
154        }
155    }
156}
157impl Request {
158    pub fn send_control(&self, message: Message, payload: Vec<u8>) -> Result<(), CoreError> {
159        let mut header = self.header.clone();
160        header.message = message;
161        self.transport
162            .control
163            .try_send(SendFrame {
164                frame: Frame { header, payload },
165                mailbox: None,
166                scope: self.transport.scope.clone(),
167            })
168            .map_err(|_| {
169                CoreError::new(
170                    ErrorCode::Unknown,
171                    "original transfer control queue unavailable; inspect without replay",
172                )
173            })
174    }
175
176    pub fn send(&self, message: Message, payload: Vec<u8>) -> Result<(), CoreError> {
177        let mut header = self.header.clone();
178        header.message = message;
179        self.transport
180            .sender
181            .send(SendFrame {
182                frame: Frame { header, payload },
183                mailbox: None,
184                scope: self.transport.scope.clone(),
185            })
186            .map_err(|_| closed())
187    }
188    pub fn credit(&self, stream: &str, bytes: usize) -> Result<(), CoreError> {
189        let mut credits = self.transport.credits.lock().unwrap();
190        if credits.len() >= 64 && !credits.contains_key(&self.header.request_id) {
191            return Err(CoreError::new(
192                ErrorCode::Busy,
193                "Remote credit owner capacity reached",
194            ));
195        }
196        let mut header = self.header.clone();
197        header.stream_id = Some(stream.into());
198        let entry =
199            credits
200                .entry(header.request_id)
201                .or_insert((header, 0, self.transport.scope.clone()));
202        if entry.0.stream_id.as_deref() != Some(stream) || entry.1 + bytes > 524288 {
203            return Err(w::invalid());
204        }
205        entry.1 += bytes;
206        Ok(())
207    }
208    pub fn next(&mut self) -> Result<Frame, CoreError> {
209        self.next_with_timeout(Duration::from_secs(120))
210    }
211    fn next_with_timeout(&mut self, timeout: Duration) -> Result<Frame, CoreError> {
212        let mut s = self.mailbox.state.lock().unwrap();
213        loop {
214            if let Some(f) = Mailbox::take(&mut s) {
215                return f;
216            }
217            let (state, timeout) = self.mailbox.ready.wait_timeout(s, timeout).unwrap();
218            s = state;
219            if timeout.timed_out() {
220                return Err(CoreError::new(
221                    ErrorCode::Unknown,
222                    "Remote request deadline exceeded; inspect original operation without replay",
223                ));
224            }
225        }
226    }
227    pub fn poll(&mut self, cx: &mut std::task::Context<'_>) -> Poll<Result<Frame, CoreError>> {
228        let mut s = self.mailbox.state.lock().unwrap();
229        if let Some(f) = Mailbox::take(&mut s) {
230            Poll::Ready(f)
231        } else {
232            s.waker = Some(cx.waker().clone());
233            Poll::Pending
234        }
235    }
236    pub fn finish(&mut self) {
237        self.finished = true;
238        self.mailbox.finished.store(true, Ordering::SeqCst);
239    }
240    pub fn connection_lost(&self) -> bool {
241        !self.transport.connected.load(Ordering::SeqCst)
242    }
243    pub fn refused_before_send(&self) -> bool {
244        self.mailbox.unsent_refusal.load(Ordering::SeqCst)
245    }
246}
247fn closed() -> CoreError {
248    CoreError::new(
249        ErrorCode::Closed,
250        "Remote connection lost; inspect accepted operations, never replay",
251    )
252}
253impl Transport {
254    pub fn connect(
255        mut socket: UnixStream,
256        target: &TargetIdentity,
257        lease: &WorkspaceLease,
258    ) -> Result<Self, CoreError> {
259        socket
260            .set_read_timeout(Some(Duration::from_secs(5)))
261            .map_err(io_error)?;
262        socket
263            .set_write_timeout(Some(Duration::from_secs(5)))
264            .map_err(io_error)?;
265        let mut h = Header {
266            major: 1,
267            minor: 0,
268            request_id: 1,
269            context: None,
270            operation_id: None,
271            stream_id: None,
272            sequence: None,
273            message: Message::Hello {
274                capabilities: buntrail_protocol::negotiation::CAPABILITIES
275                    .iter()
276                    .map(|s| s.to_string())
277                    .collect(),
278                max_header: 65536,
279                max_payload: 1048576,
280            },
281        };
282        write_frame(
283            &mut socket,
284            &Frame {
285                header: h.clone(),
286                payload: vec![],
287            },
288        )?;
289        let welcome = read_frame(&mut socket, None)?;
290        if welcome.header.request_id != 1 {
291            return Err(w::invalid());
292        }
293        let Message::Welcome {
294            target_binding,
295            runtime_generation,
296            container_start_base64,
297            capabilities,
298            ..
299        } = welcome.header.message
300        else {
301            return Err(w::invalid());
302        };
303        if !capabilities.iter().any(|s| s == "services") {
304            return Err(CoreError::new(
305                ErrorCode::Unavailable,
306                "Remote service capability unavailable",
307            ));
308        }
309        let expected = w::hex(&sha2::Sha256::digest(
310            serde_json::to_vec(target).map_err(io_error)?,
311        ));
312        if target_binding != expected {
313            return Err(CoreError::new(
314                ErrorCode::TargetChanged,
315                "Helper target binding changed",
316            ));
317        }
318        let runtime = RuntimeGeneration {
319            supervisor_boot: w::id(&runtime_generation)?,
320            container_start: buntrail_protocol::types::decode_container_start(
321                container_start_base64.as_deref(),
322            )
323            .map_err(|_| w::invalid())?,
324        };
325        if target.container.is_some() && runtime.container_start.is_none() {
326            return Err(CoreError::new(
327                ErrorCode::TargetChanged,
328                "Container helper Welcome is missing verified start evidence; no owner was registered",
329            ));
330        }
331        if runtime.container_start != lease.runtime.container_start {
332            return Err(CoreError::new(
333                ErrorCode::TargetChanged,
334                "Helper container start generation changed; preserve the original runtime and reconcile explicitly",
335            ));
336        }
337        if lease.runtime.supervisor_boot != [0; 16] && lease.runtime != runtime {
338            return Err(CoreError::new(
339                ErrorCode::TargetChanged,
340                "Helper runtime changed",
341            ));
342        }
343        let context = Context {
344            workspace_id: w::uuid(lease.workspace.0),
345            client_id: w::uuid(lease.client.0),
346            target_binding,
347            runtime_generation,
348        };
349        h.request_id = 2;
350        h.message = Message::Register {
351            lease_id: w::uuid(lease.lease.0),
352            client_id: context.client_id.clone(),
353            workspace_id: context.workspace_id.clone(),
354            target_binding: context.target_binding.clone(),
355            runtime_generation: context.runtime_generation.clone(),
356        };
357        write_frame(
358            &mut socket,
359            &Frame {
360                header: h,
361                payload: vec![],
362            },
363        )?;
364        let ack = read_frame(&mut socket, Some(&context))?;
365        if ack.header.request_id != 2 || !matches!(ack.header.message, Message::Ack {}) {
366            return Err(w::invalid());
367        }
368        socket.set_nonblocking(true).map_err(io_error)?;
369        let (tx, rx) = mpsc::sync_channel::<SendFrame>(32);
370        let (control, controls) = mpsc::sync_channel::<SendFrame>(32);
371        let connected = Arc::new(AtomicBool::new(true));
372        let live = connected.clone();
373        let scope = Arc::new(Scope {
374            context: context.clone(),
375            lease: WorkspaceLease {
376                runtime: runtime.clone(),
377                ..lease.clone()
378            },
379            active: AtomicBool::new(true),
380            release_confirmed: Mutex::new(None),
381        });
382        let scopes = Arc::new(Mutex::new(vec![Arc::downgrade(&scope)]));
383        let retired = Arc::new(Retirement::default());
384        let completion = Retire(retired.clone(), connected.clone());
385        let credits: Credits = Arc::new(Mutex::new(HashMap::new()));
386        let pending_credit = credits.clone();
387        let ids = Arc::new(std::sync::atomic::AtomicU64::new(3));
388        let worker_ids = ids.clone();
389        std::thread::Builder::new()
390            .name("buntrail-remote-wire".into())
391            .spawn(move || {
392                let _completion = completion;
393                worker(socket, rx, controls, pending_credit, worker_ids, live)
394            })
395            .map_err(io_error)?;
396        Ok(Self {
397            sender: tx,
398            credits,
399            control,
400            context,
401            scope,
402            scopes,
403            registrations: Arc::new(Mutex::new(())),
404            capabilities: Arc::new(capabilities),
405            retired,
406            runtime,
407            connected,
408            ids,
409            authentication_ready: Arc::new(AtomicBool::new(false)),
410            physical: None,
411            connection_permit: None,
412            requests: Arc::new(std::sync::atomic::AtomicUsize::new(0)),
413        })
414    }
415    pub fn register(&self, lease: WorkspaceLease) -> Result<Self, CoreError> {
416        let _registration = self.registrations.lock().map_err(io_error)?;
417        if !self.connected.load(Ordering::SeqCst) {
418            return Err(closed());
419        }
420        if lease.client != self.scope.lease.client || lease.runtime != self.runtime {
421            return Err(CoreError::new(
422                ErrorCode::TargetChanged,
423                "workspace client or runtime changed",
424            ));
425        }
426        let mut context = self.context.clone();
427        context.workspace_id = w::uuid(lease.workspace.0);
428        let scope = {
429            let mut scopes = self.scopes.lock().map_err(io_error)?;
430            scopes.retain(|s| s.strong_count() > 0);
431            if scopes
432                .iter()
433                .filter_map(Weak::upgrade)
434                .any(|s| s.lease.lease == lease.lease && s.lease != lease)
435            {
436                return Err(CoreError::new(
437                    ErrorCode::InvalidArgument,
438                    "lease already belongs to another workspace",
439                ));
440            }
441            if let Some(existing) = scopes
442                .iter()
443                .filter_map(Weak::upgrade)
444                .find(|s| s.lease == lease && s.active.load(Ordering::SeqCst))
445            {
446                existing
447            } else {
448                if scopes.len() >= 256 {
449                    return Err(CoreError::new(
450                        ErrorCode::Busy,
451                        "workspace view capacity reached; retire prior views",
452                    ));
453                }
454                let scope = Arc::new(Scope {
455                    context: context.clone(),
456                    lease: lease.clone(),
457                    active: AtomicBool::new(false),
458                    release_confirmed: Mutex::new(None),
459                });
460                scopes.push(Arc::downgrade(&scope));
461                scope
462            }
463        };
464        let view = Self {
465            context,
466            scope,
467            ..self.clone()
468        };
469        if view.scope.active.load(Ordering::SeqCst) {
470            return Ok(view);
471        }
472        let c = &view.context;
473        let mut request = view.start(
474            Message::Register {
475                lease_id: w::uuid(lease.lease.0),
476                client_id: c.client_id.clone(),
477                workspace_id: c.workspace_id.clone(),
478                target_binding: c.target_binding.clone(),
479                runtime_generation: c.runtime_generation.clone(),
480            },
481            None,
482            None,
483        )?;
484        if !matches!(request.next()?.header.message, Message::Ack {}) {
485            return Err(w::invalid());
486        }
487        request.finish();
488        view.scope.active.store(true, Ordering::SeqCst);
489        Ok(view)
490    }
491    pub fn release(
492        &self,
493        intent: buntrail_protocol::types::ReleaseIntent,
494    ) -> Result<(), CoreError> {
495        self.release_wait(intent, Duration::from_secs(120), false)
496    }
497    pub(super) fn detach(&self) -> Result<(), CoreError> {
498        self.release_wait(
499            buntrail_protocol::types::ReleaseIntent::ViewDetachment,
500            Duration::from_millis(500),
501            true,
502        )
503    }
504    fn release_wait(
505        &self,
506        intent: buntrail_protocol::types::ReleaseIntent,
507        timeout: Duration,
508        nonblocking: bool,
509    ) -> Result<(), CoreError> {
510        let _registration = if nonblocking {
511            self.registrations.try_lock().map_err(|_| {
512                self.scope.active.store(false, Ordering::SeqCst);
513                CoreError::new(
514                    ErrorCode::Unknown,
515                    "Local view detached; remote lease acknowledgement is unconfirmed",
516                )
517            })?
518        } else {
519            self.registrations.lock().map_err(io_error)?
520        };
521        if !self.scope.active.swap(false, Ordering::SeqCst) {
522            return if self
523                .scope
524                .release_confirmed
525                .lock()
526                .map_err(io_error)?
527                .as_ref()
528                == Some(&intent)
529            {
530                Ok(())
531            } else {
532                Err(CoreError::new(
533                    ErrorCode::Closed,
534                    "workspace release is unconfirmed; automatic replay is disabled",
535                ))
536            };
537        }
538        let mut request = self.start(
539            Message::Release {
540                lease_id: w::uuid(self.scope.lease.lease.0),
541                intent: intent.clone(),
542            },
543            None,
544            None,
545        )?;
546        if !matches!(
547            request.next_with_timeout(timeout)?.header.message,
548            Message::Ack {}
549        ) {
550            return Err(w::invalid());
551        }
552        request.finish();
553        *self.scope.release_confirmed.lock().map_err(io_error)? = Some(intent);
554        Ok(())
555    }
556    pub(super) fn owns_registration(&self, lease: &WorkspaceLease) -> bool {
557        self.registered() && &self.scope.lease == lease && lease.runtime == self.runtime
558    }
559    pub(super) fn registered(&self) -> bool {
560        self.scope.active.load(Ordering::SeqCst) && self.connected.load(Ordering::SeqCst)
561    }
562    pub(super) fn lifetime(&self) -> Lifetime {
563        Lifetime {
564            registrations: self.registrations.clone(),
565            connected: self.connected.clone(),
566            retired: self.retired.clone(),
567            scopes: self.scopes.clone(),
568            requests: self.requests.clone(),
569        }
570    }
571    pub fn close(&self) -> Shutdown {
572        self.connected.store(false, Ordering::SeqCst);
573        Shutdown(self.retired.clone())
574    }
575    fn frame(
576        &self,
577        message: Message,
578        operation: Option<OperationID>,
579        stream: Option<String>,
580    ) -> Frame {
581        Frame {
582            header: Header {
583                major: 1,
584                minor: 0,
585                request_id: self.ids.fetch_add(1, Ordering::SeqCst),
586                context: if matches!(message, Message::Register { .. }) {
587                    None
588                } else {
589                    Some(self.context.clone())
590                },
591                operation_id: operation.map(|o| w::uuid(o.0)),
592                stream_id: stream,
593                sequence: None,
594                message,
595            },
596            payload: vec![],
597        }
598    }
599    pub fn start(
600        &self,
601        message: Message,
602        operation: Option<OperationID>,
603        stream: Option<String>,
604    ) -> Result<Request, CoreError> {
605        self.start_payload(message, operation, stream, Vec::new())
606    }
607    pub fn start_payload(
608        &self,
609        message: Message,
610        operation: Option<OperationID>,
611        stream: Option<String>,
612        payload: Vec<u8>,
613    ) -> Result<Request, CoreError> {
614        if !self.connected.load(Ordering::SeqCst)
615            || (!self.scope.active.load(Ordering::SeqCst)
616                && !matches!(message, Message::Register { .. } | Message::Release { .. }))
617        {
618            return Err(closed());
619        }
620        let mut frame = self.frame(message, operation, stream);
621        frame.payload = payload;
622        // Legacy existing-only calls must fit the exact envelope before queueing any frame.
623        if matches!(
624            &frame.header.message,
625            Message::ServiceCall {
626                call: w::Call::GitExisting { .. }
627            }
628        ) {
629            struct HeaderBound(usize);
630            impl std::io::Write for HeaderBound {
631                fn write(&mut self, bytes: &[u8]) -> std::io::Result<usize> {
632                    if bytes.len() > 65536 - self.0 {
633                        return Err(std::io::Error::other(
634                            "legacy original exceeds header limit",
635                        ));
636                    }
637                    self.0 += bytes.len();
638                    Ok(bytes.len())
639                }
640                fn flush(&mut self) -> std::io::Result<()> {
641                    Ok(())
642                }
643            }
644            serde_json::to_writer(HeaderBound(0), &frame.header).map_err(|_| CoreError::new(ErrorCode::Unavailable, "Active helper cannot stream this complete original request; preserve its operation and supervisor"))?;
645        }
646        let header = frame.header.clone();
647        let mailbox = Mailbox::new();
648        let sender = if matches!(&frame.header.message, Message::ServiceCall { call } if call.unprepared_worktree_mutation())
649            || matches!(
650                frame.header.message,
651                Message::ProofLookupBegin { .. }
652                    | Message::ProofCancel { .. }
653                    | Message::ProofRelease { .. }
654                    | Message::ServiceCall {
655                        call: w::Call::Inspect
656                            | w::Call::GitExisting { .. }
657                            | w::Call::LeaseObserve { .. }
658                            | w::Call::TerminalEnd { .. }
659                            | w::Call::TerminalInput { .. }
660                            | w::Call::TerminalResize { .. }
661                    }
662            ) {
663            &self.control
664        } else {
665            &self.sender
666        };
667        sender
668            .try_send(SendFrame {
669                frame,
670                mailbox: Some(mailbox.clone()),
671                scope: self.scope.clone(),
672            })
673            .map_err(|_| CoreError::new(ErrorCode::Busy, "Remote request queue full"))?;
674        self.requests.fetch_add(1, Ordering::SeqCst);
675        Ok(Request {
676            header,
677            mailbox,
678            transport: self.clone(),
679            finished: false,
680        })
681    }
682}
683fn write_frame(s: &mut UnixStream, f: &Frame) -> Result<(), CoreError> {
684    s.write_all(&f.encode().map_err(|_| w::invalid())?)
685        .map_err(io_error)
686}
687fn read_frame(s: &mut UnixStream, c: Option<&Context>) -> Result<Frame, CoreError> {
688    let mut prefix = [0; 8];
689    s.read_exact(&mut prefix).map_err(io_error)?;
690    let h = u32::from_be_bytes(prefix[..4].try_into().unwrap()) as usize;
691    let p = u32::from_be_bytes(prefix[4..].try_into().unwrap()) as usize;
692    if h > 65536 || p > 1048576 {
693        return Err(w::invalid());
694    }
695    let mut bytes = vec![0; h + p + 8];
696    bytes[..8].copy_from_slice(&prefix);
697    s.read_exact(&mut bytes[8..]).map_err(io_error)?;
698    Decoder::default()
699        .push(&bytes, c)
700        .map_err(|_| w::invalid())?
701        .pop()
702        .ok_or_else(w::invalid)
703}
704
705fn worker(
706    mut socket: UnixStream,
707    rx: mpsc::Receiver<SendFrame>,
708    controls: mpsc::Receiver<SendFrame>,
709    pending_credit: Credits,
710    worker_ids: Arc<std::sync::atomic::AtomicU64>,
711    live: Arc<AtomicBool>,
712) {
713    let mut routes: HashMap<u64, (Arc<Mailbox>, Header, Arc<Scope>)> = HashMap::new();
714    let mut decoder = Decoder::default();
715    let mut output = Vec::new();
716    let mut written = 0;
717    let mut _writing_scope: Option<Arc<Scope>> = None;
718    let mut input = [0; 65536];
719    let outcome = (|| -> Result<(), CoreError> {
720        loop {
721            if !live.load(Ordering::SeqCst) {
722                return Ok(());
723            }
724            if output.is_empty() {
725                // Wire completion does not retire the consumer's bounded owner. A finished
726                // consumer has returned all credit; flush it before releasing its route.
727                {
728                    let credits = pending_credit.lock().unwrap();
729                    routes.retain(|id, (mailbox, _, _)| {
730                        !mailbox.finished.load(Ordering::SeqCst) || credits.contains_key(id)
731                    });
732                }
733                let cancelled = routes
734                    .iter()
735                    .find(|(_, (m, _, _))| m.cancelled.load(Ordering::SeqCst))
736                    .map(|(id, _)| *id);
737
738                let next = if let Some(id) = cancelled {
739                    let (_, mut header, scope) = routes.remove(&id).unwrap();
740
741                    pending_credit.lock().unwrap().remove(&id);
742
743                    header.request_id = worker_ids.fetch_add(1, Ordering::SeqCst);
744
745                    header.message = Message::Cancel { request_id: id };
746
747                    header.operation_id = None;
748                    header.stream_id = None;
749                    header.sequence = None;
750
751                    Some(SendFrame {
752                        frame: Frame {
753                            header,
754                            payload: vec![],
755                        },
756                        mailbox: None,
757                        scope,
758                    })
759                } else if let Ok(control) = controls.try_recv() {
760                    Some(control)
761                } else {
762                    let credit = {
763                        let mut credits = pending_credit.lock().unwrap();
764
765                        let id = credits.keys().next().copied();
766
767                        id.and_then(|id| credits.remove(&id))
768                    };
769
770                    if let Some((mut header, bytes, scope)) = credit {
771                        header.message = Message::Credit {
772                            bytes: bytes as u32,
773                        };
774
775                        Some(SendFrame {
776                            frame: Frame {
777                                header,
778                                payload: vec![],
779                            },
780                            mailbox: None,
781                            scope,
782                        })
783                    } else {
784                        match rx.try_recv() {
785                            Ok(v) => Some(v),
786                            Err(mpsc::TryRecvError::Empty) => None,
787                            Err(mpsc::TryRecvError::Disconnected) => return Ok(()),
788                        }
789                    }
790                };
791
792                if let Some(v) = next {
793                    if let Some(mailbox) = v.mailbox {
794                        if mailbox.cancelled.load(Ordering::SeqCst) {
795                            continue;
796                        }
797                        if routes.len() >= 32 {
798                            mailbox.refuse_unsent(CoreError::new(
799                                ErrorCode::Busy,
800                                "Remote request capacity reached",
801                            ));
802
803                            continue;
804                        }
805                        routes.insert(
806                            v.frame.header.request_id,
807                            (mailbox, v.frame.header.clone(), v.scope.clone()),
808                        );
809                    } else if !matches!(
810                        v.frame.header.message,
811                        Message::Credit { .. } | Message::Cancel { .. }
812                    ) && !routes.contains_key(&v.frame.header.request_id)
813                    {
814                        continue;
815                    }
816                    output = v.frame.encode().map_err(|_| w::invalid())?;
817                    _writing_scope = Some(v.scope);
818                    written = 0;
819                }
820            }
821            if !output.is_empty() {
822                match socket.write(&output[written..]) {
823                    Ok(0) => return Err(closed()),
824                    Ok(n) => {
825                        written += n;
826                        if written == output.len() {
827                            output.clear();
828                            _writing_scope = None;
829                        }
830                    }
831                    Err(e) if e.kind() == std::io::ErrorKind::WouldBlock => {}
832                    Err(e) => return Err(io_error(e)),
833                }
834            }
835            match socket.read(&mut input) {
836                Ok(0) => return Err(closed()),
837                Ok(n) => {
838                    for f in decoder
839                        .push_replies(&input[..n], |id| {
840                            routes.get(&id).map(|(_, header, scope)| {
841                                (scope.context.clone(), header.operation_id.clone())
842                            })
843                        })
844                        .map_err(|_| w::invalid())?
845                    {
846                        let id = f.header.request_id;
847
848                        if let Some((route, request_header, expected)) = routes.get(&id) {
849                            if f.header.operation_id != request_header.operation_id
850                                || f.header.context.as_ref() != Some(&expected.context)
851                            {
852                                return Err(w::invalid());
853                            }
854                            route.push(f)?;
855                        }
856                    }
857                }
858                Err(e) if e.kind() == std::io::ErrorKind::WouldBlock => {}
859                Err(e) => return Err(io_error(e)),
860            }
861            std::thread::sleep(Duration::from_millis(1));
862        }
863    })();
864
865    live.store(false, Ordering::SeqCst);
866
867    let error = outcome.err().unwrap_or_else(closed);
868
869    for (route, _, _) in routes.values() {
870        route.fail(error.clone());
871    }
872    for receiver in [rx, controls] {
873        while let Ok(v) = receiver.try_recv() {
874            if let Some(m) = v.mailbox {
875                m.refuse_unsent(error.clone());
876            }
877        }
878    }
879}
880
881#[cfg(test)]
882#[path = "transport_tests.rs"]
883mod tests;
884
885pub(crate) struct Lifetime {
886    registrations: Arc<Mutex<()>>,
887    pub connected: Arc<AtomicBool>,
888    retired: Arc<Retirement>,
889    scopes: Arc<Mutex<Vec<Weak<Scope>>>>,
890    requests: Arc<std::sync::atomic::AtomicUsize>,
891}
892impl Lifetime {
893    pub fn has_active_views(&self) -> bool {
894        self.scopes
895            .lock()
896            .unwrap()
897            .iter()
898            .filter_map(Weak::upgrade)
899            .any(|s| s.active.load(Ordering::SeqCst))
900    }
901    pub fn idle(&self) -> bool {
902        // Registration/release can temporarily have neither an active view nor a queued request.
903        let Ok(_transition) = self.registrations.try_lock() else {
904            return false;
905        };
906        self.requests.load(Ordering::SeqCst) == 0 && !self.has_active_views()
907    }
908    pub fn shutdown(&self) -> Shutdown {
909        self.connected.store(false, Ordering::SeqCst);
910        Shutdown(self.retired.clone())
911    }
912}
buntrail_engine/remote/transport.rs

buntrail_engine/remote/
transport.rs
