executor/distaff/src/heartbeat.rs

use std::{
    collections::BTreeMap,
    mem::ManuallyDrop,
    ops::Deref,
    ptr::NonNull,
    sync::{
        Arc,
        atomic::{AtomicBool, Ordering},
    },
    time::Instant,
};

use parking_lot::Mutex;

use crate::latch::WorkerLatch;

#[derive(Debug, Clone)]
pub struct HeartbeatList {
    inner: Arc<Mutex<HeartbeatListInner>>,
}

impl HeartbeatList {
    pub fn new() -> Self {
        Self {
            inner: Arc::new(Mutex::new(HeartbeatListInner::new())),
        }
    }

    pub fn notify_nth(&self, n: usize) {
        self.inner.lock().notify_nth(n);
    }

    pub fn notify_all(&self) {
        let mut inner = self.inner.lock();
        for (_, heartbeat) in inner.heartbeats.iter_mut() {
            heartbeat.set();
        }
    }

    pub fn len(&self) -> usize {
        self.inner.lock().len()
    }

    pub fn new_heartbeat(&self) -> OwnedHeartbeatReceiver {
        let (recv, _) = self.inner.lock().new_heartbeat();
        OwnedHeartbeatReceiver {
            list: self.clone(),
            receiver: ManuallyDrop::new(recv),
        }
    }

    pub fn inner(
        &self,
    ) -> parking_lot::lock_api::MappedMutexGuard<
        '_,
        parking_lot::RawMutex,
        BTreeMap<u64, HeartbeatSender>,
    > {
        parking_lot::MutexGuard::map(self.inner.lock(), |inner| &mut inner.heartbeats)
    }
}

#[derive(Debug)]
struct HeartbeatListInner {
    heartbeats: BTreeMap<u64, HeartbeatSender>,
    heartbeat_index: u64,
}

impl HeartbeatListInner {
    fn new() -> Self {
        Self {
            heartbeats: BTreeMap::new(),
            heartbeat_index: 0,
        }
    }

    fn iter(&self) -> std::collections::btree_map::Values<'_, u64, HeartbeatSender> {
        self.heartbeats.values()
    }

    fn notify_nth(&mut self, n: usize) {
        if let Some((_, heartbeat)) = self.heartbeats.iter_mut().nth(n) {
            heartbeat.set();
        }
    }

    fn len(&self) -> usize {
        self.heartbeats.len()
    }

    fn new_heartbeat(&mut self) -> (HeartbeatReceiver, u64) {
        let heartbeat = Heartbeat::new(self.heartbeat_index);
        let (recv, send, i) = heartbeat.into_recv_send();
        self.heartbeats.insert(i, send);
        self.heartbeat_index += 1;
        (recv, i)
    }

    fn remove_heartbeat(&mut self, receiver: HeartbeatReceiver) {
        if let Some(send) = self.heartbeats.remove(&receiver.i) {
            _ = Heartbeat::from_recv_send(receiver, send);
        }
    }
}

pub struct OwnedHeartbeatReceiver {
    list: HeartbeatList,
    receiver: ManuallyDrop<HeartbeatReceiver>,
}

impl Deref for OwnedHeartbeatReceiver {
    type Target = HeartbeatReceiver;

    fn deref(&self) -> &Self::Target {
        &self.receiver
    }
}

impl Drop for OwnedHeartbeatReceiver {
    fn drop(&mut self) {
        // SAFETY:
        // `AtomicBool` is `Sync` and `Send`, so it can be safely shared between threads.
        unsafe {
            let receiver = ManuallyDrop::take(&mut self.receiver);
            self.list.inner.lock().remove_heartbeat(receiver);
        }
    }
}

#[derive(Debug)]
pub struct Heartbeat {
    ptr: NonNull<(AtomicBool, WorkerLatch)>,
    i: u64,
}

#[derive(Debug)]
pub struct HeartbeatReceiver {
    ptr: NonNull<(AtomicBool, WorkerLatch)>,
    i: u64,
}

unsafe impl Send for HeartbeatReceiver {}

impl Drop for Heartbeat {
    fn drop(&mut self) {
        // SAFETY:
        // `AtomicBool` is `Sync` and `Send`, so it can be safely shared between threads.
        unsafe {
            let _ = Box::from_raw(self.ptr.as_ptr());
        }
    }
}

#[derive(Debug)]
pub struct HeartbeatSender {
    ptr: NonNull<(AtomicBool, WorkerLatch)>,
    pub last_heartbeat: Instant,
}

unsafe impl Send for HeartbeatSender {}

impl Heartbeat {
    fn new(i: u64) -> Heartbeat {
        // SAFETY:
        // `AtomicBool` is `Sync` and `Send`, so it can be safely shared between threads.
        let ptr = NonNull::new(Box::into_raw(Box::new((
            AtomicBool::new(true),
            WorkerLatch::new(),
        ))))
        .unwrap();
        Self { ptr, i }
    }

    pub fn from_recv_send(recv: HeartbeatReceiver, send: HeartbeatSender) -> Heartbeat {
        // SAFETY:
        // `AtomicBool` is `Sync` and `Send`, so it can be safely shared between threads.
        _ = send;
        let ptr = recv.ptr;
        let i = recv.i;
        Heartbeat { ptr, i }
    }

    pub fn into_recv_send(self) -> (HeartbeatReceiver, HeartbeatSender, u64) {
        // don't drop the `Heartbeat` yet
        let Self { ptr, i } = *ManuallyDrop::new(self);

        (
            HeartbeatReceiver { ptr, i },
            HeartbeatSender {
                ptr,
                last_heartbeat: Instant::now(),
            },
            i,
        )
    }
}

impl HeartbeatReceiver {
    pub fn take(&self) -> bool {
        unsafe {
            // SAFETY:
            // `AtomicBool` is `Sync` and `Send`, so it can be safely shared between threads.
            self.ptr.as_ref().0.swap(false, Ordering::Relaxed)
        }
    }

    pub fn wait(&self) {
        unsafe { self.ptr.as_ref().1.wait() };
    }

    pub fn raw_latch(&self) -> *const WorkerLatch {
        unsafe { &raw const self.ptr.as_ref().1 }
    }

    pub fn latch(&self) -> &WorkerLatch {
        unsafe { &self.ptr.as_ref().1 }
    }

    pub fn id(&self) -> usize {
        self.ptr.as_ptr() as usize
    }

    pub fn index(&self) -> u64 {
        self.i
    }
}

impl HeartbeatSender {
    pub fn set(&mut self) {
        // SAFETY:
        // `AtomicBool` is `Sync` and `Send`, so it can be safely shared between threads.
        unsafe { self.ptr.as_ref().0.store(true, Ordering::Relaxed) };
        self.last_heartbeat = Instant::now();
    }

    pub fn is_waiting(&self) -> bool {
        unsafe { self.ptr.as_ref().1.is_waiting() }
    }
    pub fn wake(&self) {
        unsafe { self.ptr.as_ref().1.wake() };
    }
}