lamellar 0.8.0

pub(crate) mod atomic;
pub(crate) mod collective;
pub(crate) mod comm;
pub(crate) mod fabric;
pub(crate) mod mem;
pub(crate) mod rdma;

use super::{
    comm::{CmdQStatus, CommInfo, CommShutdown},
    command_queues::CommandQueue,
    Comm, Lamellae, LamellaeInit, LamellaeShutdown, LamellaeUtil, Ser, SerializeHeader,
    SerializedData, SERIALIZE_HEADER_LEN,
};
use crate::{
    config, env_var::HeapMode, lamellae::libfabric_sys_lamellae::comm::LibfabricSysComm,
    lamellar_arch::LamellarArchRT, scheduler::Scheduler,
};

use async_trait::async_trait;
use futures_util::stream::FuturesUnordered;
use futures_util::StreamExt;
use std::sync::atomic::{AtomicU8, Ordering};
use std::sync::Arc;
use tracing::trace;

pub(crate) struct LibfabricSysBuilder {
    my_pe: usize,
    num_pes: usize,
    libfabric_sys_comm: Arc<Comm>,
}

impl LibfabricSysBuilder {
    pub(crate) fn new(provider: &str, domain: &str) -> LibfabricSysBuilder {
        let provider = if !provider.is_empty() {
            Some(provider)
        } else {
            None
        };
        let domain = if !domain.is_empty() {
            Some(domain)
        } else {
            None
        };
        let libfabric_sys_comm: Arc<Comm> =
            Arc::new(LibfabricSysComm::new(provider, domain).into());
        LibfabricSysBuilder {
            my_pe: libfabric_sys_comm.my_pe(),
            num_pes: libfabric_sys_comm.num_pes(),
            libfabric_sys_comm: libfabric_sys_comm,
        }
    }
}

impl LamellaeInit for LibfabricSysBuilder {
    fn init_fabric(&mut self) -> (usize, usize) {
        (self.my_pe, self.num_pes)
    }
    fn init_lamellae(&mut self, scheduler: Arc<Scheduler>) -> Arc<Lamellae> {
        let libfabric_sys = LibfabricSys::new(
            self.my_pe,
            self.num_pes,
            self.libfabric_sys_comm.clone(),
            scheduler.clone(),
        );
        trace!("created new libfabric_sys instance");
        let cq = libfabric_sys.cq();
        trace!("created command queue for libfabric_sys");
        let libfabric_sys = Arc::new(Lamellae::LibfabricSys(libfabric_sys));
        let libfabric_sys_clone = libfabric_sys.clone();
        let cq_clone = cq.clone();
        scheduler.submit_task(async move {
            cq_clone.recv_data(libfabric_sys_clone.clone()).await;
        });

        let cq_clone = cq.clone();
        scheduler.submit_task(async move {
            cq_clone.alloc_task().await;
        });
        let cq_clone = cq.clone();
        scheduler.submit_task(async move {
            cq_clone.panic_task().await;
        });
        libfabric_sys
    }
}

pub(crate) struct LibfabricSys {
    my_pe: usize,
    num_pes: usize,
    libfabric_sys_comm: Arc<Comm>,
    active: Arc<AtomicU8>,
    cq: Arc<CommandQueue>,
}

impl std::fmt::Debug for LibfabricSys {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(
            f,
            "LibfabricSys {{ my_pe: {}, num_pes: {},  active: {:?} }}",
            self.my_pe, self.num_pes, self.active,
        )
    }
}

impl LibfabricSys {
    fn new(
        my_pe: usize,
        num_pes: usize,
        libfabric_sys_comm: Arc<Comm>,
        scheduler: Arc<Scheduler>,
    ) -> LibfabricSys {
        // println!("my_pe {:?} num_pes {:?}",my_pe,num_pes);
        let active = Arc::new(AtomicU8::new(CmdQStatus::Active as u8));
        LibfabricSys {
            my_pe: my_pe,
            num_pes: num_pes,
            libfabric_sys_comm: libfabric_sys_comm.clone(),
            active: active.clone(),
            cq: Arc::new(CommandQueue::new(
                libfabric_sys_comm,
                scheduler,
                my_pe,
                num_pes,
                active,
            )),
        }
    }
    fn cq(&self) -> Arc<CommandQueue> {
        self.cq.clone()
    }
    pub(crate) fn wait_all_print(&self) {
        self.cq.wait_all_print();
    }
    pub(crate) fn comm(&self) -> &Comm {
        &self.libfabric_sys_comm
    }
}

impl LamellaeShutdown for LibfabricSys {
    fn shutdown(&self) {
        // println!("libfabric Lamellae shuting down");
        let _ = self.active.compare_exchange(
            CmdQStatus::Active as u8,
            CmdQStatus::ShuttingDown as u8,
            Ordering::SeqCst,
            Ordering::SeqCst,
        );
        // println!("set active to 0");
        while self.active.load(Ordering::SeqCst) != CmdQStatus::Finished as u8
            && self.active.load(Ordering::SeqCst) != CmdQStatus::Panic as u8
        {
            self.cq.scheduler.exec_task();
        }
        // println!("libfabric Lamellae shut down");
    }

    fn force_shutdown(&self) {
        self.cq.send_panic();
        self.active
            .store(CmdQStatus::Panic as u8, Ordering::Relaxed);
    }
    fn force_deinit(&self) {
        self.libfabric_sys_comm.force_shutdown();
    }
}

#[async_trait]
impl LamellaeUtil for LibfabricSys {
    async fn send_to_pes_async(
        &self,
        pe: Option<usize>,
        team: Arc<LamellarArchRT>,
        data: SerializedData,
    ) {
        if let Some(pe) = pe {
            self.cq.send_data(data, pe).await;
        } else {
            let mut futures = team
                .team_iter()
                .filter(|pe| pe != &self.my_pe)
                .map(|pe| self.cq.send_data(data.clone(), pe))
                .collect::<FuturesUnordered<_>>(); //in theory this launches all the futures before waiting...
            while let Some(_) = futures.next().await {}
        }
    }

    async fn request_new_alloc(&self, min_size: usize) {
        if config().heap_mode == HeapMode::Static {
            panic!("Error: request_new_alloc should not be called in static heap mode, please set LAMELLAR_HEAP_MODE=dynamic or increase the heap size with LAMELLAR_HEAP_SIZE environment variable");
        }
        // println!("Requesting new pool of size: {} bytes", min_size);
        self.cq.send_alloc(min_size).await;
    }

    async fn send_vec_to_pe_async(&self, pe: usize, vec_data: Vec<u8>) {
        self.cq.send_vec(vec_data, pe).await;
    }

    fn available_to_send(&self, pe: usize) -> bool {
        self.cq.available_to_send(pe)
    }
}

impl Ser for LibfabricSys {
    fn serialize_header(
        &self,
        header: Option<SerializeHeader>,
        serialized_size: usize,
    ) -> Result<SerializedData, anyhow::Error> {
        let header_size = *SERIALIZE_HEADER_LEN;
        let mut ser_data = SerializedData::new(
            self.libfabric_sys_comm.clone(),
            header_size + serialized_size,
        )?;
        crate::serialize_into(&mut ser_data.header_as_bytes_mut(), &header, false)?; //we want header to be a fixed size
        Ok(ser_data)
    }
}