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 .store_arc(Arc::clone(&store))
172 .visibility_store_arc(visibility_store)
173 .scheduler_threads(1)
174 .build()
175 .await?,
176 );
177 Ok(Fixture { engine, store })
178 }
179
180 fn schedule_config(
181 attributes: HashMap<String, SearchAttributeValue>,
182 ) -> Result<ScheduleConfig, aion_core::PayloadError> {
183 Ok(ScheduleConfig {
184 trigger: TriggerSpec::Interval {
185 period: Duration::from_secs(60),
186 },
187 overlap_policy: OverlapPolicy::Skip,
188 catch_up_policy: CatchUpPolicy::Skip,
189 workflow_type: "fixture".to_owned(),
190 input: Payload::from_json(&json!({ "fixture": true }))?,
191 search_attributes: attributes,
192 })
193 }
194
195 fn namespace_attributes(namespace: &str) -> HashMap<String, SearchAttributeValue> {
196 HashMap::from([(
197 NAMESPACE_ATTRIBUTE.to_owned(),
198 SearchAttributeValue::String(namespace.to_owned()),
199 )])
200 }
201
202 fn created_event(
203 engine: &Engine,
204 seq: u64,
205 schedule_id: &ScheduleId,
206 config: ScheduleConfig,
207 ) -> Event {
208 Event::ScheduleCreated {
209 envelope: EventEnvelope {
210 seq,
211 recorded_at: Utc::now(),
212 workflow_id: engine.schedule_coordinator_workflow_id().clone(),
213 },
214 schedule_id: schedule_id.clone(),
215 config,
216 }
217 }
218
219 async fn coordinator_head(fixture: &Fixture) -> Result<u64, Box<dyn std::error::Error>> {
223 let history = fixture
224 .store
225 .read_history(fixture.engine.schedule_coordinator_workflow_id())
226 .await?;
227 Ok(u64::try_from(history.len())?)
228 }
229
230 async fn append_coordinator_events(
231 fixture: &Fixture,
232 events: &[Event],
233 expected_head: u64,
234 ) -> Result<(), Box<dyn std::error::Error>> {
235 fixture
236 .store
237 .append(
238 WriteToken::recorder(),
239 fixture.engine.schedule_coordinator_workflow_id(),
240 events,
241 expected_head,
242 )
243 .await?;
244 Ok(())
245 }
246
247 #[tokio::test]
248 async fn history_source_reads_namespace_from_schedule_created()
249 -> Result<(), Box<dyn std::error::Error>> {
250 let fixture = fixture().await?;
251 let head = coordinator_head(&fixture).await?;
252 let schedule_id = ScheduleId::new(uuid::Uuid::from_u128(1));
253 let event = created_event(
254 &fixture.engine,
255 head + 1,
256 &schedule_id,
257 schedule_config(namespace_attributes("tenant-a"))?,
258 );
259 append_coordinator_events(&fixture, std::slice::from_ref(&event), head).await?;
260 let source = HistoryScheduleNamespaceSource::new(Arc::clone(&fixture.engine));
261
262 assert_eq!(
263 source.schedule_namespace(&schedule_id).await?,
264 Some(String::from("tenant-a"))
265 );
266 Ok(())
267 }
268
269 #[tokio::test]
270 async fn history_source_returns_none_for_unknown_and_unstamped_schedules()
271 -> Result<(), Box<dyn std::error::Error>> {
272 let fixture = fixture().await?;
273 let head = coordinator_head(&fixture).await?;
274 let unstamped = ScheduleId::new(uuid::Uuid::from_u128(2));
275 let unknown = ScheduleId::new(uuid::Uuid::from_u128(3));
276 let event = created_event(
277 &fixture.engine,
278 head + 1,
279 &unstamped,
280 schedule_config(HashMap::new())?,
281 );
282 append_coordinator_events(&fixture, std::slice::from_ref(&event), head).await?;
283 let source = HistoryScheduleNamespaceSource::new(Arc::clone(&fixture.engine));
284
285 assert_eq!(source.schedule_namespace(&unstamped).await?, None);
286 assert_eq!(source.schedule_namespace(&unknown).await?, None);
287 Ok(())
288 }
289
290 #[tokio::test]
291 async fn history_source_rejects_non_string_namespace_attribute()
292 -> Result<(), Box<dyn std::error::Error>> {
293 let fixture = fixture().await?;
294 let head = coordinator_head(&fixture).await?;
295 let schedule_id = ScheduleId::new(uuid::Uuid::from_u128(4));
296 let attributes =
297 HashMap::from([(NAMESPACE_ATTRIBUTE.to_owned(), SearchAttributeValue::Int(7))]);
298 let event = created_event(
299 &fixture.engine,
300 head + 1,
301 &schedule_id,
302 schedule_config(attributes)?,
303 );
304 append_coordinator_events(&fixture, std::slice::from_ref(&event), head).await?;
305 let source = HistoryScheduleNamespaceSource::new(Arc::clone(&fixture.engine));
306
307 let error = source.schedule_namespace(&schedule_id).await;
308
309 assert!(matches!(error, Err(ServerError::Config { .. })));
310 Ok(())
311 }
312
313 #[tokio::test]
314 async fn ownership_is_pinned_to_creation_not_latest_update()
315 -> Result<(), Box<dyn std::error::Error>> {
316 let fixture = fixture().await?;
317 let head = coordinator_head(&fixture).await?;
318 let schedule_id = ScheduleId::new(uuid::Uuid::from_u128(5));
319 let created = created_event(
320 &fixture.engine,
321 head + 1,
322 &schedule_id,
323 schedule_config(namespace_attributes("tenant-a"))?,
324 );
325 let updated = Event::ScheduleUpdated {
326 envelope: EventEnvelope {
327 seq: head + 2,
328 recorded_at: Utc::now(),
329 workflow_id: fixture.engine.schedule_coordinator_workflow_id().clone(),
330 },
331 schedule_id: schedule_id.clone(),
332 config: schedule_config(namespace_attributes("tenant-b"))?,
333 };
334 append_coordinator_events(&fixture, &[created, updated], head).await?;
335 let source = HistoryScheduleNamespaceSource::new(Arc::clone(&fixture.engine));
336
337 assert_eq!(
338 source.schedule_namespace(&schedule_id).await?,
339 Some(String::from("tenant-a"))
340 );
341 Ok(())
342 }
343
344 #[tokio::test]
345 async fn duplicate_creations_pin_the_first_recorded_owner()
346 -> Result<(), Box<dyn std::error::Error>> {
347 let fixture = fixture().await?;
352 let head = coordinator_head(&fixture).await?;
353 let schedule_id = ScheduleId::new(uuid::Uuid::from_u128(8));
354 let first = created_event(
355 &fixture.engine,
356 head + 1,
357 &schedule_id,
358 schedule_config(namespace_attributes("tenant-a"))?,
359 );
360 let second = created_event(
361 &fixture.engine,
362 head + 2,
363 &schedule_id,
364 schedule_config(namespace_attributes("tenant-b"))?,
365 );
366 append_coordinator_events(&fixture, &[first, second], head).await?;
367 let source = HistoryScheduleNamespaceSource::new(Arc::clone(&fixture.engine));
368
369 assert_eq!(
370 source.schedule_namespace(&schedule_id).await?,
371 Some(String::from("tenant-a"))
372 );
373 Ok(())
374 }
375
376 #[tokio::test]
377 async fn deleted_schedules_keep_their_recorded_owner() -> Result<(), Box<dyn std::error::Error>>
378 {
379 let fixture = fixture().await?;
380 let head = coordinator_head(&fixture).await?;
381 let schedule_id = ScheduleId::new(uuid::Uuid::from_u128(6));
382 let created = created_event(
383 &fixture.engine,
384 head + 1,
385 &schedule_id,
386 schedule_config(namespace_attributes("tenant-a"))?,
387 );
388 let deleted = Event::ScheduleDeleted {
389 envelope: EventEnvelope {
390 seq: head + 2,
391 recorded_at: Utc::now(),
392 workflow_id: fixture.engine.schedule_coordinator_workflow_id().clone(),
393 },
394 schedule_id: schedule_id.clone(),
395 };
396 append_coordinator_events(&fixture, &[created, deleted], head).await?;
397 let source = HistoryScheduleNamespaceSource::new(Arc::clone(&fixture.engine));
398
399 assert_eq!(
400 source.schedule_namespace(&schedule_id).await?,
401 Some(String::from("tenant-a"))
402 );
403 Ok(())
404 }
405
406 #[tokio::test]
407 async fn static_source_reports_recorded_namespace() -> Result<(), Box<dyn std::error::Error>> {
408 let ownership = StaticScheduleNamespaces::default();
409 let schedule_id = ScheduleId::new(uuid::Uuid::from_u128(7));
410 ownership.record(schedule_id.clone(), "tenant-a")?;
411
412 assert_eq!(
413 ownership.schedule_namespace(&schedule_id).await?,
414 Some(String::from("tenant-a"))
415 );
416 Ok(())
417 }
418}