Skip to main content

agent_os_kernel/
fd_table.rs

1use std::collections::{btree_map::Values, BTreeMap, BTreeSet};
2use std::error::Error;
3use std::fmt;
4use std::sync::atomic::{AtomicU32, AtomicU64, AtomicUsize, Ordering};
5use std::sync::{Arc, Condvar, Mutex, MutexGuard};
6
7pub const MAX_FDS_PER_PROCESS: usize = 256;
8
9pub const O_RDONLY: u32 = 0;
10pub const O_WRONLY: u32 = 1;
11pub const O_RDWR: u32 = 2;
12pub const O_CREAT: u32 = 0o100;
13pub const O_EXCL: u32 = 0o200;
14pub const O_TRUNC: u32 = 0o1000;
15pub const O_APPEND: u32 = 0o2000;
16pub const O_NONBLOCK: u32 = 0o4000;
17pub const F_DUPFD: u32 = 0;
18pub const F_GETFD: u32 = 1;
19pub const F_SETFD: u32 = 2;
20pub const F_GETFL: u32 = 3;
21pub const F_SETFL: u32 = 4;
22pub const FD_CLOEXEC: u32 = 1;
23pub const LOCK_SH: u32 = 1;
24pub const LOCK_EX: u32 = 2;
25pub const LOCK_NB: u32 = 4;
26pub const LOCK_UN: u32 = 8;
27
28pub const FILETYPE_UNKNOWN: u8 = 0;
29pub const FILETYPE_CHARACTER_DEVICE: u8 = 2;
30pub const FILETYPE_DIRECTORY: u8 = 3;
31pub const FILETYPE_REGULAR_FILE: u8 = 4;
32pub const FILETYPE_PIPE: u8 = 6;
33pub const FILETYPE_SYMBOLIC_LINK: u8 = 7;
34
35pub type FdResult<T> = Result<T, FdTableError>;
36pub type SharedFileDescription = Arc<FileDescription>;
37
38#[derive(Debug, Clone, PartialEq, Eq)]
39pub struct FdTableError {
40    code: &'static str,
41    message: String,
42}
43
44impl FdTableError {
45    pub fn code(&self) -> &'static str {
46        self.code
47    }
48
49    fn bad_file_descriptor(fd: u32) -> Self {
50        Self {
51            code: "EBADF",
52            message: format!("bad file descriptor {fd}"),
53        }
54    }
55
56    fn too_many_open_files() -> Self {
57        Self {
58            code: "EMFILE",
59            message: String::from("too many open files"),
60        }
61    }
62
63    fn invalid_argument(message: impl Into<String>) -> Self {
64        Self {
65            code: "EINVAL",
66            message: message.into(),
67        }
68    }
69
70    fn would_block(message: impl Into<String>) -> Self {
71        Self {
72            code: "EWOULDBLOCK",
73            message: message.into(),
74        }
75    }
76}
77
78impl fmt::Display for FdTableError {
79    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
80        write!(f, "{}: {}", self.code, self.message)
81    }
82}
83
84impl Error for FdTableError {}
85
86#[derive(Debug)]
87pub struct FileDescription {
88    id: u64,
89    path: String,
90    lock_target: Option<FileLockTarget>,
91    cursor: AtomicU64,
92    flags: AtomicU32,
93    ref_count: AtomicUsize,
94}
95
96impl FileDescription {
97    pub fn new(id: u64, path: impl Into<String>, flags: u32) -> Self {
98        Self::with_ref_count_and_lock(id, path, flags, 1, None)
99    }
100
101    pub fn new_with_lock(
102        id: u64,
103        path: impl Into<String>,
104        flags: u32,
105        lock_target: Option<FileLockTarget>,
106    ) -> Self {
107        Self::with_ref_count_and_lock(id, path, flags, 1, lock_target)
108    }
109
110    pub fn with_ref_count(id: u64, path: impl Into<String>, flags: u32, ref_count: usize) -> Self {
111        Self::with_ref_count_and_lock(id, path, flags, ref_count, None)
112    }
113
114    pub fn with_ref_count_and_lock(
115        id: u64,
116        path: impl Into<String>,
117        flags: u32,
118        ref_count: usize,
119        lock_target: Option<FileLockTarget>,
120    ) -> Self {
121        Self {
122            id,
123            path: path.into(),
124            lock_target,
125            cursor: AtomicU64::new(0),
126            flags: AtomicU32::new(flags),
127            ref_count: AtomicUsize::new(ref_count),
128        }
129    }
130
131    pub fn id(&self) -> u64 {
132        self.id
133    }
134
135    pub fn path(&self) -> &str {
136        &self.path
137    }
138
139    pub fn lock_target(&self) -> Option<FileLockTarget> {
140        self.lock_target
141    }
142
143    pub fn cursor(&self) -> u64 {
144        self.cursor.load(Ordering::SeqCst)
145    }
146
147    pub fn set_cursor(&self, cursor: u64) {
148        self.cursor.store(cursor, Ordering::SeqCst);
149    }
150
151    pub fn flags(&self) -> u32 {
152        self.flags.load(Ordering::SeqCst)
153    }
154
155    pub fn update_flags(&self, mask: u32, flags: u32) -> u32 {
156        let mut current = self.flags();
157        loop {
158            let next = (current & !mask) | (flags & mask);
159            match self
160                .flags
161                .compare_exchange(current, next, Ordering::SeqCst, Ordering::SeqCst)
162            {
163                Ok(_) => return next,
164                Err(observed) => current = observed,
165            }
166        }
167    }
168
169    pub fn ref_count(&self) -> usize {
170        self.ref_count.load(Ordering::SeqCst)
171    }
172
173    pub fn increment_ref_count(&self) -> usize {
174        self.ref_count.fetch_add(1, Ordering::SeqCst) + 1
175    }
176
177    pub fn decrement_ref_count(&self) -> usize {
178        let mut current = self.ref_count.load(Ordering::SeqCst);
179        loop {
180            let next = current.saturating_sub(1);
181            match self
182                .ref_count
183                .compare_exchange(current, next, Ordering::SeqCst, Ordering::SeqCst)
184            {
185                Ok(_) => return next,
186                Err(observed) => current = observed,
187            }
188        }
189    }
190}
191
192#[derive(Debug, Clone)]
193pub struct FdEntry {
194    pub fd: u32,
195    pub description: SharedFileDescription,
196    pub status_flags: u32,
197    pub fd_flags: u32,
198    pub rights: u64,
199    pub filetype: u8,
200}
201
202#[derive(Debug, Clone, Copy, PartialEq, Eq)]
203pub struct FdStat {
204    pub filetype: u8,
205    pub flags: u32,
206    pub rights: u64,
207}
208
209#[derive(Debug, Clone)]
210pub struct StdioOverride {
211    pub description: SharedFileDescription,
212    pub filetype: u8,
213}
214
215#[derive(Debug, Clone)]
216struct DescriptionFactory {
217    next_description_id: Arc<AtomicU64>,
218}
219
220impl DescriptionFactory {
221    fn new(starting_id: u64) -> Self {
222        Self {
223            next_description_id: Arc::new(AtomicU64::new(starting_id)),
224        }
225    }
226
227    fn allocate(&self, path: &str, flags: u32) -> SharedFileDescription {
228        self.allocate_with_lock(path, flags, None)
229    }
230
231    fn allocate_with_lock(
232        &self,
233        path: &str,
234        flags: u32,
235        lock_target: Option<FileLockTarget>,
236    ) -> SharedFileDescription {
237        let next_id = self.next_description_id.fetch_add(1, Ordering::SeqCst);
238        Arc::new(FileDescription::new_with_lock(
239            next_id,
240            path,
241            flags,
242            lock_target,
243        ))
244    }
245}
246
247#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
248pub struct FileLockTarget {
249    ino: u64,
250}
251
252impl FileLockTarget {
253    pub const fn new(ino: u64) -> Self {
254        Self { ino }
255    }
256
257    pub const fn ino(self) -> u64 {
258        self.ino
259    }
260}
261
262#[derive(Debug, Clone, Copy, PartialEq, Eq)]
263enum FileLockMode {
264    Shared,
265    Exclusive,
266}
267
268#[derive(Debug, Clone, Copy, PartialEq, Eq)]
269pub enum FlockOperation {
270    Shared { nonblocking: bool },
271    Exclusive { nonblocking: bool },
272    Unlock,
273}
274
275impl FlockOperation {
276    pub fn from_bits(operation: u32) -> FdResult<Self> {
277        let nonblocking = operation & LOCK_NB != 0;
278        match operation & !LOCK_NB {
279            LOCK_SH => Ok(Self::Shared { nonblocking }),
280            LOCK_EX => Ok(Self::Exclusive { nonblocking }),
281            LOCK_UN => Ok(Self::Unlock),
282            _ => Err(FdTableError::invalid_argument(format!(
283                "invalid flock operation {operation:#x}"
284            ))),
285        }
286    }
287}
288
289#[derive(Debug, Clone)]
290pub struct ProcessFdTable {
291    entries: BTreeMap<u32, FdEntry>,
292    next_fd: u32,
293    alloc_desc: DescriptionFactory,
294    max_fds: usize,
295}
296
297impl ProcessFdTable {
298    fn new(alloc_desc: DescriptionFactory, max_fds: usize) -> Self {
299        Self {
300            entries: BTreeMap::new(),
301            next_fd: 3,
302            alloc_desc,
303            max_fds,
304        }
305    }
306
307    pub fn max_fds(&self) -> usize {
308        self.max_fds
309    }
310
311    pub fn init_stdio(
312        &mut self,
313        stdin_desc: SharedFileDescription,
314        stdout_desc: SharedFileDescription,
315        stderr_desc: SharedFileDescription,
316    ) {
317        self.entries.insert(
318            0,
319            FdEntry {
320                fd: 0,
321                description: stdin_desc,
322                status_flags: 0,
323                fd_flags: 0,
324                rights: 0,
325                filetype: FILETYPE_CHARACTER_DEVICE,
326            },
327        );
328        self.entries.insert(
329            1,
330            FdEntry {
331                fd: 1,
332                description: stdout_desc,
333                status_flags: 0,
334                fd_flags: 0,
335                rights: 0,
336                filetype: FILETYPE_CHARACTER_DEVICE,
337            },
338        );
339        self.entries.insert(
340            2,
341            FdEntry {
342                fd: 2,
343                description: stderr_desc,
344                status_flags: 0,
345                fd_flags: 0,
346                rights: 0,
347                filetype: FILETYPE_CHARACTER_DEVICE,
348            },
349        );
350    }
351
352    pub fn init_stdio_with_types(
353        &mut self,
354        stdin_desc: SharedFileDescription,
355        stdin_type: u8,
356        stdout_desc: SharedFileDescription,
357        stdout_type: u8,
358        stderr_desc: SharedFileDescription,
359        stderr_type: u8,
360    ) {
361        stdin_desc.increment_ref_count();
362        stdout_desc.increment_ref_count();
363        stderr_desc.increment_ref_count();
364        self.entries.insert(
365            0,
366            FdEntry {
367                fd: 0,
368                description: stdin_desc,
369                status_flags: 0,
370                fd_flags: 0,
371                rights: 0,
372                filetype: stdin_type,
373            },
374        );
375        self.entries.insert(
376            1,
377            FdEntry {
378                fd: 1,
379                description: stdout_desc,
380                status_flags: 0,
381                fd_flags: 0,
382                rights: 0,
383                filetype: stdout_type,
384            },
385        );
386        self.entries.insert(
387            2,
388            FdEntry {
389                fd: 2,
390                description: stderr_desc,
391                status_flags: 0,
392                fd_flags: 0,
393                rights: 0,
394                filetype: stderr_type,
395            },
396        );
397    }
398
399    pub fn open(&mut self, path: &str, flags: u32) -> FdResult<u32> {
400        self.open_with_details(path, flags, FILETYPE_REGULAR_FILE, None)
401    }
402
403    pub fn open_with_filetype(&mut self, path: &str, flags: u32, filetype: u8) -> FdResult<u32> {
404        self.open_with_details(path, flags, filetype, None)
405    }
406
407    pub fn open_with_details(
408        &mut self,
409        path: &str,
410        flags: u32,
411        filetype: u8,
412        lock_target: Option<FileLockTarget>,
413    ) -> FdResult<u32> {
414        let fd = self.allocate_fd()?;
415        let description =
416            self.alloc_desc
417                .allocate_with_lock(path, description_flags(flags), lock_target);
418        self.entries.insert(
419            fd,
420            FdEntry {
421                fd,
422                description,
423                status_flags: status_flags(flags),
424                fd_flags: 0,
425                rights: 0,
426                filetype,
427            },
428        );
429        Ok(fd)
430    }
431
432    pub fn open_with(
433        &mut self,
434        description: SharedFileDescription,
435        filetype: u8,
436        target_fd: Option<u32>,
437    ) -> FdResult<u32> {
438        let fd = match target_fd {
439            Some(fd) => {
440                self.validate_fd_bounds(fd)?;
441                fd
442            }
443            None => self.allocate_fd()?,
444        };
445        description.increment_ref_count();
446        self.entries.insert(
447            fd,
448            FdEntry {
449                fd,
450                description,
451                status_flags: 0,
452                fd_flags: 0,
453                rights: 0,
454                filetype,
455            },
456        );
457        Ok(fd)
458    }
459
460    pub fn get(&self, fd: u32) -> Option<&FdEntry> {
461        self.entries.get(&fd)
462    }
463
464    pub fn close(&mut self, fd: u32) -> bool {
465        let Some(entry) = self.entries.remove(&fd) else {
466            return false;
467        };
468        entry.description.decrement_ref_count();
469        true
470    }
471
472    pub fn dup(&mut self, fd: u32) -> FdResult<u32> {
473        self.dup_with_status_flags(fd, None)
474    }
475
476    pub fn dup_with_status_flags(
477        &mut self,
478        fd: u32,
479        status_flags_override: Option<u32>,
480    ) -> FdResult<u32> {
481        let entry = self
482            .entries
483            .get(&fd)
484            .cloned()
485            .ok_or_else(|| FdTableError::bad_file_descriptor(fd))?;
486        let new_fd = self.allocate_fd()?;
487        self.duplicate_entry(
488            &entry,
489            new_fd,
490            status_flags_override.unwrap_or(entry.status_flags),
491            0,
492        )
493    }
494
495    pub fn dup2(&mut self, old_fd: u32, new_fd: u32) -> FdResult<()> {
496        let entry = self
497            .entries
498            .get(&old_fd)
499            .cloned()
500            .ok_or_else(|| FdTableError::bad_file_descriptor(old_fd))?;
501        self.validate_fd_bounds(new_fd)?;
502        if old_fd == new_fd {
503            return Ok(());
504        }
505
506        if self.entries.contains_key(&new_fd) {
507            self.close(new_fd);
508        }
509
510        self.duplicate_entry(&entry, new_fd, entry.status_flags, 0)?;
511        Ok(())
512    }
513
514    pub fn stat(&self, fd: u32) -> FdResult<FdStat> {
515        let entry = self
516            .entries
517            .get(&fd)
518            .ok_or_else(|| FdTableError::bad_file_descriptor(fd))?;
519        Ok(FdStat {
520            filetype: entry.filetype,
521            flags: visible_fd_flags(entry.description.flags(), entry.status_flags),
522            rights: entry.rights,
523        })
524    }
525
526    pub fn fcntl(&mut self, fd: u32, command: u32, arg: u32) -> FdResult<u32> {
527        match command {
528            F_DUPFD => {
529                let entry = self
530                    .entries
531                    .get(&fd)
532                    .cloned()
533                    .ok_or_else(|| FdTableError::bad_file_descriptor(fd))?;
534                let min_fd = self.validate_fcntl_dup_min(arg)?;
535                let new_fd = self.allocate_fd_from(min_fd)?;
536                self.duplicate_entry(&entry, new_fd, entry.status_flags, 0)
537            }
538            F_GETFD => {
539                let entry = self
540                    .entries
541                    .get(&fd)
542                    .ok_or_else(|| FdTableError::bad_file_descriptor(fd))?;
543                Ok(entry.fd_flags & FD_CLOEXEC)
544            }
545            F_SETFD => {
546                let entry = self
547                    .entries
548                    .get_mut(&fd)
549                    .ok_or_else(|| FdTableError::bad_file_descriptor(fd))?;
550                entry.fd_flags = arg & FD_CLOEXEC;
551                Ok(0)
552            }
553            F_GETFL => {
554                let entry = self
555                    .entries
556                    .get(&fd)
557                    .ok_or_else(|| FdTableError::bad_file_descriptor(fd))?;
558                Ok(visible_fd_flags(
559                    entry.description.flags(),
560                    entry.status_flags,
561                ))
562            }
563            F_SETFL => {
564                let entry = self
565                    .entries
566                    .get_mut(&fd)
567                    .ok_or_else(|| FdTableError::bad_file_descriptor(fd))?;
568                entry.status_flags = arg & ENTRY_STATUS_FLAG_MASK;
569                entry.description.update_flags(SHARED_STATUS_FLAG_MASK, arg);
570                Ok(0)
571            }
572            _ => Err(FdTableError::invalid_argument(format!(
573                "unsupported fcntl command {command}"
574            ))),
575        }
576    }
577
578    pub fn fork(&self) -> Self {
579        let mut child = Self::new(self.alloc_desc.clone(), self.max_fds);
580        child.next_fd = self.next_fd;
581
582        for (fd, entry) in &self.entries {
583            entry.description.increment_ref_count();
584            child.entries.insert(
585                *fd,
586                FdEntry {
587                    fd: *fd,
588                    description: Arc::clone(&entry.description),
589                    status_flags: entry.status_flags,
590                    fd_flags: entry.fd_flags,
591                    rights: entry.rights,
592                    filetype: entry.filetype,
593                },
594            );
595        }
596
597        child
598    }
599
600    pub fn close_all(&mut self) {
601        let fds: Vec<u32> = self.entries.keys().copied().collect();
602        for fd in fds {
603            self.close(fd);
604        }
605    }
606
607    pub fn len(&self) -> usize {
608        self.entries.len()
609    }
610
611    pub fn is_empty(&self) -> bool {
612        self.entries.is_empty()
613    }
614
615    pub fn iter(&self) -> Values<'_, u32, FdEntry> {
616        self.entries.values()
617    }
618
619    fn allocate_fd(&mut self) -> FdResult<u32> {
620        if self.entries.len() >= self.max_fds {
621            return Err(FdTableError::too_many_open_files());
622        }
623
624        let start = usize::try_from(self.next_fd).unwrap_or(0) % self.max_fds;
625        for offset in 0..self.max_fds {
626            let candidate = ((start + offset) % self.max_fds) as u32;
627            if !self.entries.contains_key(&candidate) {
628                self.next_fd = candidate.saturating_add(1);
629                return Ok(candidate);
630            }
631        }
632
633        Err(FdTableError::too_many_open_files())
634    }
635
636    fn allocate_fd_from(&mut self, min_fd: u32) -> FdResult<u32> {
637        if self.entries.len() >= self.max_fds {
638            return Err(FdTableError::too_many_open_files());
639        }
640
641        if min_fd as usize >= self.max_fds {
642            return Err(FdTableError::invalid_argument(format!(
643                "fd {min_fd} exceeds process fd limit"
644            )));
645        }
646
647        for candidate in min_fd..self.max_fds as u32 {
648            if !self.entries.contains_key(&candidate) {
649                self.next_fd = candidate.saturating_add(1);
650                return Ok(candidate);
651            }
652        }
653
654        Err(FdTableError::too_many_open_files())
655    }
656
657    fn duplicate_entry(
658        &mut self,
659        entry: &FdEntry,
660        new_fd: u32,
661        status_flags: u32,
662        fd_flags: u32,
663    ) -> FdResult<u32> {
664        entry.description.increment_ref_count();
665        self.entries.insert(
666            new_fd,
667            FdEntry {
668                fd: new_fd,
669                description: Arc::clone(&entry.description),
670                status_flags,
671                fd_flags,
672                rights: entry.rights,
673                filetype: entry.filetype,
674            },
675        );
676        Ok(new_fd)
677    }
678
679    fn validate_fd_bounds(&self, fd: u32) -> FdResult<()> {
680        if fd as usize >= self.max_fds {
681            return Err(FdTableError::bad_file_descriptor(fd));
682        }
683        Ok(())
684    }
685
686    fn validate_fcntl_dup_min(&self, min_fd: u32) -> FdResult<u32> {
687        if min_fd as usize >= self.max_fds {
688            return Err(FdTableError::invalid_argument(format!(
689                "fd {min_fd} exceeds process fd limit"
690            )));
691        }
692        Ok(min_fd)
693    }
694}
695
696fn description_flags(flags: u32) -> u32 {
697    flags & !status_flags(flags)
698}
699
700fn status_flags(flags: u32) -> u32 {
701    flags & ENTRY_STATUS_FLAG_MASK
702}
703
704fn visible_fd_flags(description_flags: u32, entry_status_flags: u32) -> u32 {
705    (description_flags & (0b11 | SHARED_STATUS_FLAG_MASK))
706        | (entry_status_flags & ENTRY_STATUS_FLAG_MASK)
707}
708
709const SHARED_STATUS_FLAG_MASK: u32 = O_APPEND;
710const ENTRY_STATUS_FLAG_MASK: u32 = O_NONBLOCK;
711
712impl<'a> IntoIterator for &'a ProcessFdTable {
713    type Item = &'a FdEntry;
714    type IntoIter = Values<'a, u32, FdEntry>;
715
716    fn into_iter(self) -> Self::IntoIter {
717        self.entries.values()
718    }
719}
720
721#[derive(Debug, Clone)]
722pub struct FdTableManager {
723    tables: BTreeMap<u32, ProcessFdTable>,
724    alloc_desc: DescriptionFactory,
725    max_fds: usize,
726}
727
728impl Default for FdTableManager {
729    fn default() -> Self {
730        Self {
731            tables: BTreeMap::new(),
732            alloc_desc: DescriptionFactory::new(1),
733            max_fds: MAX_FDS_PER_PROCESS,
734        }
735    }
736}
737
738impl FdTableManager {
739    pub fn new() -> Self {
740        Self::default()
741    }
742
743    pub fn with_max_fds(max_fds: usize) -> Self {
744        Self {
745            max_fds,
746            ..Self::default()
747        }
748    }
749
750    pub fn create(&mut self, pid: u32) -> &mut ProcessFdTable {
751        let mut table = ProcessFdTable::new(self.alloc_desc.clone(), self.max_fds);
752        table.init_stdio(
753            self.alloc_desc.allocate("/dev/stdin", O_RDONLY),
754            self.alloc_desc.allocate("/dev/stdout", O_WRONLY),
755            self.alloc_desc.allocate("/dev/stderr", O_WRONLY),
756        );
757        self.remove(pid);
758        self.tables.insert(pid, table);
759        self.tables
760            .get_mut(&pid)
761            .expect("newly created FD table should be stored")
762    }
763
764    pub fn create_with_stdio(
765        &mut self,
766        pid: u32,
767        stdin_override: Option<StdioOverride>,
768        stdout_override: Option<StdioOverride>,
769        stderr_override: Option<StdioOverride>,
770    ) -> &mut ProcessFdTable {
771        let mut table = ProcessFdTable::new(self.alloc_desc.clone(), self.max_fds);
772        let stdin_desc = stdin_override
773            .as_ref()
774            .map(|entry| Arc::clone(&entry.description))
775            .unwrap_or_else(|| self.alloc_desc.allocate("/dev/stdin", O_RDONLY));
776        let stdout_desc = stdout_override
777            .as_ref()
778            .map(|entry| Arc::clone(&entry.description))
779            .unwrap_or_else(|| self.alloc_desc.allocate("/dev/stdout", O_WRONLY));
780        let stderr_desc = stderr_override
781            .as_ref()
782            .map(|entry| Arc::clone(&entry.description))
783            .unwrap_or_else(|| self.alloc_desc.allocate("/dev/stderr", O_WRONLY));
784
785        table.init_stdio_with_types(
786            stdin_desc,
787            stdin_override
788                .as_ref()
789                .map(|entry| entry.filetype)
790                .unwrap_or(FILETYPE_CHARACTER_DEVICE),
791            stdout_desc,
792            stdout_override
793                .as_ref()
794                .map(|entry| entry.filetype)
795                .unwrap_or(FILETYPE_CHARACTER_DEVICE),
796            stderr_desc,
797            stderr_override
798                .as_ref()
799                .map(|entry| entry.filetype)
800                .unwrap_or(FILETYPE_CHARACTER_DEVICE),
801        );
802        self.remove(pid);
803        self.tables.insert(pid, table);
804        self.tables
805            .get_mut(&pid)
806            .expect("newly created FD table should be stored")
807    }
808
809    pub fn fork(&mut self, parent_pid: u32, child_pid: u32) -> &mut ProcessFdTable {
810        if !self.tables.contains_key(&parent_pid) {
811            return self.create(child_pid);
812        }
813
814        let child = self
815            .tables
816            .get(&parent_pid)
817            .expect("parent table presence was checked")
818            .fork();
819        self.remove(child_pid);
820        self.tables.insert(child_pid, child);
821        self.tables
822            .get_mut(&child_pid)
823            .expect("forked FD table should be stored")
824    }
825
826    pub fn get(&self, pid: u32) -> Option<&ProcessFdTable> {
827        self.tables.get(&pid)
828    }
829
830    pub fn get_mut(&mut self, pid: u32) -> Option<&mut ProcessFdTable> {
831        self.tables.get_mut(&pid)
832    }
833
834    pub fn has(&self, pid: u32) -> bool {
835        self.tables.contains_key(&pid)
836    }
837
838    pub fn len(&self) -> usize {
839        self.tables.len()
840    }
841
842    pub fn is_empty(&self) -> bool {
843        self.tables.is_empty()
844    }
845
846    pub fn total_open_fds(&self) -> usize {
847        self.tables.values().map(ProcessFdTable::len).sum()
848    }
849
850    pub fn pids(&self) -> Vec<u32> {
851        self.tables.keys().copied().collect()
852    }
853
854    pub fn remove(&mut self, pid: u32) {
855        if let Some(mut table) = self.tables.remove(&pid) {
856            table.close_all();
857        }
858    }
859}
860
861#[derive(Debug, Clone, Default)]
862pub struct FileLockManager {
863    inner: Arc<FileLockManagerInner>,
864}
865
866#[derive(Debug, Default)]
867struct FileLockManagerInner {
868    state: Mutex<FileLockState>,
869    wake: Condvar,
870}
871
872#[derive(Debug, Default)]
873struct FileLockState {
874    entries: BTreeMap<FileLockTarget, FileLockEntry>,
875}
876
877#[derive(Debug, Default)]
878struct FileLockEntry {
879    shared: BTreeSet<u64>,
880    exclusive: Option<u64>,
881}
882
883impl FileLockManager {
884    pub fn new() -> Self {
885        Self::default()
886    }
887
888    pub fn apply(
889        &self,
890        owner_id: u64,
891        target: FileLockTarget,
892        operation: FlockOperation,
893    ) -> FdResult<()> {
894        match operation {
895            FlockOperation::Shared { nonblocking } => {
896                self.acquire(owner_id, target, FileLockMode::Shared, nonblocking)
897            }
898            FlockOperation::Exclusive { nonblocking } => {
899                self.acquire(owner_id, target, FileLockMode::Exclusive, nonblocking)
900            }
901            FlockOperation::Unlock => {
902                self.release_owner(owner_id);
903                Ok(())
904            }
905        }
906    }
907
908    pub fn release_owner(&self, owner_id: u64) -> bool {
909        let mut state = lock_or_recover(&self.inner.state);
910        let mut released = false;
911        state.entries.retain(|_, entry| {
912            let entry_changed = entry.shared.remove(&owner_id) || entry.exclusive == Some(owner_id);
913            if entry.exclusive == Some(owner_id) {
914                entry.exclusive = None;
915            }
916            released |= entry_changed;
917            !entry.is_empty()
918        });
919        drop(state);
920        if released {
921            self.inner.wake.notify_all();
922        }
923        released
924    }
925
926    fn acquire(
927        &self,
928        owner_id: u64,
929        target: FileLockTarget,
930        mode: FileLockMode,
931        nonblocking: bool,
932    ) -> FdResult<()> {
933        let mut state = lock_or_recover(&self.inner.state);
934        loop {
935            let entry = state.entries.entry(target).or_default();
936            if entry.can_grant(owner_id, mode) {
937                entry.grant(owner_id, mode);
938                return Ok(());
939            }
940
941            if nonblocking {
942                return Err(FdTableError::would_block(
943                    "advisory file lock is unavailable",
944                ));
945            }
946
947            state = wait_or_recover(&self.inner.wake, state);
948        }
949    }
950}
951
952impl FileLockEntry {
953    fn can_grant(&self, owner_id: u64, mode: FileLockMode) -> bool {
954        match mode {
955            FileLockMode::Shared => self.exclusive.is_none_or(|owner| owner == owner_id),
956            FileLockMode::Exclusive => {
957                self.exclusive.is_none_or(|owner| owner == owner_id)
958                    && self.shared.iter().all(|owner| *owner == owner_id)
959            }
960        }
961    }
962
963    fn grant(&mut self, owner_id: u64, mode: FileLockMode) {
964        match mode {
965            FileLockMode::Shared => {
966                self.exclusive = None;
967                self.shared.insert(owner_id);
968            }
969            FileLockMode::Exclusive => {
970                self.shared.retain(|owner| *owner != owner_id);
971                self.exclusive = Some(owner_id);
972            }
973        }
974    }
975
976    fn is_empty(&self) -> bool {
977        self.exclusive.is_none() && self.shared.is_empty()
978    }
979}
980
981fn lock_or_recover<T>(mutex: &Mutex<T>) -> MutexGuard<'_, T> {
982    mutex
983        .lock()
984        .unwrap_or_else(|poisoned| poisoned.into_inner())
985}
986
987fn wait_or_recover<'a, T>(condvar: &Condvar, guard: MutexGuard<'a, T>) -> MutexGuard<'a, T> {
988    condvar
989        .wait(guard)
990        .unwrap_or_else(|poisoned| poisoned.into_inner())
991}