1use crate::bus::{BusManager, BusMessage};
2use crate::cache::{CacheWrite, EntityCache, PatchOrigin};
3use crate::mutation_batch::{MutationBatch, SlotContext};
4use crate::view::{ViewIndex, ViewSpec};
5use crate::websocket::frame::{apply_wire_format, Mode, SourceFrame};
6use arete_interpreter::vm::{VmContext, WholeEntityRequests};
7use arete_interpreter::{CanonicalLog, WholeEntity};
8use bytes::Bytes;
9use lru::LruCache;
10use serde_json::Value;
11use smallvec::SmallVec;
12use std::collections::HashSet;
13use std::future::Future;
14use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
15use std::sync::{Arc, Mutex as StdMutex};
16use tokio::sync::mpsc;
17use tracing::{debug, debug_span, error, instrument, warn};
18
19#[cfg(feature = "otel")]
20use crate::metrics::Metrics;
21
22const WHOLE_ENTITY_REQUEST_CAPACITY: usize = 4_096;
26
27type LastChanges = LruCache<(String, String), Option<String>>;
29
30#[derive(Clone)]
60pub struct EntityResync {
61 requests: WholeEntityRequests,
62 last_changes: Arc<StdMutex<LastChanges>>,
66 linked: Arc<AtomicBool>,
67 warned_unlinked: Arc<AtomicBool>,
68}
69
70tokio::task_local! {
71 static ACTIVE_ENTITY_RESYNC: EntityResync;
72}
73
74impl Default for EntityResync {
75 fn default() -> Self {
76 Self::new()
77 }
78}
79
80impl EntityResync {
81 pub fn new() -> Self {
82 Self {
83 requests: WholeEntityRequests::new(WHOLE_ENTITY_REQUEST_CAPACITY),
84 last_changes: Arc::new(StdMutex::new(LruCache::unbounded())),
85 linked: Arc::new(AtomicBool::new(false)),
86 warned_unlinked: Arc::new(AtomicBool::new(false)),
87 }
88 }
89
90 pub async fn scope<F: Future>(&self, future: F) -> F::Output {
93 ACTIVE_ENTITY_RESYNC.scope(self.clone(), future).await
94 }
95
96 pub fn link(&self, vm: &Arc<StdMutex<VmContext>>) {
98 vm.lock()
99 .unwrap_or_else(|poisoned| poisoned.into_inner())
100 .set_whole_entity_requests(self.requests.clone());
101 self.linked.store(true, Ordering::Release);
102 }
103
104 pub fn is_linked(&self) -> bool {
105 self.linked.load(Ordering::Acquire)
106 }
107
108 fn unmarked_origin(&self) -> PatchOrigin {
111 if self.is_linked() {
112 PatchOrigin::Change
113 } else {
114 PatchOrigin::Unknown
115 }
116 }
117
118 pub fn requests(&self) -> &WholeEntityRequests {
120 &self.requests
121 }
122
123 fn last_changes(&self) -> std::sync::MutexGuard<'_, LastChanges> {
124 self.last_changes
125 .lock()
126 .unwrap_or_else(|poisoned| poisoned.into_inner())
127 }
128
129 fn note_change(&self, export: &str, key: &str, seq: &Option<String>, refused: bool) {
134 if !self.is_linked() {
135 return;
136 }
137 let mut last_changes = self.last_changes();
138 if refused {
139 last_changes.put((export.to_string(), key.to_string()), seq.clone());
140 while last_changes.len() > WHOLE_ENTITY_REQUEST_CAPACITY {
141 last_changes.pop_lru();
142 }
143 } else if !last_changes.is_empty() {
144 if let Some(last) = last_changes.get_mut(&(export.to_string(), key.to_string())) {
145 *last = seq.clone();
146 }
147 }
148 }
149
150 fn take_last_change(&self, export: &str, key: &str) -> Option<Option<String>> {
153 let mut last_changes = self.last_changes();
154 if last_changes.is_empty() {
155 return None;
156 }
157 last_changes.pop(&(export.to_string(), key.to_string()))
158 }
159
160 fn request(&self, export: &str, key: &Value) {
161 if self.is_linked() {
162 self.requests.request(export, key);
163 } else if !self.warned_unlinked.swap(true, Ordering::Relaxed) {
164 warn!(
165 export,
166 "Entity cache refused a patch for an evicted key and no VM is registered to \
167 send the whole entity; the key stays out of snapshots until its source sends \
168 a mutation marked whole or deletes it. Raise max_entities_per_view if the \
169 working set exceeds it."
170 );
171 }
172 }
173}
174
175pub(crate) fn link_active_resync(vm: &Arc<StdMutex<VmContext>>) {
178 let _ = ACTIVE_ENTITY_RESYNC.try_with(|resync| resync.link(vm));
179}
180
181struct FrameVersions {
196 epoch: String,
197 last: AtomicU64,
198}
199
200impl FrameVersions {
201 fn new() -> Self {
202 let epoch = hex::encode(&uuid::Uuid::new_v4().as_bytes()[..4]);
205 Self {
206 epoch,
207 last: AtomicU64::new(0),
208 }
209 }
210
211 fn stamp(&self, data: &mut Value) {
214 let Value::Object(map) = data else {
215 debug!("mutation is not an object; publishing it without a version");
216 return;
217 };
218 let counter = self.last.fetch_add(1, Ordering::Relaxed) + 1;
219 map.insert(
220 "_version".to_string(),
221 Value::String(format!("{}:{counter}", self.epoch)),
222 );
223 }
224}
225
226pub struct Projector {
227 view_index: Arc<ViewIndex>,
228 bus_manager: BusManager,
229 entity_cache: EntityCache,
230 mutations_rx: mpsc::Receiver<MutationBatch>,
231 snapshot_runtime: Option<crate::snapshot::SnapshotRuntime>,
232 journal: Option<Arc<crate::journal::EventJournal>>,
233 resync: EntityResync,
234 versions: FrameVersions,
235 #[cfg(feature = "otel")]
236 metrics: Option<Arc<Metrics>>,
237}
238
239impl Projector {
240 #[cfg(feature = "otel")]
241 pub fn new(
242 view_index: Arc<ViewIndex>,
243 bus_manager: BusManager,
244 entity_cache: EntityCache,
245 mutations_rx: mpsc::Receiver<MutationBatch>,
246 metrics: Option<Arc<Metrics>>,
247 ) -> Self {
248 Self {
249 view_index,
250 bus_manager,
251 entity_cache,
252 mutations_rx,
253 snapshot_runtime: None,
254 journal: None,
255 resync: EntityResync::new(),
256 versions: FrameVersions::new(),
257 metrics,
258 }
259 }
260
261 #[cfg(not(feature = "otel"))]
262 pub fn new(
263 view_index: Arc<ViewIndex>,
264 bus_manager: BusManager,
265 entity_cache: EntityCache,
266 mutations_rx: mpsc::Receiver<MutationBatch>,
267 ) -> Self {
268 Self {
269 view_index,
270 bus_manager,
271 entity_cache,
272 mutations_rx,
273 snapshot_runtime: None,
274 journal: None,
275 resync: EntityResync::new(),
276 versions: FrameVersions::new(),
277 }
278 }
279
280 pub fn with_snapshot_runtime(
282 mut self,
283 snapshot_runtime: crate::snapshot::SnapshotRuntime,
284 ) -> Self {
285 self.snapshot_runtime = Some(snapshot_runtime);
286 self
287 }
288
289 pub fn with_journal(mut self, journal: Arc<crate::journal::EventJournal>) -> Self {
291 self.journal = Some(journal);
292 self
293 }
294
295 pub fn with_entity_resync(mut self, resync: EntityResync) -> Self {
297 self.resync = resync;
298 self
299 }
300
301 pub async fn run(mut self) {
302 debug!("Projector started");
303
304 let mut json_buffer = Vec::with_capacity(4096);
305
306 while let Some(mut batch) = self.mutations_rx.recv().await {
307 let mut log = CanonicalLog::new();
308 log.set("phase", "projector");
309
310 let batch_size = batch.len();
311 let slot_context = batch.slot_context;
312 let batch_span = debug_span!(
313 parent: &batch.span,
314 "projector.batch",
315 batch.mutations = batch_size,
316 batch.position = tracing::field::Empty,
317 frames_published = tracing::field::Empty,
318 );
319 if let (Some(context), false) = (slot_context, batch_span.is_disabled()) {
320 batch_span.record("batch.position", context.to_seq_string());
321 }
322 let _span_guard = batch_span.enter();
323 let mut frames_published = 0u32;
324 let mut errors = 0u32;
325
326 if let Some(ctx) = batch.event_context.as_ref() {
327 log.set("program", &ctx.program)
328 .set("event_kind", &ctx.event_kind)
329 .set("event_type", &ctx.event_type)
330 .set("account", &ctx.account)
331 .set("accounts_count", ctx.accounts_count);
332 }
333
334 let arriving_whole: HashSet<(String, String)> = batch
337 .mutations
338 .iter()
339 .filter(|mutation| mutation.is_whole_entity())
340 .map(|mutation| (mutation.export.clone(), Self::extract_key(&mutation.key)))
341 .collect();
342
343 for mutation in std::mem::take(&mut batch.mutations).into_iter() {
344 #[cfg(feature = "otel")]
345 let export = mutation.export.clone();
346
347 match self
348 .process_mutation(mutation, slot_context, &arriving_whole, &mut json_buffer)
349 .await
350 {
351 Ok(count) => frames_published += count,
352 Err(e) => {
353 error!("Failed to process mutation: {}", e);
354 errors += 1;
355 }
356 }
357
358 #[cfg(feature = "otel")]
359 if let Some(ref metrics) = self.metrics {
360 metrics.record_mutation_processed(&export);
361 }
362 }
363
364 if batch_size > 0 {
369 if let Some(snapshot_runtime) = &self.snapshot_runtime {
370 snapshot_runtime.record_applied_batch(slot_context.map(|ctx| ctx.slot));
371 }
372 }
373 drop(batch.snapshot_guard.take());
374
375 if let Some(ack) = batch.flush_ack.take() {
378 let _ = ack.send(());
379 }
380
381 log.set("batch_size", batch_size)
382 .set("frames_published", frames_published)
383 .set("errors", errors);
384 batch_span.record("frames_published", frames_published);
385
386 #[cfg(feature = "otel")]
387 if let Some(ref metrics) = self.metrics {
388 metrics.record_projector_latency(log.duration_ms());
389 }
390
391 log.emit();
392 }
393
394 debug!("Projector stopped");
395 }
396
397 #[instrument(
398 name = "projector.mutation",
399 level = "debug",
400 skip(self, mutation, slot_context, arriving_whole, json_buffer),
401 fields(export = %mutation.export)
402 )]
403 async fn process_mutation(
404 &self,
405 mut mutation: arete_interpreter::Mutation,
406 slot_context: Option<SlotContext>,
407 arriving_whole: &HashSet<(String, String)>,
408 json_buffer: &mut Vec<u8>,
409 ) -> anyhow::Result<u32> {
410 let specs = self.view_index.by_export(&mutation.export);
411
412 if specs.is_empty() {
413 return Ok(0);
414 }
415
416 let mark = mutation.take_whole_entity_mark();
417 let whole = mark == Some(WholeEntity::Resent);
418 let origin = match mark {
419 Some(WholeEntity::Created) => PatchOrigin::Creation,
420 _ => self.resync.unmarked_origin(),
421 };
422 let key = Self::extract_key(&mutation.key);
423 let arete_interpreter::Mutation {
424 export,
425 key: source_key,
426 mut patch,
427 append,
428 } = mutation;
429
430 let batch_seq = slot_context.map(|ctx| ctx.to_seq_string());
433 let seq = if whole {
434 self.resync
435 .take_last_change(&export, &key)
436 .unwrap_or(batch_seq)
437 } else {
438 batch_seq
439 };
440 if let (Some(seq), Value::Object(map)) = (&seq, &mut patch) {
441 map.insert("_seq".to_string(), Value::String(seq.clone()));
442 }
443
444 let matching_specs: SmallVec<[&ViewSpec; 4]> = specs
445 .iter()
446 .filter(|spec| spec.filters.matches(&key))
447 .collect();
448
449 let match_count = matching_specs.len();
450 if match_count == 0 {
451 return Ok(0);
452 }
453
454 let mut frames_published = 0u32;
455 let mut refused = false;
456
457 for (i, spec) in matching_specs.into_iter().enumerate() {
458 let is_last = i == match_count - 1;
459 let patch_data = if is_last {
460 std::mem::take(&mut patch)
461 } else {
462 patch.clone()
463 };
464
465 let mut projected = spec.projection.apply(patch_data);
469 self.versions.stamp(&mut projected);
470 let mut wire_data = projected.clone();
471 apply_wire_format(&mut wire_data, &spec.wire_format);
472
473 let seq = seq.clone();
474
475 if whole {
476 frames_published += self
477 .apply_whole_entity(spec, &key, projected, wire_data, seq, json_buffer)
478 .await?;
479 continue;
480 }
481
482 let journal = self
487 .journal
488 .as_ref()
489 .filter(|journal| journal.is_enabled() && spec.mode == Mode::Append);
490
491 let mut frame = SourceFrame {
492 mode: spec.mode,
493 export: spec.id.clone(),
494 op: "patch",
495 key: key.clone(),
496 data: wire_data,
497 append: append.clone(),
498 seq,
499 offset: None,
500 };
501
502 let retained = match journal {
503 Some(journal) => {
504 journal
505 .append_with(&spec.id, &key, |offset| {
506 frame.offset = Some(offset);
507 json_buffer.clear();
508 serde_json::to_writer(&mut *json_buffer, &frame)?;
509 Ok::<_, anyhow::Error>(Arc::new(Bytes::copy_from_slice(json_buffer)))
510 })
511 .await?
512 }
513 None => None,
514 };
515 let payload = match retained {
516 Some((_offset, payload)) => payload,
517 None => {
520 frame.offset = None;
521 json_buffer.clear();
522 serde_json::to_writer(&mut *json_buffer, &frame)?;
523 Arc::new(Bytes::copy_from_slice(json_buffer))
524 }
525 };
526
527 let write = self
528 .entity_cache
529 .upsert_with_append(&spec.id, &key, projected, &frame.append, origin)
530 .await;
531
532 match write {
533 CacheWrite::Refused { patch } => {
534 refused = true;
535 if spec.mode == Mode::List {
536 self.merge_into_held_derived_entities(&spec.id, &key, patch, &frame.append)
537 .await;
538 }
539 }
540 CacheWrite::Merged | CacheWrite::Created => {
541 if spec.mode == Mode::List {
542 self.update_derived_view_caches(&spec.id, &key).await;
543 }
544 }
545 }
546
547 let message = Arc::new(BusMessage {
548 key: key.clone(),
549 entity: spec.id.clone(),
550 payload,
551 });
552
553 self.publish_frame(spec, message).await;
554 frames_published += 1;
555
556 #[cfg(feature = "otel")]
557 if let Some(ref metrics) = self.metrics {
558 let mode_str = match spec.mode {
559 Mode::List => "list",
560 Mode::State => "state",
561 Mode::Append => "append",
562 };
563 metrics.record_frame_published(mode_str, &spec.export);
564 }
565 }
566
567 if !whole {
568 self.resync.note_change(&export, &key, &seq, refused);
569 }
570 if refused && !arriving_whole.contains(&(export.clone(), key.clone())) {
571 self.resync.request(&export, &source_key);
572 }
573
574 Ok(frames_published)
575 }
576
577 async fn apply_whole_entity(
586 &self,
587 spec: &ViewSpec,
588 key: &str,
589 projected: Value,
590 wire_data: Value,
591 seq: Option<String>,
592 json_buffer: &mut Vec<u8>,
593 ) -> anyhow::Result<u32> {
594 self.entity_cache
595 .store_whole(&spec.id, key, projected)
596 .await;
597 match spec.mode {
598 Mode::Append => return Ok(0),
599 Mode::List => self.update_derived_view_caches(&spec.id, key).await,
600 Mode::State => {}
601 }
602 let frame = SourceFrame {
603 mode: spec.mode,
604 export: spec.id.clone(),
605 op: "upsert",
606 key: key.to_string(),
607 data: wire_data,
608 append: Vec::new(),
609 seq,
610 offset: None,
611 };
612 json_buffer.clear();
613 serde_json::to_writer(&mut *json_buffer, &frame)?;
614 let message = Arc::new(BusMessage {
615 key: key.to_string(),
616 entity: spec.id.clone(),
617 payload: Arc::new(Bytes::copy_from_slice(json_buffer)),
618 });
619 self.publish_frame(spec, message).await;
620 Ok(1)
621 }
622
623 pub(crate) fn extract_key(key: &serde_json::Value) -> String {
624 key.as_str()
625 .map(|s| s.to_string())
626 .or_else(|| key.as_u64().map(|n| n.to_string()))
627 .or_else(|| key.as_i64().map(|n| n.to_string()))
628 .or_else(|| {
629 key.as_array().and_then(|arr| {
630 let bytes: Vec<u8> = arr
631 .iter()
632 .filter_map(|v| v.as_u64().map(|n| n as u8))
633 .collect();
634 if bytes.len() == arr.len() {
635 Some(hex::encode(&bytes))
636 } else {
637 None
638 }
639 })
640 })
641 .unwrap_or_else(|| key.to_string())
642 }
643
644 async fn update_derived_view_caches(&self, source_view_id: &str, entity_key: &str) {
645 let derived_views = self.view_index.get_derived_views_for_source(source_view_id);
646 if derived_views.is_empty() {
647 return;
648 }
649
650 let entity_data = match self.entity_cache.get(source_view_id, entity_key).await {
651 Some(data) => data,
652 None => return,
653 };
654
655 let max_entries = self.entity_cache.max_entities_per_view();
658 let sorted_caches = self.view_index.sorted_caches();
659 let mut caches = sorted_caches.write().await;
660
661 let mut keeping: SmallVec<[&str; 4]> = SmallVec::new();
666 for spec in &derived_views {
667 let Some(cache) = caches.get_mut(&spec.id) else {
668 continue;
669 };
670 let passes = spec
671 .pipeline
672 .as_ref()
673 .and_then(|pipeline| pipeline.filter.as_ref())
674 .is_none_or(|filter| filter.matches(&entity_data));
675 if !passes {
676 cache.remove(entity_key);
677 continue;
678 }
679 if cache.would_keep(entity_key, &entity_data, max_entries) {
680 keeping.push(spec.id.as_str());
681 }
682 }
683 let mut entity_data = Some(entity_data);
684 for (index, view_id) in keeping.iter().enumerate() {
685 let Some(cache) = caches.get_mut(*view_id) else {
686 continue;
687 };
688 let entity = if index + 1 == keeping.len() {
689 entity_data.take()
690 } else {
691 entity_data.clone()
692 };
693 let Some(entity) = entity else {
694 continue;
695 };
696 cache.upsert_bounded(entity_key.to_string(), entity, max_entries);
697 debug!(
698 "Updated sorted cache for derived view {} with key {}",
699 view_id, entity_key
700 );
701 }
702 }
703
704 async fn merge_into_held_derived_entities(
714 &self,
715 source_view_id: &str,
716 entity_key: &str,
717 patch: Value,
718 append_paths: &[String],
719 ) {
720 let derived_views = self.view_index.get_derived_views_for_source(source_view_id);
721 if derived_views.is_empty() {
722 return;
723 }
724 let max_entries = self.entity_cache.max_entities_per_view();
725 let sorted_caches = self.view_index.sorted_caches();
726 let mut caches = sorted_caches.write().await;
727 for spec in &derived_views {
728 let Some(cache) = caches.get_mut(&spec.id) else {
729 continue;
730 };
731 let Some(mut entity) = cache.get(entity_key).cloned() else {
732 continue;
733 };
734 self.entity_cache
735 .merge_patch(&mut entity, patch.clone(), append_paths);
736 let passes = spec
737 .pipeline
738 .as_ref()
739 .and_then(|pipeline| pipeline.filter.as_ref())
740 .is_none_or(|filter| filter.matches(&entity));
741 if passes {
742 cache.remove(entity_key);
746 cache.upsert_bounded(entity_key.to_string(), entity, max_entries);
747 } else {
748 cache.remove(entity_key);
749 }
750 }
751 }
752
753 #[instrument(
754 name = "projector.publish",
755 level = "debug",
756 skip(self, spec, message),
757 fields(view_id = %spec.id, mode = ?spec.mode)
758 )]
759 async fn publish_frame(&self, spec: &ViewSpec, message: Arc<BusMessage>) {
760 match spec.mode {
761 Mode::State => {
762 self.bus_manager
763 .publish_state(&spec.id, &message.key, message.payload.clone())
764 .await;
765 }
766 Mode::List | Mode::Append => {
767 self.bus_manager.publish_list(&spec.id, message).await;
768 }
769 }
770 }
771}
772
773#[cfg(test)]
774mod tests {
775 use super::*;
776
777 fn linked() -> EntityResync {
778 let resync = EntityResync::new();
779 resync.link(&Arc::new(StdMutex::new(VmContext::new())));
780 resync
781 }
782
783 fn seq(slot: u64) -> Option<String> {
784 Some(SlotContext::new(slot, 0).to_seq_string())
785 }
786
787 #[test]
790 fn a_resend_keeps_the_latest_change_seen_before_it() {
791 let resync = linked();
792 resync.note_change("Round", "1", &seq(5), true);
793 resync.note_change("Round", "1", &seq(6), true);
794 resync.note_change("Round", "2", &seq(7), false);
795 assert_eq!(resync.take_last_change("Round", "2"), None, "never refused");
796 resync.note_change("Round", "1", &seq(8), false);
797 assert_eq!(resync.take_last_change("Round", "1"), Some(seq(8)));
798 assert_eq!(resync.take_last_change("Round", "1"), None);
799 }
800
801 #[test]
802 fn remembered_positions_are_bounded_like_requests() {
803 let resync = linked();
804 for key in 0..=WHOLE_ENTITY_REQUEST_CAPACITY {
805 resync.note_change("Round", &key.to_string(), &seq(key as u64), true);
806 }
807 assert_eq!(
808 resync.take_last_change("Round", "0"),
809 None,
810 "the oldest goes"
811 );
812 let newest = WHOLE_ENTITY_REQUEST_CAPACITY;
813 assert_eq!(
814 resync.take_last_change("Round", &newest.to_string()),
815 Some(seq(newest as u64))
816 );
817 }
818
819 #[test]
822 fn nothing_is_remembered_without_a_vm() {
823 let resync = EntityResync::new();
824 resync.note_change("Round", "1", &seq(5), true);
825 assert_eq!(resync.take_last_change("Round", "1"), None);
826 }
827}