Skip to main content

fory_core/util/
sync.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18use std::cell::UnsafeCell;
19use std::ops::{Deref, DerefMut};
20use std::sync::atomic::{AtomicBool, Ordering};
21use std::thread;
22
23pub struct Spinlock<T> {
24    data: UnsafeCell<T>,
25    flag: AtomicBool,
26}
27
28unsafe impl<T: Send> Send for Spinlock<T> {}
29unsafe impl<T: Sync> Sync for Spinlock<T> {}
30
31impl<T> Spinlock<T> {
32    pub fn new(data: T) -> Self {
33        Spinlock {
34            data: UnsafeCell::new(data),
35            flag: AtomicBool::new(false),
36        }
37    }
38
39    pub fn lock(&self) -> SpinlockGuard<'_, T> {
40        let mut spins = 0;
41        while self
42            .flag
43            .compare_exchange(false, true, Ordering::Acquire, Ordering::Relaxed)
44            .is_err()
45        {
46            // Spin for a few iterations
47            if spins < 10 {
48                std::hint::spin_loop();
49                spins += 1;
50            } else {
51                // Then yield to the scheduler
52                thread::yield_now();
53                spins = 0; // reset spin counter
54            }
55        }
56        SpinlockGuard { lock: self }
57    }
58
59    fn unlock(&self) {
60        self.flag.store(false, Ordering::Release);
61    }
62}
63
64#[allow(clippy::needless_lifetimes)]
65pub struct SpinlockGuard<'a, T> {
66    lock: &'a Spinlock<T>,
67}
68
69#[allow(clippy::needless_lifetimes)]
70impl<'a, T> Drop for SpinlockGuard<'a, T> {
71    fn drop(&mut self) {
72        self.lock.unlock();
73    }
74}
75
76#[allow(clippy::needless_lifetimes)]
77impl<'a, T> Deref for SpinlockGuard<'a, T> {
78    type Target = T;
79    fn deref(&self) -> &Self::Target {
80        unsafe { &*self.lock.data.get() }
81    }
82}
83
84#[allow(clippy::needless_lifetimes)]
85impl<'a, T> DerefMut for SpinlockGuard<'a, T> {
86    fn deref_mut(&mut self) -> &mut Self::Target {
87        unsafe { &mut *self.lock.data.get() }
88    }
89}