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.
45use core::mem::replace;
6use core::task::Waker;
78use thiserror::Error;
910/// 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")]
15NotWriteable,
1617/// The locker's ordinal did not match the given ordinal.
18#[error("write expected a locker with ordinal {expected}, but found {actual}")]
19MismatchedOrdinal { expected: u64, actual: u64 },
20}
2122/// 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}
3031impl<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.
35pub fn write(&mut self, ordinal: u64, value: T) -> Result<bool, LockerError> {
36match self {
37Self::Free(_) | Self::Ready(_) | Self::Finished => Err(LockerError::NotWriteable),
38Self::Pending { expected_ordinal, read_waker } => {
39if *expected_ordinal != ordinal {
40return Err(LockerError::MismatchedOrdinal {
41 expected: *expected_ordinal,
42 actual: ordinal,
43 });
44 }
45if let Some(waker) = read_waker.take() {
46 waker.wake();
47 }
48*self = Locker::Ready(value);
49Ok(false)
50 }
51Self::Canceled => Ok(true),
52 }
53 }
5455/// 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.
59pub fn read(&mut self, waker: &Waker) -> Option<T> {
60match self {
61Self::Free(_) | Self::Canceled | Self::Finished => unreachable!(),
62Self::Pending { read_waker, .. } => {
63*read_waker = Some(waker.clone());
64None
65}
66Self::Ready(_) => {
67let Self::Ready(result) = replace(self, Self::Finished) else { unreachable!() };
68Some(result)
69 }
70 }
71 }
7273/// Cancels the pending read for the locker.
74 ///
75 /// Returns `true` if the locker was already written and can now be freed.
76pub fn cancel(&mut self) -> bool {
77match self {
78Self::Free(_) | Self::Canceled | Self::Finished => unreachable!(),
79Self::Pending { .. } => {
80*self = Self::Canceled;
81false
82}
83Self::Ready(_) => {
84*self = Self::Canceled;
85true
86}
87 }
88 }
89}
9091/// 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}
9899impl<T> Lockers<T> {
100/// Returns a new `Lockers`.
101pub fn new() -> Self {
102Self { lockers: Vec::new(), next_free: 0 }
103 }
104105/// Allocates a fresh locker, returning its index.
106pub fn alloc(&mut self, ordinal: u64) -> u32 {
107let new_locker = Locker::Pending { expected_ordinal: ordinal, read_waker: None };
108109if self.next_free < self.lockers.len() {
110let locker = replace(&mut self.lockers[self.next_free], new_locker);
111let Locker::Free(next_free) = locker else {
112panic!("unexpected allocation in free list");
113 };
114 replace(&mut self.next_free, next_free) as u32
115 } else {
116let result = self.lockers.len();
117self.lockers.push(new_locker);
118self.next_free = self.lockers.len();
119 result as u32
120 }
121 }
122123/// Frees the locker with the given index.
124pub fn free(&mut self, index: u32) {
125self.lockers[index as usize] = Locker::Free(self.next_free);
126self.next_free = index as usize;
127 }
128129/// Wakes up all of the `Pending` lockers.
130pub fn wake_all(&mut self) {
131for locker in self.lockers.iter_mut() {
132if let Locker::Pending { read_waker, .. } = locker {
133if let Some(waker) = read_waker.take() {
134 waker.wake();
135 }
136 }
137 }
138 }
139140/// Gets the locker corresponding to the given index.
141pub fn get(&mut self, index: u32) -> Option<&mut Locker<T>> {
142let locker = self.lockers.get_mut(index as usize)?;
143if matches!(locker, Locker::Free(_)) {
144None
145} else {
146Some(locker)
147 }
148 }
149}