1use std::collections::{HashMap, HashSet};
10use std::sync::atomic::{AtomicBool, Ordering};
11use std::sync::{Arc, Condvar, Mutex};
12use std::time::{Duration, SystemTime};
13
14use dex_protocol::dex::{
15 AttributeWrite, ChannelInfo, ChannelMessage, ConditionResults, ConditionStatus,
16 Context as ProtoContext, Kv, Value as ProtoValue, value,
17};
18
19use crate::persistence::PersistenceKind;
20use crate::registry::{RegisteredFlow, decode_instance, physical_name};
21use crate::value_mapper;
22use crate::{
23 Attribute, AttributeMap, Channel, ChannelMap, FlowResult, HandlerError, HandlerResult, Value,
24};
25
26#[derive(Clone, Copy, Eq, PartialEq)]
27pub(crate) enum InvocationMethod {
28 WaitFor,
29 Execute,
30 Rpc,
31}
32
33pub struct Context {
39 method: InvocationMethod,
40 flow: RegisteredFlow,
41 metadata: ProtoContext,
42 attributes: HashMap<String, ProtoValue>,
43 locals: HashMap<String, ProtoValue>,
44 condition_results: Option<ConditionResults>,
45 channel_infos: HashMap<String, ChannelInfo>,
46 attribute_writes: HashMap<String, AttributeWrite>,
47 local_writes: HashMap<String, Kv>,
48 events: Vec<Kv>,
49 event_names: HashSet<String>,
50 publications: Vec<ChannelMessage>,
51 cancellation: InvocationCancellation,
52}
53
54impl Context {
55 pub fn flow_id(&self) -> &str {
57 &self.metadata.flow_id
58 }
59
60 pub fn run_id(&self) -> &str {
62 &self.metadata.run_id
63 }
64
65 pub fn flow_started_at(&self) -> SystemTime {
67 system_time(self.metadata.flow_started_timestamp)
68 }
69
70 pub fn step_execution_id(&self) -> &str {
72 &self.metadata.step_execution_id
73 }
74
75 pub fn from_step_execution_id(&self) -> &str {
77 &self.metadata.from_step_execution_id
78 }
79
80 pub fn first_attempt_at(&self) -> SystemTime {
82 system_time(self.metadata.first_attempt_timestamp)
83 }
84
85 pub fn attempt(&self) -> u32 {
87 u32::try_from(self.metadata.attempt).unwrap_or_default()
88 }
89
90 pub fn has_any_timer_fired(&self) -> bool {
92 self.condition_results.as_ref().is_some_and(|results| {
93 results
94 .timer_results
95 .iter()
96 .any(|timer| timer.condition_status == ConditionStatus::Completed as i32)
97 })
98 }
99
100 pub fn has_timer_fired(&self, index: usize) -> bool {
104 self.condition_results
105 .as_ref()
106 .and_then(|results| results.timer_results.get(index))
107 .is_some_and(|timer| timer.condition_status == ConditionStatus::Completed as i32)
108 }
109
110 pub fn wait_for_method_failed(&self) -> bool {
112 self.condition_results
113 .as_ref()
114 .is_some_and(|results| results.wait_for_failed)
115 }
116
117 pub(crate) fn sub_flow_result(&self, index: usize) -> HandlerResult<FlowResult> {
118 if self.method != InvocationMethod::Execute {
119 return Err(HandlerError::new(
120 "SubFlow results are available only during execute",
121 ));
122 }
123 let result = self
124 .condition_results
125 .as_ref()
126 .and_then(|results| results.sub_flow_results.get(index))
127 .ok_or_else(|| {
128 HandlerError::new(format!("SubFlow result index is unavailable: {index}"))
129 })?;
130 FlowResult::from_proto(result).map_err(|error| HandlerError::new(error.to_string()))
131 }
132
133 pub(crate) fn sub_flow_id(&self, index: usize) -> HandlerResult<String> {
134 self.sub_flow_result(index)?;
135 Ok(format!(
136 "SubFlow:{}-{}-{index}",
137 self.flow_id(),
138 self.step_execution_id()
139 ))
140 }
141
142 pub fn wait_for_cancellation(&self) {
146 self.cancellation.wait();
147 }
148
149 pub fn is_cancelled(&self) -> bool {
151 self.cancellation.is_cancelled()
152 }
153
154 pub fn set_step_execution_local<T: Value>(&mut self, key: &str, value: T) -> HandlerResult<()> {
160 require_name(key, "step-execution local key")?;
161 let value = value_mapper::encode_handler(&value)?;
162 self.local_writes.insert(
163 key.to_string(),
164 Kv {
165 key: key.to_string(),
166 value: Some(value),
167 },
168 );
169 Ok(())
170 }
171
172 pub fn step_execution_local<T: Value>(&self, key: &str) -> HandlerResult<T> {
178 require_name(key, "step-execution local key")?;
179 let value = self
180 .local_writes
181 .get(key)
182 .and_then(|entry| entry.value.as_ref())
183 .or_else(|| self.locals.get(key))
184 .ok_or_else(|| HandlerError::new(format!("step-execution local {key} is missing")))?;
185 value_mapper::decode_handler(value)
186 }
187
188 pub fn record_event<T: Value>(&mut self, name: &str, value: T) -> HandlerResult<()> {
196 require_name(name, "event name")?;
197 if !self.event_names.insert(name.to_string()) {
198 return Err(HandlerError::new(format!(
199 "event was already recorded: {name}"
200 )));
201 }
202 self.events.push(Kv {
203 key: name.to_string(),
204 value: Some(value_mapper::encode_handler(&value)?),
205 });
206 Ok(())
207 }
208
209 pub fn get_attribute<T: Value>(&self, attribute: &Attribute<T>) -> HandlerResult<Option<T>> {
211 self.get_attribute_value(attribute.name(), PersistenceKind::Attribute, None)
212 }
213
214 pub fn get_attribute_map<T: Value>(
216 &self,
217 attribute: &AttributeMap<T>,
218 instance: &str,
219 ) -> HandlerResult<Option<T>> {
220 self.get_attribute_value(
221 attribute.name(),
222 PersistenceKind::AttributeMap,
223 Some(instance),
224 )
225 }
226
227 pub fn set_attribute<T: Value>(
229 &mut self,
230 attribute: &Attribute<T>,
231 value: T,
232 ) -> HandlerResult<()> {
233 self.set_attribute_value(
234 attribute.name(),
235 PersistenceKind::Attribute,
236 None,
237 value_mapper::encode_handler(&value)?,
238 attribute.index().map(|index| index.proto_config(false)),
239 attribute.sync_config(),
240 )
241 }
242
243 pub fn set_attribute_map<T: Value>(
245 &mut self,
246 attribute: &AttributeMap<T>,
247 instance: &str,
248 value: T,
249 ) -> HandlerResult<()> {
250 self.set_attribute_value(
251 attribute.name(),
252 PersistenceKind::AttributeMap,
253 Some(instance),
254 value_mapper::encode_handler(&value)?,
255 attribute.index().map(|index| index.proto_config(true)),
256 attribute.sync_config(),
257 )
258 }
259
260 pub fn delete_attribute<T>(&mut self, attribute: &Attribute<T>) -> HandlerResult<()> {
262 self.set_attribute_value(
263 attribute.name(),
264 PersistenceKind::Attribute,
265 None,
266 value_mapper::deletion(),
267 attribute.index().map(|index| index.proto_config(false)),
268 attribute.sync_config(),
269 )
270 }
271
272 pub fn delete_attribute_map<T>(
274 &mut self,
275 attribute: &AttributeMap<T>,
276 instance: &str,
277 ) -> HandlerResult<()> {
278 self.set_attribute_value(
279 attribute.name(),
280 PersistenceKind::AttributeMap,
281 Some(instance),
282 value_mapper::deletion(),
283 attribute.index().map(|index| index.proto_config(true)),
284 attribute.sync_config(),
285 )
286 }
287
288 pub(crate) fn attribute_map_keys<T>(
289 &self,
290 attribute: &AttributeMap<T>,
291 ) -> HandlerResult<Vec<String>> {
292 self.registered_name(
293 attribute.name(),
294 PersistenceKind::AttributeMap,
295 Some("key-check"),
296 )?;
297 let prefix = format!("{}/", attribute.name());
298 let mut physical_keys = self
299 .attributes
300 .keys()
301 .filter(|key| key.starts_with(&prefix))
302 .cloned()
303 .collect::<HashSet<_>>();
304 for (key, write) in &self.attribute_writes {
305 if !key.starts_with(&prefix) {
306 continue;
307 }
308 if matches!(
309 write.value.as_ref().and_then(|value| value.kind.as_ref()),
310 Some(value::Kind::NullValue(_))
311 ) {
312 physical_keys.remove(key);
313 } else {
314 physical_keys.insert(key.clone());
315 }
316 }
317 sorted_instance_keys(&prefix, physical_keys)
318 }
319
320 pub fn publish<T: Value>(&mut self, channel: &Channel<T>, value: T) -> HandlerResult<()> {
322 self.publish_value(
323 channel.name(),
324 PersistenceKind::Channel,
325 None,
326 value_mapper::encode_handler(&value)?,
327 )
328 }
329
330 pub fn publish_map<T: Value>(
332 &mut self,
333 channel: &ChannelMap<T>,
334 instance: &str,
335 value: T,
336 ) -> HandlerResult<()> {
337 self.publish_value(
338 channel.name(),
339 PersistenceKind::ChannelMap,
340 Some(instance),
341 value_mapper::encode_handler(&value)?,
342 )
343 }
344
345 pub fn channel_size<T>(&self, channel: &Channel<T>) -> HandlerResult<usize> {
347 self.channel_size_value(channel.name(), PersistenceKind::Channel, None)
348 }
349
350 pub fn channel_map_size<T>(
352 &self,
353 channel: &ChannelMap<T>,
354 instance: &str,
355 ) -> HandlerResult<usize> {
356 self.channel_size_value(channel.name(), PersistenceKind::ChannelMap, Some(instance))
357 }
358
359 pub(crate) fn channel_map_keys<T>(
360 &self,
361 channel: &ChannelMap<T>,
362 ) -> HandlerResult<Vec<String>> {
363 if self.method != InvocationMethod::Rpc {
364 return Err(HandlerError::new(
365 "ChannelMap introspection requires an RPC invocation",
366 ));
367 }
368 self.registered_name(
369 channel.name(),
370 PersistenceKind::ChannelMap,
371 Some("key-check"),
372 )?;
373 let prefix = format!("{}/", channel.name());
374 let physical_keys = self
375 .channel_infos
376 .iter()
377 .filter(|(key, info)| key.starts_with(&prefix) && info.size > 0)
378 .map(|(key, _)| key.clone())
379 .collect::<HashSet<_>>();
380 sorted_instance_keys(&prefix, physical_keys)
381 }
382
383 pub fn channel_results<T: Value>(&self, channel: &Channel<T>) -> HandlerResult<Vec<T>> {
385 self.channel_results_value(channel.name(), PersistenceKind::Channel, None)
386 }
387
388 pub fn channel_map_results<T: Value>(
390 &self,
391 channel: &ChannelMap<T>,
392 instance: &str,
393 ) -> HandlerResult<Vec<T>> {
394 self.channel_results_value(channel.name(), PersistenceKind::ChannelMap, Some(instance))
395 }
396
397 pub(crate) fn new(
398 method: InvocationMethod,
399 flow: RegisteredFlow,
400 metadata: ProtoContext,
401 attributes: Vec<Kv>,
402 locals: Vec<Kv>,
403 condition_results: Option<ConditionResults>,
404 channel_infos: HashMap<String, ChannelInfo>,
405 ) -> HandlerResult<Self> {
406 Ok(Self {
407 method,
408 flow,
409 metadata,
410 attributes: map_values("Attribute", attributes)?,
411 locals: map_values("step-execution local", locals)?,
412 condition_results,
413 channel_infos,
414 attribute_writes: HashMap::new(),
415 local_writes: HashMap::new(),
416 events: Vec::new(),
417 event_names: HashSet::new(),
418 publications: Vec::new(),
419 cancellation: InvocationCancellation::new(),
420 })
421 }
422
423 pub(crate) fn cancellation(&self) -> InvocationCancellation {
424 self.cancellation.clone()
425 }
426
427 pub(crate) fn take_outputs(
428 self,
429 ) -> (Vec<AttributeWrite>, Vec<Kv>, Vec<Kv>, Vec<ChannelMessage>) {
430 (
431 self.attribute_writes.into_values().collect(),
432 self.local_writes.into_values().collect(),
433 self.events,
434 self.publications,
435 )
436 }
437
438 fn get_attribute_value<T: Value>(
439 &self,
440 name: &str,
441 kind: PersistenceKind,
442 instance: Option<&str>,
443 ) -> HandlerResult<Option<T>> {
444 let key = self.registered_name(name, kind, instance)?;
445 if let Some(write) = self.attribute_writes.get(&key) {
446 let value = write
447 .value
448 .as_ref()
449 .ok_or_else(|| HandlerError::new(format!("Attribute write {key} has no Value")))?;
450 if matches!(value.kind, Some(value::Kind::NullValue(_))) {
451 return Ok(None);
452 }
453 return value_mapper::decode_handler(value).map(Some);
454 }
455 self.attributes
456 .get(&key)
457 .map(value_mapper::decode_handler)
458 .transpose()
459 }
460
461 fn set_attribute_value(
462 &mut self,
463 name: &str,
464 kind: PersistenceKind,
465 instance: Option<&str>,
466 value: ProtoValue,
467 index_config: Option<dex_protocol::dex::IndexConfig>,
468 sync_config: Option<dex_protocol::dex::AttributeSyncConfig>,
469 ) -> HandlerResult<()> {
470 let key = self.registered_name(name, kind, instance)?;
471 self.attribute_writes.insert(
472 key.clone(),
473 AttributeWrite {
474 key,
475 value: Some(value),
476 index_config,
477 sync_config,
478 },
479 );
480 Ok(())
481 }
482
483 fn publish_value(
484 &mut self,
485 name: &str,
486 kind: PersistenceKind,
487 instance: Option<&str>,
488 value: ProtoValue,
489 ) -> HandlerResult<()> {
490 let channel_name = self.registered_name(name, kind, instance)?;
491 self.publications.push(ChannelMessage {
492 channel_name: channel_name.clone(),
493 value: Some(value),
494 });
495 if self.method == InvocationMethod::Rpc {
496 self.channel_infos
497 .entry(channel_name)
498 .and_modify(|info| info.size += 1)
499 .or_insert(ChannelInfo { size: 1 });
500 }
501 Ok(())
502 }
503
504 fn channel_size_value(
505 &self,
506 name: &str,
507 kind: PersistenceKind,
508 instance: Option<&str>,
509 ) -> HandlerResult<usize> {
510 let channel_name = self.registered_name(name, kind, instance)?;
511 Ok(self
512 .channel_infos
513 .get(&channel_name)
514 .map_or(0, |info| usize::try_from(info.size).unwrap_or_default()))
515 }
516
517 fn channel_results_value<T: Value>(
518 &self,
519 name: &str,
520 kind: PersistenceKind,
521 instance: Option<&str>,
522 ) -> HandlerResult<Vec<T>> {
523 let channel_name = self.registered_name(name, kind, instance)?;
524 let Some(results) = &self.condition_results else {
525 return Ok(Vec::new());
526 };
527 let mut decoded = Vec::new();
528 for result in &results.channel_results {
529 if result.channel_name == channel_name
530 && result.condition_status == ConditionStatus::Completed as i32
531 {
532 for value in &result.values {
533 decoded.push(value_mapper::decode_handler(value)?);
534 }
535 }
536 }
537 Ok(decoded)
538 }
539
540 fn registered_name(
541 &self,
542 name: &str,
543 kind: PersistenceKind,
544 instance: Option<&str>,
545 ) -> HandlerResult<String> {
546 let definition = self.flow.persistence.get(name).ok_or_else(|| {
547 HandlerError::new(format!(
548 "persistence definition does not belong to Flow: {name}"
549 ))
550 })?;
551 if definition.kind != kind {
552 return Err(HandlerError::new(format!(
553 "persistence definition kind does not match: {name}"
554 )));
555 }
556 match kind {
557 PersistenceKind::AttributeMap | PersistenceKind::ChannelMap => instance
558 .map(|instance| physical_name(name, instance))
559 .ok_or_else(|| HandlerError::new(format!("{name} requires an instance"))),
560 PersistenceKind::Attribute | PersistenceKind::Channel if instance.is_none() => {
561 Ok(name.to_string())
562 }
563 _ => Err(HandlerError::new(format!(
564 "static persistence definition cannot use an instance: {name}"
565 ))),
566 }
567 }
568}
569
570fn sorted_instance_keys(
571 prefix: &str,
572 physical_keys: HashSet<String>,
573) -> HandlerResult<Vec<String>> {
574 let mut keys = physical_keys
575 .iter()
576 .map(|key| decode_instance(key, prefix).map_err(HandlerError::new))
577 .collect::<HandlerResult<Vec<_>>>()?;
578 keys.sort();
579 Ok(keys)
580}
581
582#[derive(Clone)]
583pub(crate) struct InvocationCancellation {
584 inner: Arc<InvocationCancellationInner>,
585}
586
587struct InvocationCancellationInner {
588 cancelled: AtomicBool,
589 mutex: Mutex<()>,
590 condition: Condvar,
591}
592
593impl InvocationCancellation {
594 fn new() -> Self {
595 Self {
596 inner: Arc::new(InvocationCancellationInner {
597 cancelled: AtomicBool::new(false),
598 mutex: Mutex::new(()),
599 condition: Condvar::new(),
600 }),
601 }
602 }
603
604 fn wait(&self) {
605 let mut guard = self.inner.mutex.lock().expect("cancellation lock");
606 while !self.is_cancelled() {
607 guard = self.inner.condition.wait(guard).expect("cancellation lock");
608 }
609 }
610
611 fn is_cancelled(&self) -> bool {
612 self.inner.cancelled.load(Ordering::Acquire)
613 }
614
615 pub(crate) fn cancel(&self) {
616 let _guard = self.inner.mutex.lock().expect("cancellation lock");
617 self.inner.cancelled.store(true, Ordering::Release);
618 self.inner.condition.notify_all();
619 }
620}
621
622fn system_time(timestamp: i64) -> SystemTime {
623 if timestamp >= 0 {
624 SystemTime::UNIX_EPOCH + Duration::from_secs(timestamp as u64)
625 } else {
626 SystemTime::UNIX_EPOCH - Duration::from_secs(timestamp.unsigned_abs())
627 }
628}
629
630fn map_values(kind: &str, entries: Vec<Kv>) -> HandlerResult<HashMap<String, ProtoValue>> {
631 let mut mapped = HashMap::new();
632 for entry in entries {
633 require_name(&entry.key, kind)?;
634 let value = entry
635 .value
636 .ok_or_else(|| HandlerError::new(format!("{kind} {} has no Value", entry.key)))?;
637 if mapped.insert(entry.key.clone(), value).is_some() {
638 return Err(HandlerError::new(format!("duplicate {kind} {}", entry.key)));
639 }
640 }
641 Ok(mapped)
642}
643
644fn require_name(value: &str, kind: &str) -> HandlerResult<()> {
645 if value.is_empty() {
646 Err(HandlerError::new(format!("{kind} is required")))
647 } else {
648 Ok(())
649 }
650}
651
652#[cfg(test)]
653mod tests {
654 use super::{Context, InvocationMethod};
655 use crate::{
656 AttributeMap, ChannelMap, Flow, PersistenceSchema, Registry, registry::physical_name,
657 value_mapper,
658 };
659 use dex_protocol::dex::{ChannelInfo, Context as ProtoContext, Kv};
660 use std::collections::HashMap;
661
662 struct MapFlow {
663 attributes: AttributeMap<String>,
664 channels: ChannelMap<String>,
665 }
666
667 impl Flow for MapFlow {
668 type StartInput = ();
669
670 fn persistence(&self) -> PersistenceSchema {
671 PersistenceSchema::new()
672 .attribute_map(&self.attributes)
673 .channel_map(&self.channels)
674 }
675 }
676
677 #[test]
678 fn map_introspection_tracks_buffered_changes() {
679 let flow = MapFlow {
680 attributes: AttributeMap::new("items"),
681 channels: ChannelMap::new("messages"),
682 };
683 let registry = Registry::new().register(flow).expect("register map Flow");
684 let registered = registry.flow("MapFlow").expect("lookup map Flow").clone();
685 let attributes = AttributeMap::<String>::new("items");
686 let channels = ChannelMap::<String>::new("messages");
687 let special = "special / key";
688 let mut context = Context::new(
689 InvocationMethod::Rpc,
690 registered,
691 ProtoContext::default(),
692 vec![
693 Kv {
694 key: physical_name("items", special),
695 value: Some(value_mapper::encode_handler(&"initial".to_string()).unwrap()),
696 },
697 Kv {
698 key: physical_name("items", "z"),
699 value: Some(value_mapper::encode_handler(&"remove".to_string()).unwrap()),
700 },
701 ],
702 Vec::new(),
703 None,
704 HashMap::from([
705 (physical_name("messages", special), ChannelInfo { size: 1 }),
706 (physical_name("messages", "empty"), ChannelInfo { size: 0 }),
707 ]),
708 )
709 .expect("create RPC Context");
710
711 assert_eq!(
712 vec![special.to_string(), "z".to_string()],
713 attributes.all_instance_keys(&context).unwrap()
714 );
715 attributes
716 .set(&mut context, "a", "added".to_string())
717 .unwrap();
718 attributes.delete(&mut context, "z").unwrap();
719 assert_eq!(
720 vec!["a".to_string(), special.to_string()],
721 attributes.all_instance_keys(&context).unwrap()
722 );
723 assert_eq!(2, attributes.map_size(&context).unwrap());
724
725 assert_eq!(
726 vec![special.to_string()],
727 channels.all_instance_keys(&context).unwrap()
728 );
729 channels
730 .publish(&mut context, "a", "published".to_string())
731 .unwrap();
732 assert_eq!(
733 vec!["a".to_string(), special.to_string()],
734 channels.all_instance_keys(&context).unwrap()
735 );
736 assert_eq!(2, channels.map_size(&context).unwrap());
737 }
738}