kunobi_daemon/
control.rs

1//! In-memory health and drain handlers. The listener owns peer authentication.
2use crate::{
3    Lifecycle,
4    wire::{self, Control, Health, Hello, MessageKind, operation},
5};
6use std::{
7    io,
8    sync::{
9        Arc,
10        atomic::{AtomicBool, Ordering},
11    },
12};
13use tokio::io::{AsyncRead, AsyncWrite};
14
15mod watch;
16pub use watch::WatchClient;
17
18/// One process's lifecycle control, independent of application capacity and I/O.
19pub struct ControlService {
20    lifecycle: Arc<Lifecycle>,
21    ready: AtomicBool,
22    events: watch::Publisher,
23    process_id: u32,
24    generation: u64,
25    build: String,
26    revision: u64,
27}
28impl ControlService {
29    /// Start in the initializing state. The consumer marks readiness explicitly.
30    pub fn new(lifecycle: Arc<Lifecycle>, generation: u64, build: String, revision: u64) -> Self {
31        Self {
32            lifecycle,
33            ready: AtomicBool::new(false),
34            events: watch::Publisher::new(),
35            process_id: std::process::id(),
36            generation,
37            build,
38            revision,
39        }
40    }
41
42    /// Mark application initialization complete. Drain always takes precedence.
43    pub fn mark_ready(&self) {
44        if !self.ready.swap(true, Ordering::AcqRel) {
45            self.events.notify();
46        }
47    }
48
49    /// Read a bounded snapshot without waiting for storage, writers or draining.
50    pub fn snapshot(&self) -> Health {
51        let state = self.lifecycle.snapshot();
52        Health {
53            process_id: self.process_id,
54            generation: self.generation,
55            build: self.build.clone(),
56            revision: self.revision,
57            ready: self.ready.load(Ordering::Acquire) && !state.draining,
58            draining: state.draining,
59            active: state.active as u64,
60            ..Default::default()
61        }
62    }
63
64    /// Handle a control request without admitting application work or waiting for drain.
65    pub fn handle(&self, request: &Control) -> io::Result<Control> {
66        if request.kind != MessageKind::Lifecycle
67            || !request.payload.is_empty()
68            || !request.token.is_empty()
69            || request.offset.is_some()
70            || (request.generation != 0 && request.generation != self.generation)
71        {
72            return Err(io::Error::new(
73                io::ErrorKind::InvalidData,
74                "invalid lifecycle request",
75            ));
76        }
77        match request.operation {
78            operation::HEALTH => {}
79            operation::DRAIN => {
80                self.lifecycle.start_drain();
81            }
82            _ => {
83                return Err(io::Error::new(
84                    io::ErrorKind::InvalidData,
85                    "unsupported lifecycle request",
86                ));
87            }
88        }
89        self.snapshot().response(request)
90    }
91
92    /// Serve one authenticated connection under an absolute setup/control deadline.
93    /// WATCH continues beyond setup, with a five-second deadline per event write.
94    ///
95    /// No application operation is accepted here. The caller reserves control
96    /// admission before invoking this method and drops the stream after an error.
97    pub async fn serve<S: AsyncRead + AsyncWrite + Unpin>(
98        &self,
99        stream: S,
100        offer: &Hello,
101        deadline: tokio::time::Instant,
102    ) -> io::Result<()> {
103        let subscription = tokio::time::timeout_at(deadline, async {
104            let mut session = wire::AsyncSession::accept(stream, offer).await?;
105            let request = session.receive().await?;
106            if request.operation == operation::WATCH && request.kind == MessageKind::Lifecycle {
107                self.validate_watch(&request)?;
108                return Ok(Some((session, request)));
109            }
110            if session.agreement().capabilities & wire::capability::HEALTH_DETAILS == 0 {
111                return Err(io::Error::new(
112                    io::ErrorKind::InvalidData,
113                    "typed health not negotiated",
114                ));
115            }
116            session.send(&self.handle(&request)?).await?;
117            Ok(None)
118        })
119        .await
120        .map_err(|_| io::Error::from(io::ErrorKind::TimedOut))??;
121        if let Some((session, request)) = subscription {
122            self.serve_watch(session, request).await?;
123        }
124        Ok(())
125    }
126}
kunobi_daemon/control.rs

kunobi_daemon/
control.rs
