Skip to main content

camel_function/provider/
mod.rs

1use crate::pool::{RunnerHandle, RunnerPoolKey};
2use camel_api::{Exchange, function::*};
3use std::time::Duration;
4
5mod sealed {
6    pub trait Sealed {}
7}
8
9/// Health status of a function provider instance.
10/// Named `FunctionHealthStatus` to avoid collision with `camel_api::FunctionHealthStatus`.
11#[derive(Debug, Clone)]
12pub enum FunctionHealthStatus {
13    Healthy,
14    Unhealthy(String),
15}
16
17#[derive(Debug, thiserror::Error)]
18pub enum ProviderError {
19    #[error("spawn failed: {0}")]
20    SpawnFailed(String),
21    #[error("health check failed: {0}")]
22    HealthFailed(String),
23    #[error("register failed: {0}")]
24    RegisterFailed(String),
25    #[error("unregister failed: {0}")]
26    UnregisterFailed(String),
27    #[error("invoke failed: {0}")]
28    InvokeFailed(String),
29    #[error("shutdown failed: {0}")]
30    ShutdownFailed(String),
31    #[error("boot timeout")]
32    BootTimeout,
33}
34
35#[async_trait::async_trait]
36pub(crate) trait FunctionProvider: Send + Sync + sealed::Sealed {
37    async fn spawn(&self, key: &RunnerPoolKey) -> Result<RunnerHandle, ProviderError>;
38    async fn shutdown(&self, handle: RunnerHandle) -> Result<(), ProviderError>;
39    async fn health(&self, handle: &RunnerHandle) -> Result<FunctionHealthStatus, ProviderError>;
40    async fn register(
41        &self,
42        handle: &RunnerHandle,
43        def: &FunctionDefinition,
44    ) -> Result<(), ProviderError>;
45    async fn unregister(&self, handle: &RunnerHandle, id: &FunctionId)
46    -> Result<(), ProviderError>;
47    async fn invoke(
48        &self,
49        handle: &RunnerHandle,
50        id: &FunctionId,
51        ex: &Exchange,
52        timeout: Duration,
53    ) -> Result<ExchangePatch, ProviderError>;
54}
55
56pub mod container;
57pub mod fake {
58    use super::*;
59    use std::collections::{HashMap, HashSet};
60    use std::sync::atomic::{AtomicUsize, Ordering};
61    use std::sync::{Arc, Mutex};
62    use tokio_util::sync::CancellationToken;
63
64    #[derive(Debug, Clone, Default)]
65    pub struct FakeProviderConfig {
66        pub fail_on_spawn: bool,
67        pub fail_on_register: usize,
68        pub fail_on_health: bool,
69        pub fail_on_shutdown: bool,
70        pub invoke_response: Option<ExchangePatch>,
71        pub invoke_delay: Option<std::time::Duration>,
72    }
73
74    #[derive(Debug, Clone)]
75    pub enum FakeCall {
76        Spawn(RunnerPoolKey),
77        Shutdown(RunnerPoolKey),
78        Health(String),
79        Register(String, FunctionId),
80        Unregister(String, FunctionId),
81        Invoke(String, FunctionId),
82    }
83
84    pub struct FakeProvider {
85        pub config: Arc<Mutex<FakeProviderConfig>>,
86        pub calls: Arc<Mutex<Vec<FakeCall>>>,
87        pub registered: Arc<Mutex<HashMap<String, HashSet<FunctionId>>>>,
88        pub spawned: Arc<Mutex<Vec<RunnerPoolKey>>>,
89        pub shutdowns: Arc<Mutex<Vec<RunnerPoolKey>>>,
90        register_ok_count: Arc<Mutex<usize>>,
91        spawn_count: AtomicUsize,
92    }
93
94    impl FakeProvider {
95        pub fn new(config: FakeProviderConfig) -> Self {
96            Self {
97                config: Arc::new(Mutex::new(config)),
98                calls: Arc::new(Mutex::new(Vec::new())),
99                registered: Arc::new(Mutex::new(HashMap::new())),
100                spawned: Arc::new(Mutex::new(Vec::new())),
101                shutdowns: Arc::new(Mutex::new(Vec::new())),
102                register_ok_count: Arc::new(Mutex::new(0)),
103                spawn_count: AtomicUsize::new(0),
104            }
105        }
106
107        pub fn spawn_count(&self) -> usize {
108            self.spawn_count.load(Ordering::SeqCst)
109        }
110    }
111
112    impl super::sealed::Sealed for FakeProvider {}
113
114    #[async_trait::async_trait]
115    impl FunctionProvider for FakeProvider {
116        async fn spawn(&self, key: &RunnerPoolKey) -> Result<RunnerHandle, ProviderError> {
117            self.spawn_count.fetch_add(1, Ordering::SeqCst);
118            self.calls
119                .lock()
120                .expect("calls") // allow-unwrap
121                .push(FakeCall::Spawn(key.clone()));
122            self.spawned.lock().expect("spawned").push(key.clone()); // allow-unwrap
123            if self.config.lock().expect("config").fail_on_spawn {
124                // allow-unwrap
125                return Err(ProviderError::SpawnFailed("configured".into()));
126            }
127            Ok(RunnerHandle {
128                id: format!("fake-{}", key.runtime),
129                state: Arc::new(Mutex::new(crate::pool::RunnerState::Booting)),
130                cancel: CancellationToken::new(),
131            })
132        }
133
134        async fn shutdown(&self, handle: RunnerHandle) -> Result<(), ProviderError> {
135            self.calls
136                .lock()
137                .expect("calls") // allow-unwrap
138                .push(FakeCall::Shutdown(RunnerPoolKey {
139                    runtime: handle.id.replace("fake-", ""),
140                }));
141            self.shutdowns
142                .lock()
143                .expect("shutdowns") // allow-unwrap
144                .push(RunnerPoolKey {
145                    runtime: handle.id.replace("fake-", ""),
146                });
147            if self.config.lock().expect("config").fail_on_shutdown {
148                // allow-unwrap
149                return Err(ProviderError::ShutdownFailed(
150                    "configured shutdown failure".into(),
151                ));
152            }
153            Ok(())
154        }
155
156        async fn health(
157            &self,
158            handle: &RunnerHandle,
159        ) -> Result<FunctionHealthStatus, ProviderError> {
160            self.calls
161                .lock()
162                .expect("calls") // allow-unwrap
163                .push(FakeCall::Health(handle.id.clone()));
164            if self.config.lock().expect("config").fail_on_health {
165                // allow-unwrap
166                return Ok(FunctionHealthStatus::Unhealthy("configured".into()));
167            }
168            Ok(FunctionHealthStatus::Healthy)
169        }
170
171        async fn register(
172            &self,
173            handle: &RunnerHandle,
174            def: &FunctionDefinition,
175        ) -> Result<(), ProviderError> {
176            self.calls
177                .lock()
178                .expect("calls") // allow-unwrap
179                .push(FakeCall::Register(handle.id.clone(), def.id.clone()));
180            let mut count = self.register_ok_count.lock().expect("count"); // allow-unwrap
181            let cfg = self.config.lock().expect("config").clone(); // allow-unwrap
182            if cfg.fail_on_register > 0 && *count >= cfg.fail_on_register {
183                return Err(ProviderError::RegisterFailed("configured".into()));
184            }
185            *count += 1;
186            self.registered
187                .lock()
188                .expect("registered") // allow-unwrap
189                .entry(handle.id.clone())
190                .or_default()
191                .insert(def.id.clone());
192            Ok(())
193        }
194
195        async fn unregister(
196            &self,
197            handle: &RunnerHandle,
198            id: &FunctionId,
199        ) -> Result<(), ProviderError> {
200            self.calls
201                .lock()
202                .expect("calls") // allow-unwrap
203                .push(FakeCall::Unregister(handle.id.clone(), id.clone()));
204            if let Some(set) = self
205                .registered
206                .lock()
207                .expect("registered") // allow-unwrap
208                .get_mut(&handle.id)
209            {
210                set.remove(id);
211            }
212            Ok(())
213        }
214
215        async fn invoke(
216            &self,
217            handle: &RunnerHandle,
218            id: &FunctionId,
219            _ex: &Exchange,
220            _timeout: Duration,
221        ) -> Result<ExchangePatch, ProviderError> {
222            self.calls
223                .lock()
224                .expect("calls") // allow-unwrap
225                .push(FakeCall::Invoke(handle.id.clone(), id.clone()));
226            let exists = self
227                .registered
228                .lock()
229                .expect("registered") // allow-unwrap
230                .get(&handle.id)
231                .map(|s| s.contains(id))
232                .unwrap_or(false);
233            if !exists {
234                return Err(ProviderError::InvokeFailed("not registered".into()));
235            }
236            let cfg = self.config.lock().expect("config").clone(); // allow-unwrap
237            if let Some(delay) = cfg.invoke_delay {
238                tokio::time::sleep(delay).await;
239            }
240            Ok(cfg.invoke_response.unwrap_or_default())
241        }
242    }
243}