Skip to main content

agentos_kernel/
pipe_manager.rs

1use crate::fd_table::{
2    FdResult, FileDescription, ProcessFdTable, SharedFileDescription, FILETYPE_PIPE, O_RDONLY,
3    O_WRONLY,
4};
5use crate::poll::{PollEvents, PollNotifier, POLLERR, POLLHUP, POLLIN, POLLOUT};
6use std::collections::{BTreeMap, VecDeque};
7use std::error::Error;
8use std::fmt;
9use std::sync::{Arc, Condvar, Mutex, MutexGuard};
10use std::time::Duration;
11use web_time::Instant;
12
13pub const MAX_PIPE_BUFFER_BYTES: usize = 65_536;
14pub const PIPE_BUF_BYTES: usize = 4_096;
15
16pub type PipeResult<T> = Result<T, PipeError>;
17
18#[derive(Debug, Clone, PartialEq, Eq)]
19pub struct PipeError {
20    code: &'static str,
21    message: String,
22}
23
24impl PipeError {
25    pub fn code(&self) -> &'static str {
26        self.code
27    }
28
29    fn bad_file_descriptor(message: impl Into<String>) -> Self {
30        Self {
31            code: "EBADF",
32            message: message.into(),
33        }
34    }
35
36    fn broken_pipe(message: impl Into<String>) -> Self {
37        Self {
38            code: "EPIPE",
39            message: message.into(),
40        }
41    }
42
43    fn would_block(message: impl Into<String>) -> Self {
44        Self {
45            code: "EAGAIN",
46            message: message.into(),
47        }
48    }
49}
50
51impl fmt::Display for PipeError {
52    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
53        write!(f, "{}: {}", self.code, self.message)
54    }
55}
56
57impl Error for PipeError {}
58
59#[derive(Debug, Clone)]
60pub struct PipeEnd {
61    pub description: SharedFileDescription,
62    pub filetype: u8,
63}
64
65#[derive(Debug, Clone)]
66pub struct PipePair {
67    pub read: PipeEnd,
68    pub write: PipeEnd,
69}
70
71#[derive(Debug, Clone, Copy, PartialEq, Eq)]
72struct PipeRef {
73    pipe_id: u64,
74    end: PipeSide,
75}
76
77#[derive(Debug, Clone, Copy, PartialEq, Eq)]
78enum PipeSide {
79    Read,
80    Write,
81}
82
83#[derive(Debug, Default)]
84struct PendingRead {
85    length: usize,
86    result: Option<Option<Vec<u8>>>,
87}
88
89#[derive(Debug, Default)]
90struct PipeState {
91    buffer: VecDeque<Vec<u8>>,
92    closed_read: bool,
93    closed_write: bool,
94    waiting_reads: VecDeque<u64>,
95}
96
97#[derive(Debug)]
98struct PipeManagerState {
99    pipes: BTreeMap<u64, PipeState>,
100    desc_to_pipe: BTreeMap<u64, PipeRef>,
101    waiters: BTreeMap<u64, PendingRead>,
102    next_pipe_id: u64,
103    next_desc_id: u64,
104    next_waiter_id: u64,
105}
106
107impl Default for PipeManagerState {
108    fn default() -> Self {
109        Self {
110            pipes: BTreeMap::new(),
111            desc_to_pipe: BTreeMap::new(),
112            waiters: BTreeMap::new(),
113            next_pipe_id: 1,
114            next_desc_id: 100_000,
115            next_waiter_id: 1,
116        }
117    }
118}
119
120#[derive(Debug)]
121struct PipeManagerInner {
122    state: Mutex<PipeManagerState>,
123    waiters: Condvar,
124}
125
126#[derive(Debug, Clone)]
127pub struct PipeManager {
128    inner: Arc<PipeManagerInner>,
129    notifier: Option<PollNotifier>,
130}
131
132impl Default for PipeManager {
133    fn default() -> Self {
134        Self {
135            inner: Arc::new(PipeManagerInner {
136                state: Mutex::new(PipeManagerState::default()),
137                waiters: Condvar::new(),
138            }),
139            notifier: None,
140        }
141    }
142}
143
144impl PipeManager {
145    pub fn new() -> Self {
146        Self::default()
147    }
148
149    pub(crate) fn with_notifier(notifier: PollNotifier) -> Self {
150        Self {
151            notifier: Some(notifier),
152            ..Self::default()
153        }
154    }
155
156    pub fn create_pipe(&self) -> PipePair {
157        let mut state = lock_or_recover(&self.inner.state);
158        let pipe_id = state.next_pipe_id;
159        state.next_pipe_id += 1;
160
161        let read_id = state.next_desc_id;
162        state.next_desc_id += 1;
163        let write_id = state.next_desc_id;
164        state.next_desc_id += 1;
165
166        state.pipes.insert(pipe_id, PipeState::default());
167        state.desc_to_pipe.insert(
168            read_id,
169            PipeRef {
170                pipe_id,
171                end: PipeSide::Read,
172            },
173        );
174        state.desc_to_pipe.insert(
175            write_id,
176            PipeRef {
177                pipe_id,
178                end: PipeSide::Write,
179            },
180        );
181        drop(state);
182
183        PipePair {
184            read: PipeEnd {
185                description: Arc::new(FileDescription::with_ref_count(
186                    read_id,
187                    format!("pipe:{pipe_id}:read"),
188                    O_RDONLY,
189                    0,
190                )),
191                filetype: FILETYPE_PIPE,
192            },
193            write: PipeEnd {
194                description: Arc::new(FileDescription::with_ref_count(
195                    write_id,
196                    format!("pipe:{pipe_id}:write"),
197                    O_WRONLY,
198                    0,
199                )),
200                filetype: FILETYPE_PIPE,
201            },
202        }
203    }
204
205    pub fn poll(&self, description_id: u64, requested: PollEvents) -> PipeResult<PollEvents> {
206        let state = lock_or_recover(&self.inner.state);
207        let pipe_ref = state
208            .desc_to_pipe
209            .get(&description_id)
210            .copied()
211            .ok_or_else(|| PipeError::bad_file_descriptor("not a pipe end"))?;
212        let pipe = state
213            .pipes
214            .get(&pipe_ref.pipe_id)
215            .ok_or_else(|| PipeError::bad_file_descriptor("pipe not found"))?;
216
217        let mut events = PollEvents::empty();
218        match pipe_ref.end {
219            PipeSide::Read => {
220                if requested.intersects(POLLIN) && !pipe.buffer.is_empty() {
221                    events |= POLLIN;
222                }
223                if pipe.closed_write {
224                    events |= POLLHUP;
225                }
226            }
227            PipeSide::Write => {
228                if pipe.closed_read {
229                    events |= POLLERR;
230                } else if requested.intersects(POLLOUT)
231                    && (available_capacity(pipe) > 0 || !pipe.waiting_reads.is_empty())
232                {
233                    events |= POLLOUT;
234                }
235            }
236        }
237
238        Ok(events)
239    }
240
241    pub fn write(&self, description_id: u64, data: impl AsRef<[u8]>) -> PipeResult<usize> {
242        self.write_with_mode(description_id, data, true)
243    }
244
245    pub fn write_blocking(&self, description_id: u64, data: impl AsRef<[u8]>) -> PipeResult<usize> {
246        self.write_with_mode(description_id, data, false)
247    }
248
249    pub fn write_with_mode(
250        &self,
251        description_id: u64,
252        data: impl AsRef<[u8]>,
253        nonblocking: bool,
254    ) -> PipeResult<usize> {
255        let payload = data.as_ref();
256        let mut state = lock_or_recover(&self.inner.state);
257        let pipe_ref = state
258            .desc_to_pipe
259            .get(&description_id)
260            .copied()
261            .ok_or_else(|| PipeError::bad_file_descriptor("not a pipe write end"))?;
262        if pipe_ref.end != PipeSide::Write {
263            return Err(PipeError::bad_file_descriptor("not a pipe write end"));
264        }
265
266        loop {
267            let waiter_id = {
268                let pipe = state
269                    .pipes
270                    .get_mut(&pipe_ref.pipe_id)
271                    .ok_or_else(|| PipeError::bad_file_descriptor("pipe not found"))?;
272                if pipe.closed_write {
273                    return Err(PipeError::broken_pipe("write end closed"));
274                }
275                if pipe.closed_read {
276                    return Err(PipeError::broken_pipe("read end closed"));
277                }
278                pipe.waiting_reads.pop_front()
279            };
280
281            if let Some(waiter_id) = waiter_id {
282                let waiter_length = match state.waiters.get(&waiter_id) {
283                    Some(waiter) => waiter.length,
284                    None => continue,
285                };
286                let delivered_len = waiter_length.min(payload.len());
287                let delivered = payload[..delivered_len].to_vec();
288                let remainder = &payload[delivered_len..];
289
290                if !remainder.is_empty() {
291                    let pipe = state
292                        .pipes
293                        .get_mut(&pipe_ref.pipe_id)
294                        .ok_or_else(|| PipeError::bad_file_descriptor("pipe not found"))?;
295                    pipe.buffer.push_back(remainder.to_vec());
296                }
297
298                if let Some(waiter) = state.waiters.get_mut(&waiter_id) {
299                    waiter.result = Some(Some(delivered));
300                    self.notify_waiters_and_pollers();
301                    return Ok(payload.len());
302                }
303                continue;
304            }
305
306            let current_buffer_size = {
307                let pipe = state
308                    .pipes
309                    .get(&pipe_ref.pipe_id)
310                    .ok_or_else(|| PipeError::bad_file_descriptor("pipe not found"))?;
311                buffer_size(&pipe.buffer)
312            };
313            let available = MAX_PIPE_BUFFER_BYTES.saturating_sub(current_buffer_size);
314
315            if payload.len() <= PIPE_BUF_BYTES {
316                if available >= payload.len() {
317                    let pipe = state
318                        .pipes
319                        .get_mut(&pipe_ref.pipe_id)
320                        .ok_or_else(|| PipeError::bad_file_descriptor("pipe not found"))?;
321                    pipe.buffer.push_back(payload.to_vec());
322                    self.notify_waiters_and_pollers();
323                    return Ok(payload.len());
324                }
325            } else if available > 0 {
326                let chunk_len = available.min(payload.len());
327                let pipe = state
328                    .pipes
329                    .get_mut(&pipe_ref.pipe_id)
330                    .ok_or_else(|| PipeError::bad_file_descriptor("pipe not found"))?;
331                pipe.buffer.push_back(payload[..chunk_len].to_vec());
332                self.notify_waiters_and_pollers();
333                return Ok(chunk_len);
334            }
335
336            if nonblocking {
337                return Err(PipeError::would_block("pipe buffer full"));
338            }
339
340            state = wait_or_recover(&self.inner.waiters, state);
341        }
342    }
343
344    pub fn read(&self, description_id: u64, length: usize) -> PipeResult<Option<Vec<u8>>> {
345        self.read_with_timeout(description_id, length, None)
346    }
347
348    pub fn read_with_timeout(
349        &self,
350        description_id: u64,
351        length: usize,
352        timeout: Option<Duration>,
353    ) -> PipeResult<Option<Vec<u8>>> {
354        let mut state = lock_or_recover(&self.inner.state);
355        let pipe_ref = state
356            .desc_to_pipe
357            .get(&description_id)
358            .copied()
359            .ok_or_else(|| PipeError::bad_file_descriptor("not a pipe read end"))?;
360        if pipe_ref.end != PipeSide::Read {
361            return Err(PipeError::bad_file_descriptor("not a pipe read end"));
362        }
363
364        let mut waiter_id = None;
365        let deadline = timeout.map(|duration| Instant::now() + duration);
366
367        loop {
368            if let Some(id) = waiter_id {
369                if let Some(waiter) = state.waiters.get_mut(&id) {
370                    if let Some(result) = waiter.result.take() {
371                        state.waiters.remove(&id);
372                        return Ok(result);
373                    }
374                }
375            }
376
377            {
378                let pipe = state
379                    .pipes
380                    .get_mut(&pipe_ref.pipe_id)
381                    .ok_or_else(|| PipeError::bad_file_descriptor("pipe not found"))?;
382
383                if !pipe.buffer.is_empty() {
384                    let result = drain_buffer(&mut pipe.buffer, length);
385                    self.notify_waiters_and_pollers();
386                    return Ok(Some(result));
387                }
388
389                if pipe.closed_write {
390                    if let Some(id) = waiter_id {
391                        state.waiters.remove(&id);
392                    }
393                    return Ok(None);
394                }
395            }
396
397            let id = if let Some(id) = waiter_id {
398                id
399            } else {
400                let next = state.next_waiter_id;
401                state.next_waiter_id += 1;
402                state.waiters.insert(
403                    next,
404                    PendingRead {
405                        length,
406                        result: None,
407                    },
408                );
409                let Some(pipe) = state.pipes.get_mut(&pipe_ref.pipe_id) else {
410                    state.waiters.remove(&next);
411                    return Err(PipeError::bad_file_descriptor("pipe not found"));
412                };
413                pipe.waiting_reads.push_back(next);
414                self.notify_waiters_and_pollers();
415                waiter_id = Some(next);
416                next
417            };
418
419            let Some(deadline) = deadline else {
420                state = wait_or_recover(&self.inner.waiters, state);
421                if !state.waiters.contains_key(&id) {
422                    waiter_id = None;
423                }
424                continue;
425            };
426
427            let now = Instant::now();
428            if now >= deadline {
429                if let Some(id) = waiter_id.take() {
430                    state.waiters.remove(&id);
431                    if let Some(pipe) = state.pipes.get_mut(&pipe_ref.pipe_id) {
432                        pipe.waiting_reads.retain(|queued| *queued != id);
433                    }
434                    self.notify_waiters_and_pollers();
435                }
436                return Err(PipeError::would_block("pipe read timed out"));
437            }
438
439            let remaining = deadline.saturating_duration_since(now);
440            let (next_state, wait_result) =
441                wait_timeout_or_recover(&self.inner.waiters, state, remaining);
442            state = next_state;
443            if !state.waiters.contains_key(&id) {
444                waiter_id = None;
445            }
446            if wait_result.timed_out() {
447                if let Some(id) = waiter_id.take() {
448                    state.waiters.remove(&id);
449                    if let Some(pipe) = state.pipes.get_mut(&pipe_ref.pipe_id) {
450                        pipe.waiting_reads.retain(|queued| *queued != id);
451                    }
452                    self.notify_waiters_and_pollers();
453                }
454                return Err(PipeError::would_block("pipe read timed out"));
455            }
456        }
457    }
458
459    pub fn close(&self, description_id: u64) {
460        let mut state = lock_or_recover(&self.inner.state);
461        let Some(pipe_ref) = state.desc_to_pipe.remove(&description_id) else {
462            return;
463        };
464
465        let (waiter_ids, remove_pipe, should_notify) =
466            if let Some(pipe) = state.pipes.get_mut(&pipe_ref.pipe_id) {
467                match pipe_ref.end {
468                    PipeSide::Read => {
469                        pipe.closed_read = true;
470                        (Vec::new(), pipe.closed_read && pipe.closed_write, true)
471                    }
472                    PipeSide::Write => {
473                        pipe.closed_write = true;
474                        let waiter_ids = pipe.waiting_reads.drain(..).collect::<Vec<_>>();
475                        (waiter_ids, pipe.closed_read && pipe.closed_write, true)
476                    }
477                }
478            } else {
479                (Vec::new(), false, false)
480            };
481
482        for waiter_id in waiter_ids {
483            if let Some(waiter) = state.waiters.get_mut(&waiter_id) {
484                waiter.result = Some(None);
485            }
486        }
487
488        if remove_pipe {
489            state.pipes.remove(&pipe_ref.pipe_id);
490        }
491        if should_notify {
492            self.notify_waiters_and_pollers();
493        }
494    }
495
496    pub fn is_pipe(&self, description_id: u64) -> bool {
497        lock_or_recover(&self.inner.state)
498            .desc_to_pipe
499            .contains_key(&description_id)
500    }
501
502    pub fn pipe_id_for(&self, description_id: u64) -> Option<u64> {
503        lock_or_recover(&self.inner.state)
504            .desc_to_pipe
505            .get(&description_id)
506            .map(|pipe_ref| pipe_ref.pipe_id)
507    }
508
509    pub fn pipe_count(&self) -> usize {
510        lock_or_recover(&self.inner.state).pipes.len()
511    }
512
513    pub fn buffered_bytes(&self) -> usize {
514        lock_or_recover(&self.inner.state)
515            .pipes
516            .values()
517            .map(|pipe| buffer_size(&pipe.buffer))
518            .sum()
519    }
520
521    pub fn waiting_reader_count(&self, description_id: u64) -> PipeResult<usize> {
522        let state = lock_or_recover(&self.inner.state);
523        let pipe_ref = state
524            .desc_to_pipe
525            .get(&description_id)
526            .copied()
527            .ok_or_else(|| PipeError::bad_file_descriptor("not a pipe end"))?;
528        let pipe = state
529            .pipes
530            .get(&pipe_ref.pipe_id)
531            .ok_or_else(|| PipeError::bad_file_descriptor("pipe not found"))?;
532        Ok(pipe.waiting_reads.len())
533    }
534
535    pub fn pending_read_waiter_count(&self) -> usize {
536        lock_or_recover(&self.inner.state).waiters.len()
537    }
538
539    pub fn create_pipe_fds(&self, fd_table: &mut ProcessFdTable) -> FdResult<(u32, u32)> {
540        let pipe = self.create_pipe();
541        let read_fd =
542            fd_table.open_with(Arc::clone(&pipe.read.description), FILETYPE_PIPE, None)?;
543        match fd_table.open_with(Arc::clone(&pipe.write.description), FILETYPE_PIPE, None) {
544            Ok(write_fd) => Ok((read_fd, write_fd)),
545            Err(error) => {
546                fd_table.close(read_fd);
547                self.close(pipe.read.description.id());
548                self.close(pipe.write.description.id());
549                Err(error)
550            }
551        }
552    }
553
554    fn notify_waiters_and_pollers(&self) {
555        self.inner.waiters.notify_all();
556        if let Some(notifier) = &self.notifier {
557            notifier.notify();
558        }
559    }
560}
561
562fn buffer_size(buffer: &VecDeque<Vec<u8>>) -> usize {
563    buffer.iter().map(Vec::len).sum()
564}
565
566fn available_capacity(pipe: &PipeState) -> usize {
567    MAX_PIPE_BUFFER_BYTES.saturating_sub(buffer_size(&pipe.buffer))
568}
569
570fn drain_buffer(buffer: &mut VecDeque<Vec<u8>>, length: usize) -> Vec<u8> {
571    let mut chunks = Vec::new();
572    let mut remaining = length;
573
574    while remaining > 0 {
575        let Some(chunk) = buffer.pop_front() else {
576            break;
577        };
578        if chunk.len() <= remaining {
579            remaining -= chunk.len();
580            chunks.push(chunk);
581        } else {
582            let (head, tail) = chunk.split_at(remaining);
583            chunks.push(head.to_vec());
584            buffer.push_front(tail.to_vec());
585            remaining = 0;
586        }
587    }
588
589    if chunks.len() == 1 {
590        return chunks.pop().expect("single chunk should exist");
591    }
592
593    let total = chunks.iter().map(Vec::len).sum();
594    let mut result = Vec::with_capacity(total);
595    for chunk in chunks {
596        result.extend_from_slice(&chunk);
597    }
598    result
599}
600
601fn lock_or_recover<'a, T>(mutex: &'a Mutex<T>) -> MutexGuard<'a, T> {
602    match mutex.lock() {
603        Ok(guard) => guard,
604        Err(poisoned) => poisoned.into_inner(),
605    }
606}
607
608fn wait_or_recover<'a, T>(condvar: &Condvar, guard: MutexGuard<'a, T>) -> MutexGuard<'a, T> {
609    match condvar.wait(guard) {
610        Ok(guard) => guard,
611        Err(poisoned) => poisoned.into_inner(),
612    }
613}
614
615fn wait_timeout_or_recover<'a, T>(
616    condvar: &Condvar,
617    guard: MutexGuard<'a, T>,
618    timeout: Duration,
619) -> (MutexGuard<'a, T>, std::sync::WaitTimeoutResult) {
620    match condvar.wait_timeout(guard, timeout) {
621        Ok(result) => result,
622        Err(poisoned) => poisoned.into_inner(),
623    }
624}