fidl_next_protocol/
lockers.rs

1// Copyright 2024 The Fuchsia Authors. All rights reserved.
2// Use of this source code is governed by a BSD-style license that can be
3// found in the LICENSE file.
4
5use core::mem::replace;
6use core::task::Waker;
7
8use thiserror::Error;
9
10/// Errors which may occur when attempting to write to a locker.
11#[derive(Debug, Error)]
12pub enum LockerError {
13    /// The locker was not in a writeable state.
14    #[error("the locker was not in a writeable state")]
15    NotWriteable,
16
17    /// The locker's ordinal did not match the given ordinal.
18    #[error("write expected a locker with ordinal {expected}, but found {actual}")]
19    MismatchedOrdinal { expected: u64, actual: u64 },
20}
21
22/// A dual-custody memory location which can be written to and read.
23pub enum Locker<T> {
24    Free(usize),
25    Pending { expected_ordinal: u64, read_waker: Option<Waker> },
26    Ready(T),
27    Canceled,
28    Finished,
29}
30
31impl<T> Locker<T> {
32    /// Writes the given value to the locker.
33    ///
34    /// On success, returns `true` if the reader canceled and the locker can now be freed.
35    pub fn write(&mut self, ordinal: u64, value: T) -> Result<bool, LockerError> {
36        match self {
37            Self::Free(_) | Self::Ready(_) | Self::Finished => Err(LockerError::NotWriteable),
38            Self::Pending { expected_ordinal, read_waker } => {
39                if *expected_ordinal != ordinal {
40                    return Err(LockerError::MismatchedOrdinal {
41                        expected: *expected_ordinal,
42                        actual: ordinal,
43                    });
44                }
45                if let Some(waker) = read_waker.take() {
46                    waker.wake();
47                }
48                *self = Locker::Ready(value);
49                Ok(false)
50            }
51            Self::Canceled => Ok(true),
52        }
53    }
54
55    /// Retrieves the result written to the locker, if any.
56    ///
57    /// On success, this finishes the locker and allows it to be freed. If the locker was pending,
58    /// the given waker is registered.
59    pub fn read(&mut self, waker: &Waker) -> Option<T> {
60        match self {
61            Self::Free(_) | Self::Canceled | Self::Finished => unreachable!(),
62            Self::Pending { read_waker, .. } => {
63                *read_waker = Some(waker.clone());
64                None
65            }
66            Self::Ready(_) => {
67                let Self::Ready(result) = replace(self, Self::Finished) else { unreachable!() };
68                Some(result)
69            }
70        }
71    }
72
73    /// Cancels the pending read for the locker.
74    ///
75    /// Returns `true` if the locker was already written and can now be freed.
76    pub fn cancel(&mut self) -> bool {
77        match self {
78            Self::Free(_) | Self::Canceled | Self::Finished => unreachable!(),
79            Self::Pending { .. } => {
80                *self = Self::Canceled;
81                false
82            }
83            Self::Ready(_) => {
84                *self = Self::Canceled;
85                true
86            }
87        }
88    }
89}
90
91/// A free list of [`Locker`]s.
92///
93/// Allocated lockers are assigned a unique ID which is not reused until the locker is freed.
94pub struct Lockers<T> {
95    lockers: Vec<Locker<T>>,
96    next_free: usize,
97}
98
99impl<T> Lockers<T> {
100    /// Returns a new `Lockers`.
101    pub fn new() -> Self {
102        Self { lockers: Vec::new(), next_free: 0 }
103    }
104
105    /// Allocates a fresh locker, returning its index.
106    pub fn alloc(&mut self, ordinal: u64) -> u32 {
107        let new_locker = Locker::Pending { expected_ordinal: ordinal, read_waker: None };
108
109        if self.next_free < self.lockers.len() {
110            let locker = replace(&mut self.lockers[self.next_free], new_locker);
111            let Locker::Free(next_free) = locker else {
112                panic!("unexpected allocation in free list");
113            };
114            replace(&mut self.next_free, next_free) as u32
115        } else {
116            let result = self.lockers.len();
117            self.lockers.push(new_locker);
118            self.next_free = self.lockers.len();
119            result as u32
120        }
121    }
122
123    /// Frees the locker with the given index.
124    pub fn free(&mut self, index: u32) {
125        self.lockers[index as usize] = Locker::Free(self.next_free);
126        self.next_free = index as usize;
127    }
128
129    /// Wakes up all of the `Pending` lockers.
130    pub fn wake_all(&mut self) {
131        for locker in self.lockers.iter_mut() {
132            if let Locker::Pending { read_waker, .. } = locker {
133                if let Some(waker) = read_waker.take() {
134                    waker.wake();
135                }
136            }
137        }
138    }
139
140    /// Gets the locker corresponding to the given index.
141    pub fn get(&mut self, index: u32) -> Option<&mut Locker<T>> {
142        let locker = self.lockers.get_mut(index as usize)?;
143        if matches!(locker, Locker::Free(_)) {
144            None
145        } else {
146            Some(locker)
147        }
148    }
149}