aionbot_core/
runtime.rs

1use std::{cell::UnsafeCell, sync::Arc};
2
3use anyhow::Result;
4use state::TypeMap;
5
6use crate::{entry::Entry, event::Event, handler::Handler, plugin::AionPlugin, types::SetupFn};
7
8#[derive(Default)]
9pub struct StateManager(pub(crate) TypeMap!(Send + Sync));
10
11impl StateManager {
12    pub fn new() -> Self {
13        StateManager(<TypeMap![Send + Sync]>::new())
14    }
15
16    pub fn set<T: Send + Sync + 'static>(&self, state: T) {
17        self.0.set::<T>(state);
18    }
19
20    pub fn get<T: Send + Sync + 'static>(&self) -> &T {
21        self.0.get::<T>()
22    }
23
24    pub fn try_get<T: Send + Sync + 'static>(&self) -> Option<&T> {
25        self.0.try_get::<T>()
26    }
27
28    pub fn get_owned<T: Clone + Send + Sync + 'static>(&self) -> T {
29        self.0.get::<T>().to_owned()
30    }
31}
32
33pub struct Builder<R: Runtime + Default> {
34    handler: UnsafeCell<Handler>,
35    runtime: R,
36    state: Arc<StateManager>,
37    setup: Option<SetupFn<R>>,
38}
39
40impl<R> Builder<R>
41where
42    R: Runtime + Default + Send,
43{
44    pub fn setup(&mut self, setup: SetupFn<R>) {
45        self.setup = Some(setup);
46    }
47
48    pub fn invoke_handler<E: IntoIterator<Item = Entry>>(mut self, entries: E) -> Self {
49        self.handler.get_mut().extend(entries);
50        self
51    }
52
53    pub fn plugin(self, plugin: AionPlugin) -> Self {
54        self.invoke_handler(plugin.entries().to_vec())
55    }
56
57    pub fn manage<T: Send + Sync + 'static>(self, state: T) -> Self {
58        self.state.set(state);
59        self
60    }
61
62    async fn prepare(&mut self) -> Result<()> {
63        log::debug!("Preparing for runtime...");
64        self.runtime.prepare().await?;
65        if let Some(setup) = self.setup.take() {
66            log::debug!("Setting up runtime...");
67            self.runtime.setup(setup);
68        }
69        log::debug!("Finalizing runtime...");
70        self.runtime.finalize().await?;
71        Ok(())
72    }
73
74    pub async fn run(&mut self) -> Result<()> {
75        self.prepare().await?;
76
77        loop {
78            match self.runtime.run().await? {
79                RuntimeStatus::Exit => break,
80                RuntimeStatus::Next => {}
81                RuntimeStatus::Restart => {
82                    log::info!("Restarting bot runtime...");
83                    self.runtime.prepare().await?;
84                }
85                RuntimeStatus::Event(event) => {
86                    let handler = unsafe { self.handler.get().as_mut() }.unwrap();
87                    tokio::spawn(async move {
88                        if let Err(e) = handler.input(Arc::new(event)).await {
89                            log::error!("Error handling event: {}", e);
90                        };
91                    });
92                }
93            }
94        }
95        Ok(())
96    }
97}
98
99impl<R> Default for Builder<R>
100where
101    R: Runtime + Default + Send + 'static,
102{
103    fn default() -> Self {
104        let manager = Arc::new(StateManager::new());
105        let runtime = R::default().set_manager(manager.clone());
106        Self {
107            handler: UnsafeCell::new(Handler::empty()),
108            runtime,
109            state: Arc::clone(&manager),
110            setup: None,
111        }
112    }
113}
114
115pub trait Runtime {
116    #[must_use]
117    fn set_manager(self, manager: Arc<StateManager>) -> Self;
118
119    fn manager(&self) -> &StateManager;
120
121    fn prepare(&mut self) -> impl std::future::Future<Output = Result<()>> + Send {
122        async move { Ok(()) }
123    }
124
125    fn setup(&mut self, setup: SetupFn<Self>) {
126        setup(self)
127    }
128
129    fn finalize(&mut self) -> impl std::future::Future<Output = Result<()>> + Send {
130        async move { Ok(()) }
131    }
132
133    fn run(&mut self) -> impl std::future::Future<Output = Result<RuntimeStatus>> + Send;
134}
135
136pub enum RuntimeStatus {
137    Next,
138    Exit,
139    Restart,
140    Event(Box<dyn Event>),
141}
aionbot_core/runtime.rs

aionbot_core/
runtime.rs
