AppFlowy/frontend/rust-lib/lib-dispatch/src/dispatcher.rs

use crate::runtime::AFPluginRuntime;
use crate::{
    errors::{DispatchError, Error, InternalError},
    module::{as_plugin_map, AFPlugin, AFPluginMap, AFPluginRequest},
    response::AFPluginEventResponse,
    service::{AFPluginServiceFactory, Service},
};
use derivative::*;
use futures_core::future::BoxFuture;
use futures_util::task::Context;
use pin_project::pin_project;
use std::{future::Future, sync::Arc};
use tokio::macros::support::{Pin, Poll};

pub struct AFPluginDispatcher {
    plugins: AFPluginMap,
    runtime: AFPluginRuntime,
}

impl AFPluginDispatcher {
    pub fn construct<F>(runtime: AFPluginRuntime, module_factory: F) -> AFPluginDispatcher
    where
        F: FnOnce() -> Vec<AFPlugin>,
    {
        let plugins = module_factory();
        tracing::trace!("{}", plugin_info(&plugins));
        AFPluginDispatcher {
            plugins: as_plugin_map(plugins),
            runtime,
        }
    }

    pub fn async_send<Req>(dispatch: Arc<AFPluginDispatcher>, request: Req) -> DispatchFuture<AFPluginEventResponse>
    where
        Req: std::convert::Into<AFPluginRequest>,
    {
        AFPluginDispatcher::async_send_with_callback(dispatch, request, |_| Box::pin(async {}))
    }

    pub fn async_send_with_callback<Req, Callback>(
        dispatch: Arc<AFPluginDispatcher>,
        request: Req,
        callback: Callback,
    ) -> DispatchFuture<AFPluginEventResponse>
    where
        Req: std::convert::Into<AFPluginRequest>,
        Callback: FnOnce(AFPluginEventResponse) -> BoxFuture<'static, ()> + 'static + Send + Sync,
    {
        let request: AFPluginRequest = request.into();
        let plugins = dispatch.plugins.clone();
        let service = Box::new(DispatchService { plugins });
        tracing::trace!("Async event: {:?}", &request.event);
        let service_ctx = DispatchContext {
            request,
            callback: Some(Box::new(callback)),
        };
        let join_handle = dispatch.runtime.spawn(async move {
            service.call(service_ctx).await.unwrap_or_else(|e| {
                tracing::error!("Dispatch runtime error: {:?}", e);
                InternalError::Other(format!("{:?}", e)).as_response()
            })
        });

        DispatchFuture {
            fut: Box::pin(async move {
                join_handle.await.unwrap_or_else(|e| {
                    let msg = format!("EVENT_DISPATCH join error: {:?}", e);
                    tracing::error!("{}", msg);
                    let error = InternalError::JoinError(msg);
                    error.as_response()
                })
            }),
        }
    }

    pub fn sync_send(dispatch: Arc<AFPluginDispatcher>, request: AFPluginRequest) -> AFPluginEventResponse {
        futures::executor::block_on(async {
            AFPluginDispatcher::async_send_with_callback(dispatch, request, |_| Box::pin(async {})).await
        })
    }

    pub fn spawn<F>(&self, f: F)
    where
        F: Future<Output = ()> + Send + 'static,
    {
        self.runtime.spawn(f);
    }
}

#[pin_project]
pub struct DispatchFuture<T: Send + Sync> {
    #[pin]
    pub fut: Pin<Box<dyn Future<Output = T> + Sync + Send>>,
}

impl<T> Future for DispatchFuture<T>
where
    T: Send + Sync,
{
    type Output = T;

    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        let this = self.as_mut().project();
        Poll::Ready(futures_core::ready!(this.fut.poll(cx)))
    }
}

pub type BoxFutureCallback = Box<dyn FnOnce(AFPluginEventResponse) -> BoxFuture<'static, ()> + 'static + Send + Sync>;

#[derive(Derivative)]
#[derivative(Debug)]
pub struct DispatchContext {
    pub request: AFPluginRequest,
    #[derivative(Debug = "ignore")]
    pub callback: Option<BoxFutureCallback>,
}

impl DispatchContext {
    pub(crate) fn into_parts(self) -> (AFPluginRequest, Option<BoxFutureCallback>) {
        let DispatchContext { request, callback } = self;
        (request, callback)
    }
}

pub(crate) struct DispatchService {
    pub(crate) plugins: AFPluginMap,
}

impl Service<DispatchContext> for DispatchService {
    type Response = AFPluginEventResponse;
    type Error = DispatchError;
    type Future = BoxFuture<'static, Result<Self::Response, Self::Error>>;

    #[cfg_attr(
        feature = "use_tracing",
        tracing::instrument(name = "DispatchService", level = "debug", skip(self, ctx))
    )]
    fn call(&self, ctx: DispatchContext) -> Self::Future {
        let module_map = self.plugins.clone();
        let (request, callback) = ctx.into_parts();

        Box::pin(async move {
            let result = {
                // print_module_map_info(&module_map);
                match module_map.get(&request.event) {
                    Some(module) => {
                        tracing::trace!("Handle event: {:?} by {:?}", &request.event, module.name);
                        let fut = module.new_service(());
                        let service_fut = fut.await?.call(request);
                        service_fut.await
                    }
                    None => {
                        let msg = format!("Can not find the event handler. {:?}", request);
                        tracing::error!("{}", msg);
                        Err(InternalError::HandleNotFound(msg).into())
                    }
                }
            };

            let response = result.unwrap_or_else(|e| e.into());
            tracing::trace!("Dispatch result: {:?}", response);
            if let Some(callback) = callback {
                callback(response.clone()).await;
            }

            Ok(response)
        })
    }
}

#[allow(dead_code)]
fn plugin_info(plugins: &[AFPlugin]) -> String {
    let mut info = format!("{} plugins loaded\n", plugins.len());
    for module in plugins {
        info.push_str(&format!("-> {} loaded \n", module.name));
    }
    info
}

#[allow(dead_code)]
fn print_plugins(plugins: &AFPluginMap) {
    plugins.iter().for_each(|(k, v)| {
        tracing::info!("Event: {:?} plugin : {:?}", k, v.name);
    })
}