1use std::collections::HashMap;
14use std::sync::{Arc, RwLock};
15
16use aion::Engine;
17use aion_core::{Event, ScheduleId, SearchAttributeValue};
18use async_trait::async_trait;
19
20use crate::error::ServerError;
21
22use super::resolver::NAMESPACE_ATTRIBUTE;
23
24#[async_trait]
30pub trait ScheduleNamespaceSource: Send + Sync {
31 async fn schedule_namespace(
39 &self,
40 schedule_id: &ScheduleId,
41 ) -> Result<Option<String>, ServerError>;
42}
43
44pub(crate) struct HistoryScheduleNamespaceSource {
52 engine: Arc<Engine>,
53}
54
55impl HistoryScheduleNamespaceSource {
56 pub(crate) const fn new(engine: Arc<Engine>) -> Self {
58 Self { engine }
59 }
60}
61
62#[async_trait]
63impl ScheduleNamespaceSource for HistoryScheduleNamespaceSource {
64 async fn schedule_namespace(
65 &self,
66 schedule_id: &ScheduleId,
67 ) -> Result<Option<String>, ServerError> {
68 let history = self
73 .engine
74 .store()
75 .read_history(self.engine.schedule_coordinator_workflow_id())
76 .await
77 .map_err(ServerError::from)?;
78 for event in &history {
79 if let Event::ScheduleCreated {
80 schedule_id: created_id,
81 config,
82 ..
83 } = event
84 && created_id == schedule_id
85 {
86 return match config.search_attributes.get(NAMESPACE_ATTRIBUTE) {
88 Some(SearchAttributeValue::String(namespace)) => Ok(Some(namespace.clone())),
89 Some(other) => Err(ServerError::Config {
90 message: format!(
91 "schedule {schedule_id} recorded a non-string {NAMESPACE_ATTRIBUTE} search attribute: {other:?}"
92 ),
93 }),
94 None => Ok(None),
95 };
96 }
97 }
98 Ok(None)
99 }
100}
101
102#[derive(Clone, Default)]
105pub struct StaticScheduleNamespaces {
106 inner: Arc<RwLock<HashMap<ScheduleId, String>>>,
107}
108
109impl StaticScheduleNamespaces {
110 pub fn record(&self, schedule_id: ScheduleId, namespace: &str) -> Result<(), ServerError> {
116 let mut ownership = self
117 .inner
118 .write()
119 .map_err(|_| ServerError::lock_poisoned("namespace schedule ownership"))?;
120 ownership.insert(schedule_id, namespace.to_owned());
121 Ok(())
122 }
123}
124
125#[async_trait]
126impl ScheduleNamespaceSource for StaticScheduleNamespaces {
127 async fn schedule_namespace(
128 &self,
129 schedule_id: &ScheduleId,
130 ) -> Result<Option<String>, ServerError> {
131 let ownership = self
132 .inner
133 .read()
134 .map_err(|_| ServerError::lock_poisoned("namespace schedule ownership"))?;
135 Ok(ownership.get(schedule_id).cloned())
136 }
137}
138
139#[cfg(test)]
140mod tests {
141 use std::collections::HashMap;
142 use std::sync::Arc;
143 use std::time::Duration;
144
145 use aion::{Engine, EngineBuilder};
146 use aion_core::{
147 CatchUpPolicy, Event, EventEnvelope, OverlapPolicy, Payload, ScheduleConfig, ScheduleId,
148 SearchAttributeValue, TriggerSpec,
149 };
150 use aion_store::{EventStore, InMemoryStore, WriteToken, visibility::VisibilityStore};
151 use chrono::Utc;
152 use serde_json::json;
153
154 use super::{
155 HistoryScheduleNamespaceSource, NAMESPACE_ATTRIBUTE, ScheduleNamespaceSource,
156 StaticScheduleNamespaces,
157 };
158 use crate::error::ServerError;
159
160 struct Fixture {
161 engine: Arc<Engine>,
162 store: Arc<dyn EventStore>,
163 }
164
165 async fn fixture() -> Result<Fixture, aion::EngineError> {
166 let backing = Arc::new(InMemoryStore::default());
167 let store: Arc<dyn EventStore> = backing.clone();
168 let visibility_store: Arc<dyn VisibilityStore> = backing;
169 let engine = Arc::new(
170 EngineBuilder::new()
171 .stop_drain_timeout(std::time::Duration::from_secs(5))
172 .store_arc(Arc::clone(&store))
173 .visibility_store_arc(visibility_store)
174 .scheduler_threads(1)
175 .build()
176 .await?,
177 );
178 Ok(Fixture { engine, store })
179 }
180
181 fn schedule_config(
182 attributes: HashMap<String, SearchAttributeValue>,
183 ) -> Result<ScheduleConfig, aion_core::PayloadError> {
184 Ok(ScheduleConfig {
185 trigger: TriggerSpec::Interval {
186 period: Duration::from_secs(60),
187 },
188 overlap_policy: OverlapPolicy::Skip,
189 catch_up_policy: CatchUpPolicy::Skip,
190 workflow_type: "fixture".to_owned(),
191 input: Payload::from_json(&json!({ "fixture": true }))?,
192 search_attributes: attributes,
193 })
194 }
195
196 fn namespace_attributes(namespace: &str) -> HashMap<String, SearchAttributeValue> {
197 HashMap::from([(
198 NAMESPACE_ATTRIBUTE.to_owned(),
199 SearchAttributeValue::String(namespace.to_owned()),
200 )])
201 }
202
203 fn created_event(
204 engine: &Engine,
205 seq: u64,
206 schedule_id: &ScheduleId,
207 config: ScheduleConfig,
208 ) -> Event {
209 Event::ScheduleCreated {
210 envelope: EventEnvelope {
211 seq,
212 recorded_at: Utc::now(),
213 workflow_id: engine.schedule_coordinator_workflow_id().clone(),
214 },
215 schedule_id: schedule_id.clone(),
216 config,
217 }
218 }
219
220 async fn coordinator_head(fixture: &Fixture) -> Result<u64, Box<dyn std::error::Error>> {
224 let history = fixture
225 .store
226 .read_history(fixture.engine.schedule_coordinator_workflow_id())
227 .await?;
228 Ok(u64::try_from(history.len())?)
229 }
230
231 async fn append_coordinator_events(
232 fixture: &Fixture,
233 events: &[Event],
234 expected_head: u64,
235 ) -> Result<(), Box<dyn std::error::Error>> {
236 fixture
237 .store
238 .append(
239 WriteToken::recorder(),
240 fixture.engine.schedule_coordinator_workflow_id(),
241 events,
242 expected_head,
243 )
244 .await?;
245 Ok(())
246 }
247
248 #[tokio::test]
249 async fn history_source_reads_namespace_from_schedule_created()
250 -> Result<(), Box<dyn std::error::Error>> {
251 let fixture = fixture().await?;
252 let head = coordinator_head(&fixture).await?;
253 let schedule_id = ScheduleId::new(uuid::Uuid::from_u128(1));
254 let event = created_event(
255 &fixture.engine,
256 head + 1,
257 &schedule_id,
258 schedule_config(namespace_attributes("tenant-a"))?,
259 );
260 append_coordinator_events(&fixture, std::slice::from_ref(&event), head).await?;
261 let source = HistoryScheduleNamespaceSource::new(Arc::clone(&fixture.engine));
262
263 assert_eq!(
264 source.schedule_namespace(&schedule_id).await?,
265 Some(String::from("tenant-a"))
266 );
267 Ok(())
268 }
269
270 #[tokio::test]
271 async fn history_source_returns_none_for_unknown_and_unstamped_schedules()
272 -> Result<(), Box<dyn std::error::Error>> {
273 let fixture = fixture().await?;
274 let head = coordinator_head(&fixture).await?;
275 let unstamped = ScheduleId::new(uuid::Uuid::from_u128(2));
276 let unknown = ScheduleId::new(uuid::Uuid::from_u128(3));
277 let event = created_event(
278 &fixture.engine,
279 head + 1,
280 &unstamped,
281 schedule_config(HashMap::new())?,
282 );
283 append_coordinator_events(&fixture, std::slice::from_ref(&event), head).await?;
284 let source = HistoryScheduleNamespaceSource::new(Arc::clone(&fixture.engine));
285
286 assert_eq!(source.schedule_namespace(&unstamped).await?, None);
287 assert_eq!(source.schedule_namespace(&unknown).await?, None);
288 Ok(())
289 }
290
291 #[tokio::test]
292 async fn history_source_rejects_non_string_namespace_attribute()
293 -> Result<(), Box<dyn std::error::Error>> {
294 let fixture = fixture().await?;
295 let head = coordinator_head(&fixture).await?;
296 let schedule_id = ScheduleId::new(uuid::Uuid::from_u128(4));
297 let attributes =
298 HashMap::from([(NAMESPACE_ATTRIBUTE.to_owned(), SearchAttributeValue::Int(7))]);
299 let event = created_event(
300 &fixture.engine,
301 head + 1,
302 &schedule_id,
303 schedule_config(attributes)?,
304 );
305 append_coordinator_events(&fixture, std::slice::from_ref(&event), head).await?;
306 let source = HistoryScheduleNamespaceSource::new(Arc::clone(&fixture.engine));
307
308 let error = source.schedule_namespace(&schedule_id).await;
309
310 assert!(matches!(error, Err(ServerError::Config { .. })));
311 Ok(())
312 }
313
314 #[tokio::test]
315 async fn ownership_is_pinned_to_creation_not_latest_update()
316 -> Result<(), Box<dyn std::error::Error>> {
317 let fixture = fixture().await?;
318 let head = coordinator_head(&fixture).await?;
319 let schedule_id = ScheduleId::new(uuid::Uuid::from_u128(5));
320 let created = created_event(
321 &fixture.engine,
322 head + 1,
323 &schedule_id,
324 schedule_config(namespace_attributes("tenant-a"))?,
325 );
326 let updated = Event::ScheduleUpdated {
327 envelope: EventEnvelope {
328 seq: head + 2,
329 recorded_at: Utc::now(),
330 workflow_id: fixture.engine.schedule_coordinator_workflow_id().clone(),
331 },
332 schedule_id: schedule_id.clone(),
333 config: schedule_config(namespace_attributes("tenant-b"))?,
334 };
335 append_coordinator_events(&fixture, &[created, updated], head).await?;
336 let source = HistoryScheduleNamespaceSource::new(Arc::clone(&fixture.engine));
337
338 assert_eq!(
339 source.schedule_namespace(&schedule_id).await?,
340 Some(String::from("tenant-a"))
341 );
342 Ok(())
343 }
344
345 #[tokio::test]
346 async fn duplicate_creations_pin_the_first_recorded_owner()
347 -> Result<(), Box<dyn std::error::Error>> {
348 let fixture = fixture().await?;
353 let head = coordinator_head(&fixture).await?;
354 let schedule_id = ScheduleId::new(uuid::Uuid::from_u128(8));
355 let first = created_event(
356 &fixture.engine,
357 head + 1,
358 &schedule_id,
359 schedule_config(namespace_attributes("tenant-a"))?,
360 );
361 let second = created_event(
362 &fixture.engine,
363 head + 2,
364 &schedule_id,
365 schedule_config(namespace_attributes("tenant-b"))?,
366 );
367 append_coordinator_events(&fixture, &[first, second], head).await?;
368 let source = HistoryScheduleNamespaceSource::new(Arc::clone(&fixture.engine));
369
370 assert_eq!(
371 source.schedule_namespace(&schedule_id).await?,
372 Some(String::from("tenant-a"))
373 );
374 Ok(())
375 }
376
377 #[tokio::test]
378 async fn deleted_schedules_keep_their_recorded_owner() -> Result<(), Box<dyn std::error::Error>>
379 {
380 let fixture = fixture().await?;
381 let head = coordinator_head(&fixture).await?;
382 let schedule_id = ScheduleId::new(uuid::Uuid::from_u128(6));
383 let created = created_event(
384 &fixture.engine,
385 head + 1,
386 &schedule_id,
387 schedule_config(namespace_attributes("tenant-a"))?,
388 );
389 let deleted = Event::ScheduleDeleted {
390 envelope: EventEnvelope {
391 seq: head + 2,
392 recorded_at: Utc::now(),
393 workflow_id: fixture.engine.schedule_coordinator_workflow_id().clone(),
394 },
395 schedule_id: schedule_id.clone(),
396 };
397 append_coordinator_events(&fixture, &[created, deleted], head).await?;
398 let source = HistoryScheduleNamespaceSource::new(Arc::clone(&fixture.engine));
399
400 assert_eq!(
401 source.schedule_namespace(&schedule_id).await?,
402 Some(String::from("tenant-a"))
403 );
404 Ok(())
405 }
406
407 #[tokio::test]
408 async fn static_source_reports_recorded_namespace() -> Result<(), Box<dyn std::error::Error>> {
409 let ownership = StaticScheduleNamespaces::default();
410 let schedule_id = ScheduleId::new(uuid::Uuid::from_u128(7));
411 ownership.record(schedule_id.clone(), "tenant-a")?;
412
413 assert_eq!(
414 ownership.schedule_namespace(&schedule_id).await?,
415 Some(String::from("tenant-a"))
416 );
417 Ok(())
418 }
419}