1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
// Copyright 2020 The ChromiumOS Authors
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

use std::cell::UnsafeCell;
use std::future::Future;
use std::mem;
use std::pin::Pin;
use std::ptr::NonNull;
use std::sync::atomic::AtomicBool;
use std::sync::atomic::AtomicU8;
use std::sync::atomic::Ordering;
use std::sync::Arc;
use std::task::Context;
use std::task::Poll;
use std::task::Waker;

use intrusive_collections::intrusive_adapter;
use intrusive_collections::linked_list::LinkedList;
use intrusive_collections::linked_list::LinkedListOps;
use intrusive_collections::DefaultLinkOps;
use intrusive_collections::LinkOps;

use super::super::sync::SpinLock;

// An atomic version of a LinkedListLink. See https://github.com/Amanieu/intrusive-rs/issues/47 for
// more details.
#[repr(align(128))]
pub struct AtomicLink {
    prev: UnsafeCell<Option<NonNull<AtomicLink>>>,
    next: UnsafeCell<Option<NonNull<AtomicLink>>>,
    linked: AtomicBool,
}

impl AtomicLink {
    fn new() -> AtomicLink {
        AtomicLink {
            linked: AtomicBool::new(false),
            prev: UnsafeCell::new(None),
            next: UnsafeCell::new(None),
        }
    }

    fn is_linked(&self) -> bool {
        self.linked.load(Ordering::Relaxed)
    }
}

impl DefaultLinkOps for AtomicLink {
    type Ops = AtomicLinkOps;

    const NEW: Self::Ops = AtomicLinkOps;
}

// SAFETY:
// Safe because the only way to mutate `AtomicLink` is via the `LinkedListOps` trait whose methods
// are all unsafe and require that the caller has first called `acquire_link` (and had it return
// true) to use them safely.
unsafe impl Send for AtomicLink {}
// SAFETY: See safety comment for impl Send
unsafe impl Sync for AtomicLink {}

#[derive(Copy, Clone, Default)]
pub struct AtomicLinkOps;

// TODO(b/315998194): Add safety comment
#[allow(clippy::undocumented_unsafe_blocks)]
unsafe impl LinkOps for AtomicLinkOps {
    type LinkPtr = NonNull<AtomicLink>;

    unsafe fn acquire_link(&mut self, ptr: Self::LinkPtr) -> bool {
        !ptr.as_ref().linked.swap(true, Ordering::Acquire)
    }

    unsafe fn release_link(&mut self, ptr: Self::LinkPtr) {
        ptr.as_ref().linked.store(false, Ordering::Release)
    }
}

// TODO(b/315998194): Add safety comment
#[allow(clippy::undocumented_unsafe_blocks)]
unsafe impl LinkedListOps for AtomicLinkOps {
    unsafe fn next(&self, ptr: Self::LinkPtr) -> Option<Self::LinkPtr> {
        *ptr.as_ref().next.get()
    }

    unsafe fn prev(&self, ptr: Self::LinkPtr) -> Option<Self::LinkPtr> {
        *ptr.as_ref().prev.get()
    }

    unsafe fn set_next(&mut self, ptr: Self::LinkPtr, next: Option<Self::LinkPtr>) {
        *ptr.as_ref().next.get() = next;
    }

    unsafe fn set_prev(&mut self, ptr: Self::LinkPtr, prev: Option<Self::LinkPtr>) {
        *ptr.as_ref().prev.get() = prev;
    }
}

#[derive(Clone, Copy)]
pub enum Kind {
    Shared,
    Exclusive,
}

enum State {
    Init,
    Waiting(Waker),
    Woken,
    Finished,
    Processing,
}

// Indicates the queue to which the waiter belongs. It is the responsibility of the Mutex and
// Condvar implementations to update this value when adding/removing a Waiter from their respective
// waiter lists.
#[repr(u8)]
#[derive(Debug, Eq, PartialEq)]
pub enum WaitingFor {
    // The waiter is either not linked into  a waiter list or it is linked into a temporary list.
    None = 0,
    // The waiter is linked into the Mutex's waiter list.
    Mutex = 1,
    // The waiter is linked into the Condvar's waiter list.
    Condvar = 2,
}

// Represents a thread currently blocked on a Condvar or on acquiring a Mutex.
pub struct Waiter {
    link: AtomicLink,
    state: SpinLock<State>,
    cancel: fn(usize, &Waiter, bool),
    cancel_data: usize,
    kind: Kind,
    waiting_for: AtomicU8,
}

impl Waiter {
    // Create a new, initialized Waiter.
    //
    // `kind` should indicate whether this waiter represent a thread that is waiting for a shared
    // lock or an exclusive lock.
    //
    // `cancel` is the function that is called when a `WaitFuture` (returned by the `wait()`
    // function) is dropped before it can complete. `cancel_data` is used as the first parameter of
    // the `cancel` function. The second parameter is the `Waiter` that was canceled and the third
    // parameter indicates whether the `WaitFuture` was dropped after it was woken (but before it
    // was polled to completion). A value of `false` for the third parameter may already be stale
    // by the time the cancel function runs and so does not guarantee that the waiter was not woken.
    // In this case, implementations should still check if the Waiter was woken. However, a value of
    // `true` guarantees that the waiter was already woken up so no additional checks are necessary.
    // In this case, the cancel implementation should wake up the next waiter in its wait list, if
    // any.
    //
    // `waiting_for` indicates the waiter list to which this `Waiter` will be added. See the
    // documentation of the `WaitingFor` enum for the meaning of the different values.
    pub fn new(
        kind: Kind,
        cancel: fn(usize, &Waiter, bool),
        cancel_data: usize,
        waiting_for: WaitingFor,
    ) -> Waiter {
        Waiter {
            link: AtomicLink::new(),
            state: SpinLock::new(State::Init),
            cancel,
            cancel_data,
            kind,
            waiting_for: AtomicU8::new(waiting_for as u8),
        }
    }

    // The kind of lock that this `Waiter` is waiting to acquire.
    pub fn kind(&self) -> Kind {
        self.kind
    }

    // Returns true if this `Waiter` is currently linked into a waiter list.
    pub fn is_linked(&self) -> bool {
        self.link.is_linked()
    }

    // Indicates the waiter list to which this `Waiter` belongs.
    pub fn is_waiting_for(&self) -> WaitingFor {
        match self.waiting_for.load(Ordering::Acquire) {
            0 => WaitingFor::None,
            1 => WaitingFor::Mutex,
            2 => WaitingFor::Condvar,
            v => panic!("Unknown value for `WaitingFor`: {}", v),
        }
    }

    // Change the waiter list to which this `Waiter` belongs. This will panic if called when the
    // `Waiter` is still linked into a waiter list.
    pub fn set_waiting_for(&self, waiting_for: WaitingFor) {
        self.waiting_for.store(waiting_for as u8, Ordering::Release);
    }

    // Reset the Waiter back to its initial state. Panics if this `Waiter` is still linked into a
    // waiter list.
    pub fn reset(&self, waiting_for: WaitingFor) {
        debug_assert!(!self.is_linked(), "Cannot reset `Waiter` while linked");
        self.set_waiting_for(waiting_for);

        let mut state = self.state.lock();
        if let State::Waiting(waker) = mem::replace(&mut *state, State::Init) {
            mem::drop(state);
            mem::drop(waker);
        }
    }

    // Wait until woken up by another thread.
    pub fn wait(&self) -> WaitFuture<'_> {
        WaitFuture { waiter: self }
    }

    // Wake up the thread associated with this `Waiter`. Panics if `waiting_for()` does not return
    // `WaitingFor::None` or if `is_linked()` returns true.
    pub fn wake(&self) {
        debug_assert!(!self.is_linked(), "Cannot wake `Waiter` while linked");
        debug_assert_eq!(self.is_waiting_for(), WaitingFor::None);

        let mut state = self.state.lock();

        if let State::Waiting(waker) = mem::replace(&mut *state, State::Woken) {
            mem::drop(state);
            waker.wake();
        }
    }
}

pub struct WaitFuture<'w> {
    waiter: &'w Waiter,
}

impl<'w> Future for WaitFuture<'w> {
    type Output = ();

    fn poll(self: Pin<&mut Self>, cx: &mut Context) -> Poll<Self::Output> {
        let mut state = self.waiter.state.lock();

        match mem::replace(&mut *state, State::Processing) {
            State::Init => {
                *state = State::Waiting(cx.waker().clone());

                Poll::Pending
            }
            State::Waiting(old_waker) => {
                *state = State::Waiting(cx.waker().clone());
                mem::drop(state);
                mem::drop(old_waker);

                Poll::Pending
            }
            State::Woken => {
                *state = State::Finished;
                Poll::Ready(())
            }
            State::Finished => {
                panic!("Future polled after returning Poll::Ready");
            }
            State::Processing => {
                panic!("Unexpected waker state");
            }
        }
    }
}

impl<'w> Drop for WaitFuture<'w> {
    fn drop(&mut self) {
        let state = self.waiter.state.lock();

        match *state {
            State::Finished => {}
            State::Processing => panic!("Unexpected waker state"),
            State::Woken => {
                mem::drop(state);

                // We were woken but not polled.  Wake up the next waiter.
                (self.waiter.cancel)(self.waiter.cancel_data, self.waiter, true);
            }
            _ => {
                mem::drop(state);

                // Not woken.  No need to wake up any waiters.
                (self.waiter.cancel)(self.waiter.cancel_data, self.waiter, false);
            }
        }
    }
}

intrusive_adapter!(pub WaiterAdapter = Arc<Waiter>: Waiter { link: AtomicLink });

pub type WaiterList = LinkedList<WaiterAdapter>;