cordis_core/

context.rs

use std::{
    any::TypeId,
    collections::HashMap,
    future::Future,
    marker::PhantomData,
    ops::Deref,
    sync::{Arc, Weak},
};

use futures::FutureExt;
use tokio_util::sync::CancellationToken;

use crate::{
    Error, Result, ServiceHandle, ServiceKey,
    runtime::{EventCallback, QueryCallback, Runtime},
    scope::ScopeInner,
    service::{ServiceEntry, ServiceId, boxed_service},
};

/// Cheaply clonable access to the shared runtime.
#[derive(Clone)]
pub struct Context {
    pub(crate) runtime: Arc<Runtime>,
    scope_id: u64,
    isolations: Arc<HashMap<TypeId, u64>>,
}

impl Context {
    pub(crate) fn root(runtime: Arc<Runtime>) -> Self {
        Self {
            runtime,
            scope_id: 0,
            isolations: Arc::new(HashMap::new()),
        }
    }

    pub(crate) fn for_scope(runtime: Arc<Runtime>, scope_id: u64) -> Self {
        Self {
            runtime,
            scope_id,
            isolations: Arc::new(HashMap::new()),
        }
    }

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

    /// Creates a lightweight child context sharing the same runtime.
    pub fn child(&self) -> Self {
        Self {
            runtime: self.runtime.clone(),
            scope_id: self.runtime.next_id(),
            isolations: self.isolations.clone(),
        }
    }

    /// Creates a child where service `K` resolves in a fresh isolated slot;
    /// all other service types continue to inherit their current slots.
    pub fn isolate<K: ServiceKey>(&self) -> Self {
        let mut isolations = (*self.isolations).clone();
        isolations.insert(TypeId::of::<K>(), self.runtime.next_id());
        Self {
            runtime: self.runtime.clone(),
            scope_id: self.scope_id,
            isolations: Arc::new(isolations),
        }
    }

    fn service_id<K: ServiceKey>(&self) -> ServiceId {
        ServiceId {
            key: TypeId::of::<K>(),
            isolation: self
                .isolations
                .get(&TypeId::of::<K>())
                .copied()
                .unwrap_or(0),
        }
    }

    pub fn get<K: ServiceKey>(&self) -> Result<Arc<K::Value>> {
        let services = self.runtime.services.lock().expect("service lock poisoned");
        let entry = services
            .get(&self.service_id::<K>())
            .filter(|entry| entry.active || entry.owner == self.scope_id)
            .ok_or(Error::MissingService { name: K::NAME })?;
        entry
            .value
            .downcast_ref::<Arc<K::Value>>()
            .cloned()
            .ok_or(Error::ServiceTypeMismatch { name: entry.name })
    }

    pub fn try_get<K: ServiceKey>(&self) -> Option<Arc<K::Value>> {
        let services = self.runtime.services.lock().expect("service lock poisoned");
        let entry = services.get(&self.service_id::<K>())?;
        if !entry.active && entry.owner != self.scope_id {
            return None;
        }
        entry.value.downcast_ref::<Arc<K::Value>>().cloned()
    }

    /// Invokes listeners serially in registration order.
    pub async fn emit<E: Event>(&self, event: E) -> Result<()> {
        self.runtime.emit_serial(event).await
    }

    /// Invokes all listeners concurrently.
    pub async fn parallel<E: Event>(&self, event: E) -> Result<()> {
        self.runtime.emit_parallel(event).await
    }

    /// Runs query handlers in registration order and returns the first answer.
    pub async fn query<Q: Query>(&self, query: Q) -> Result<Option<Q::Response>> {
        self.runtime.query(query).await
    }
}

#[derive(Clone, Copy, Debug)]
pub struct Ready;

#[derive(Clone, Copy, Debug)]
pub struct Fork {
    pub plugin: crate::PluginId,
    pub activation: crate::ActivationId,
}

#[derive(Clone, Copy, Debug)]
pub struct Dispose {
    pub plugin: crate::PluginId,
    pub activation: crate::ActivationId,
}

/// Any thread-safe static value can be used as an event payload.
pub trait Event: Send + Sync + 'static {}
impl<T: Send + Sync + 'static> Event for T {}

/// A typed bail/query event.
pub trait Query: Send + Sync + 'static {
    type Response: Send + Sync + 'static;
}

/// Context passed to a plugin. Registrations are automatically owned by its
/// scope and survive when their returned handles are ignored.
#[derive(Clone)]
pub struct PluginContext {
    context: Context,
    scope: Arc<ScopeInner>,
}

impl PluginContext {
    pub(crate) fn new(context: Context, scope: Arc<ScopeInner>) -> Self {
        Self { context, scope }
    }

    pub fn isolate<K: ServiceKey>(&self) -> Self {
        Self {
            context: self.context.isolate::<K>(),
            scope: self.scope.clone(),
        }
    }

    pub fn provide<K: ServiceKey>(&self, value: Arc<K::Value>) -> Result<ServiceHandle<K>> {
        let id = self.context.service_id::<K>();
        let token = self.context.runtime.next_service_generation();
        let generation = self.context.runtime.next_service_generation();
        {
            let mut services = self
                .context
                .runtime
                .services
                .lock()
                .expect("service lock poisoned");
            if services.contains_key(&id) {
                return Err(Error::DuplicateService { name: K::NAME });
            }
            services.insert(
                id,
                ServiceEntry {
                    value: boxed_service::<K>(value),
                    owner: self.scope.id,
                    token,
                    generation,
                    name: K::NAME,
                    active: false,
                },
            );
        }

        let weak = Arc::downgrade(&self.context.runtime);
        let owner = self.scope.id;
        if let Err(error) = self.scope.push(Box::new(move || {
            Box::pin(async move {
                if let Some(runtime) = weak.upgrade() {
                    runtime.remove_service(id, owner, token);
                }
                Ok(())
            })
        })) {
            self.context.runtime.remove_service(id, owner, token);
            return Err(error);
        }

        Ok(ServiceHandle {
            id,
            runtime: Arc::downgrade(&self.context.runtime),
            owner,
            token,
            _key: PhantomData,
        })
    }

    pub fn on<E, H, Fut>(&self, handler: H) -> Result<ListenerHandle>
    where
        E: Event,
        H: Fn(Context, Arc<E>) -> Fut + Send + Sync + 'static,
        Fut: Future<Output = Result<()>> + Send + 'static,
    {
        let event_key = TypeId::of::<E>();
        let context = self.context.clone();
        let callback: Arc<EventCallback> = Arc::new(move |event| {
            let result = event.downcast::<E>();
            let context = context.clone();
            match result {
                Ok(event) => handler(context, event).boxed(),
                Err(_) => {
                    async { Err(Error::Cleanup("event payload type mismatch".into())) }.boxed()
                }
            }
        });
        let id = self
            .context
            .runtime
            .add_listener(event_key, self.scope.id, callback);
        let weak = Arc::downgrade(&self.context.runtime);
        if let Err(error) = self.scope.push(Box::new(move || {
            Box::pin(async move {
                if let Some(runtime) = weak.upgrade() {
                    runtime.remove_listener(event_key, id);
                }
                Ok(())
            })
        })) {
            self.context.runtime.remove_listener(event_key, id);
            return Err(error);
        }
        Ok(ListenerHandle {
            runtime: Arc::downgrade(&self.context.runtime),
            key: event_key,
            id,
            query: false,
        })
    }

    pub fn on_query<Q, H, Fut>(&self, handler: H) -> Result<ListenerHandle>
    where
        Q: Query,
        H: Fn(Context, Arc<Q>) -> Fut + Send + Sync + 'static,
        Fut: Future<Output = Result<Option<Q::Response>>> + Send + 'static,
    {
        let key = TypeId::of::<Q>();
        let context = self.context.clone();
        let callback: Arc<QueryCallback> = Arc::new(move |query| {
            let result = query.downcast::<Q>();
            let context = context.clone();
            match result {
                Ok(query) => handler(context, query)
                    .map(|result| {
                        result.map(|response| {
                            response.map(|value| {
                                Box::new(value) as Box<dyn std::any::Any + Send + Sync>
                            })
                        })
                    })
                    .boxed(),
                Err(_) => {
                    async { Err(Error::Cleanup("query payload type mismatch".into())) }.boxed()
                }
            }
        });
        let id = self
            .context
            .runtime
            .add_query_listener(key, self.scope.id, callback);
        let weak = Arc::downgrade(&self.context.runtime);
        if let Err(error) = self.scope.push(Box::new(move || {
            Box::pin(async move {
                if let Some(runtime) = weak.upgrade() {
                    runtime.remove_query_listener(key, id);
                }
                Ok(())
            })
        })) {
            self.context.runtime.remove_query_listener(key, id);
            return Err(error);
        }
        Ok(ListenerHandle {
            runtime: Arc::downgrade(&self.context.runtime),
            key,
            id,
            query: true,
        })
    }

    pub fn defer<F, Fut>(&self, cleanup: F) -> Result<()>
    where
        F: FnOnce() -> Fut + Send + 'static,
        Fut: Future<Output = Result<()>> + Send + 'static,
    {
        self.scope.push(Box::new(move || Box::pin(cleanup())))
    }

    pub fn manage<R: Resource>(&self, resource: R) -> Result<()> {
        let resource = Arc::new(std::sync::Mutex::new(Some(resource)));
        let start_resource = resource.clone();
        self.scope.on_commit(Box::new(move || {
            Box::pin(async move {
                let resource = start_resource
                    .lock()
                    .expect("resource lock poisoned")
                    .take()
                    .ok_or_else(|| Error::cleanup("resource already consumed"))?;
                let result = resource.start().await;
                *start_resource.lock().expect("resource lock poisoned") = Some(resource);
                result
            })
        }))?;

        self.scope.push(Box::new(move || {
            Box::pin(async move {
                let resource = resource.lock().expect("resource lock poisoned").take();
                if let Some(resource) = resource {
                    resource.cancel();
                    Box::new(resource).dispose().await?;
                }
                Ok(())
            })
        }))
    }

    /// Registers a cooperative task. It is spawned only when the activation
    /// commits. Disposal waits five seconds, then aborts an unresponsive task.
    pub fn spawn<F, Fut>(&self, task: F) -> Result<TaskHandle>
    where
        F: FnOnce(CancellationToken) -> Fut + Send + 'static,
        Fut: Future<Output = Result<()>> + Send + 'static,
    {
        let token = CancellationToken::new();
        let join = Arc::new(std::sync::Mutex::new(None));
        let start_join = join.clone();
        let child = token.child_token();
        self.scope.on_commit(Box::new(move || {
            Box::pin(async move {
                let future = std::panic::AssertUnwindSafe(task(child))
                    .catch_unwind()
                    .map(|result| match result {
                        Ok(result) => result,
                        Err(payload) => Err(Error::panic(payload)),
                    });
                *start_join.lock().expect("task lock poisoned") = Some(tokio::spawn(future));
                Ok(())
            })
        }))?;

        let cleanup_token = token.clone();
        self.scope.push(Box::new(move || {
            Box::pin(async move {
                cleanup_token.cancel();
                let join = join.lock().expect("task lock poisoned").take();
                if let Some(mut join) = join {
                    match tokio::time::timeout(std::time::Duration::from_secs(5), &mut join).await {
                        Ok(result) => result??,
                        Err(_) => {
                            join.abort();
                            let _ = join.await;
                            return Err(Error::TaskTimeout { seconds: 5 });
                        }
                    }
                }
                Ok(())
            })
        }))?;
        Ok(TaskHandle { token })
    }
}

impl Deref for PluginContext {
    type Target = Context;
    fn deref(&self) -> &Self::Target {
        &self.context
    }
}

/// Non-owning early-cancellation handle. `Drop` intentionally does nothing.
pub struct ListenerHandle {
    runtime: Weak<Runtime>,
    key: TypeId,
    id: u64,
    query: bool,
}

impl ListenerHandle {
    pub fn cancel(&self) -> bool {
        self.runtime.upgrade().is_some_and(|runtime| {
            if self.query {
                runtime.remove_query_listener(self.key, self.id)
            } else {
                runtime.remove_listener(self.key, self.id)
            }
        })
    }
}

pub struct TaskHandle {
    token: CancellationToken,
}
impl TaskHandle {
    pub fn cancel(&self) {
        self.token.cancel();
    }
}

pub trait Resource: Send + Sync + 'static {
    fn start(&self) -> impl Future<Output = Result<()>> + Send {
        async { Ok(()) }
    }

    fn cancel(&self) {}

    fn dispose(self: Box<Self>) -> impl Future<Output = Result<()>> + Send;
}