use super::*;
use mj_core::native_agent::{NativeAgentHistoryPage, NativeAgentSummary, NativeAgentView};
use mj_core::snapshot_map::SnapshotMap;
use std::collections::BTreeSet;
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct RuntimeFeedHealth {
pub refresh_error: Option<String>,
pub native_error: Option<String>,
}
impl RuntimeFeedHealth {
pub fn error(&self) -> Option<String> {
self.refresh_error
.clone()
.or_else(|| self.native_error.clone())
}
}
#[derive(Default)]
pub(super) struct NativeAgentLoader {
desired: SnapshotMap<String, NativeAgentSummary>,
views: SnapshotMap<String, NativeAgentView>,
failures: BTreeMap<String, (NativeAgentSummary, String)>,
retry: BTreeSet<String>,
active: BTreeSet<String>,
pending: BTreeSet<String>,
tasks: tokio::task::JoinSet<(String, NativeAgentSummary, Result<NativeAgentView>)>,
}
impl NativeAgentLoader {
#[cfg(test)]
pub fn update(&mut self, summaries: Vec<NativeAgentSummary>) {
self.update_snapshot(
summaries
.into_iter()
.map(|summary| (summary.agent.view_id(), summary))
.collect(),
);
}
pub fn update_snapshot(&mut self, summaries: SnapshotMap<String, NativeAgentSummary>) {
for (id, summary) in self.desired.changes(&summaries) {
if summary.is_some_and(|summary| {
self.views
.get(id)
.is_none_or(|view| !summary.is_satisfied_by(view))
}) {
self.pending.insert(id.clone());
} else {
self.pending.remove(id);
}
if summary.is_none_or(|summary| {
self.views
.get(id)
.is_some_and(|view| summary.generation_ordinal != view.generation_ordinal)
}) {
self.views.remove(id);
}
self.failures.remove(id);
}
self.desired = summaries;
self.failures
.retain(|id, (summary, _)| self.desired.get(id) == Some(summary));
self.retry = self.failures.keys().cloned().collect();
}
#[cfg(test)]
pub fn views(&self) -> Vec<NativeAgentView> {
self.views.values().cloned().collect()
}
pub fn views_snapshot(&self) -> SnapshotMap<String, NativeAgentView> {
self.views.clone()
}
pub fn error(&self) -> Option<String> {
self.failures
.values()
.next()
.map(|(_, error)| error.clone())
}
fn pending(&self) -> Option<(String, NativeAgentSummary)> {
self.pending
.iter()
.filter_map(|id| self.desired.get(id).map(|summary| (id, summary)))
.find(|(id, summary)| {
!self.active.contains(*id)
&& (!self.failures.contains_key(*id) || self.retry.contains(*id))
&& self
.views
.get(*id)
.is_none_or(|view| !summary.is_satisfied_by(view))
})
.map(|(id, summary)| (id.clone(), summary.clone()))
}
pub fn has_work(&self) -> bool {
!self.tasks.is_empty() || self.pending().is_some()
}
pub async fn next(&mut self) {
self.next_with(load_native_projection).await;
}
async fn next_with<L, F>(&mut self, load: L)
where
L: Fn(NativeAgentSummary) -> F + Clone + Send + 'static,
F: Future<Output = Result<NativeAgentView>> + Send + 'static,
{
while self.tasks.len() < 4 {
let Some((id, summary)) = self.pending() else {
break;
};
self.active.insert(id.clone());
self.retry.remove(&id);
let load = load.clone();
self.tasks.spawn(async move {
let result = load(summary.clone()).await;
(id, summary, result)
});
}
let Some(completed) = self.tasks.join_next().await else {
return;
};
let (id, requested, result) = match completed {
Ok(completed) => completed,
Err(error) => {
tracing::warn!(%error, "native projection task failed");
self.tasks = tokio::task::JoinSet::new();
for id in std::mem::take(&mut self.active) {
if let Some(summary) = self.desired.get(&id) {
self.failures.insert(
id,
(
summary.clone(),
format!("Native agent reader failed: {error}"),
),
);
}
}
return;
}
};
self.active.remove(&id);
let Some(desired) = self.desired.get(&id) else {
return;
};
if desired.generation_ordinal != requested.generation_ordinal {
return;
}
match result {
Ok(view) if desired.is_satisfied_by(&view) => {
self.failures.remove(&id);
self.pending.remove(&id);
self.views.insert(id, view);
}
Ok(_) => {} Err(error) if desired == &requested => {
let message = format!(
"Could not refresh native agent {}: {error:#}",
desired.agent.name
);
tracing::warn!(%id, %message);
self.failures.insert(id, (requested, message));
}
Err(error) => {
tracing::debug!(%id, %error, "newer native publication superseded a failed read");
}
}
}
}
async fn load_native_projection(summary: NativeAgentSummary) -> Result<NativeAgentView> {
converge_native_projection(&summary, || {
let owner = summary.agent.owner_session_id.clone();
let child = summary.agent.session_id.clone();
native_read(move || crate::database::load_native_agent_view(&owner, &child, 200))
})
.await
}
async fn converge_native_projection<F, Fut>(
summary: &NativeAgentSummary,
mut read: F,
) -> Result<NativeAgentView>
where
F: FnMut() -> Fut,
Fut: Future<Output = Result<Option<NativeAgentView>>>,
{
for attempt in 0..=runtime_feed::PROJECTION_CONVERGENCE_RETRIES {
let view = read().await?;
if let Some(view) = view
&& summary.is_satisfied_by(&view)
{
return Ok(view);
}
if attempt < runtime_feed::PROJECTION_CONVERGENCE_RETRIES {
tokio::time::sleep(runtime_feed::PROJECTION_CONVERGENCE_RETRY_DELAY).await;
}
}
bail!(
"stored projection did not converge to generation {} at event {}",
summary.generation_ordinal,
summary.projection_ordinal
)
}
async fn native_read<T: Send + 'static>(
read: impl FnOnce() -> Result<T> + Send + 'static,
) -> Result<T> {
let permit = Arc::clone(&runtime_feed::PROJECTION_READERS)
.acquire_owned()
.await
.context("native projection readers stopped")?;
tokio::task::spawn_blocking(move || {
let _permit = permit;
let result = read();
if let Err(error) = &result {
tracing::warn!(%error, "native projection read failed");
}
result
})
.await
.context("native projection reader panicked")?
}
pub async fn load_native_agent_history(
owner: String,
child: String,
before: Option<(u64, String)>,
) -> Result<NativeAgentHistoryPage> {
native_read(move || crate::database::native_agent_history(&owner, &child, before)).await
}
#[cfg(test)]
mod tests {
use super::*;
use mj_core::native_agent::{NativeAgent, NativeAgentCapabilities, NativeAgentState};
use std::sync::atomic::AtomicUsize;
fn view(child: &str, generation: u64, ordinal: u64) -> NativeAgentView {
let agent = NativeAgent {
availability: Default::default(),
availability_reason: None,
stable_id: None,
owner_session_id: "owner".into(),
session_id: child.into(),
parent_session_id: None,
name: child.into(),
task: "Inspect".into(),
capabilities: NativeAgentCapabilities::default(),
state: NativeAgentState::Running,
};
let mut projection = MaterializedSession::empty(agent.view_id());
projection.applied_event_ordinal = ordinal;
projection.applied_event_digest = format!("digest-{ordinal}");
NativeAgentView {
generation_ordinal: generation,
agent,
projection,
}
}
#[tokio::test]
async fn unchanged_views_are_cached_and_changed_or_replayed_views_reload() {
let mut loader = NativeAgentLoader::default();
let original = view("child", 1, 2);
loader.update(vec![NativeAgentSummary::of(&original)]);
let answer = original.clone();
loader
.next_with(move |_| {
let answer = answer.clone();
async { Ok(answer) }
})
.await;
assert_eq!(loader.views(), vec![original.clone()]);
loader.update(vec![NativeAgentSummary::of(&original)]);
assert!(!loader.has_work());
let changed = view("child", 1, 3);
loader.update(vec![NativeAgentSummary::of(&changed)]);
let answer = changed.clone();
loader
.next_with(move |_| {
let answer = answer.clone();
async { Ok(answer) }
})
.await;
assert_eq!(loader.views(), vec![changed]);
let replayed = view("child", 4, 5);
loader.update(vec![NativeAgentSummary::of(&replayed)]);
assert!(loader.views().is_empty());
let answer = replayed.clone();
loader
.next_with(move |_| {
let answer = answer.clone();
async { Ok(answer) }
})
.await;
assert_eq!(loader.views(), vec![replayed]);
loader.update(Vec::new());
assert!(loader.views().is_empty());
assert!(!loader.has_work());
}
#[tokio::test]
async fn failed_reads_keep_last_good_view_and_recover_without_hiding_the_error_early() {
let mut loader = NativeAgentLoader::default();
let original = view("child", 1, 2);
loader.update(vec![NativeAgentSummary::of(&original)]);
let answer = original.clone();
loader
.next_with(move |_| {
let answer = answer.clone();
async { Ok(answer) }
})
.await;
let changed = view("child", 1, 3);
loader.update(vec![NativeAgentSummary::of(&changed)]);
loader.next_with(|_| async { bail!("read failed") }).await;
assert_eq!(loader.views(), vec![original]);
assert!(loader.error().unwrap().contains("read failed"));
assert!(!loader.has_work());
loader.update(vec![NativeAgentSummary::of(&changed)]);
assert!(loader.error().is_some());
let answer = changed.clone();
loader
.next_with(move |_| {
let answer = answer.clone();
async { Ok(answer) }
})
.await;
assert!(loader.error().is_none());
assert_eq!(loader.views(), vec![changed]);
}
#[tokio::test]
async fn slow_reads_are_bounded_and_cannot_resurrect_removed_children() {
let mut loader = NativeAgentLoader::default();
loader.update(
(0..6)
.map(|id| NativeAgentSummary::of(&view(&id.to_string(), 1, 2)))
.collect(),
);
let started = Arc::new(AtomicUsize::new(0));
let release = Arc::new(tokio::sync::Semaphore::new(0));
let load = {
let started = started.clone();
let release = release.clone();
move |summary: NativeAgentSummary| {
let started = started.clone();
let release = release.clone();
async move {
started.fetch_add(1, Ordering::SeqCst);
release.acquire().await.unwrap().forget();
Ok(view(&summary.agent.session_id, 1, 2))
}
}
};
assert!(
tokio::time::timeout(Duration::from_millis(50), loader.next_with(load.clone()))
.await
.is_err()
);
assert_eq!(started.load(Ordering::SeqCst), 4);
loader.update(Vec::new());
release.add_permits(4);
while loader.has_work() {
loader.next_with(load.clone()).await;
}
assert!(loader.views().is_empty());
assert_eq!(started.load(Ordering::SeqCst), 4);
}
#[tokio::test]
async fn a_slow_child_does_not_delay_a_ready_sibling() {
let mut loader = NativeAgentLoader::default();
loader.update(vec![
NativeAgentSummary::of(&view("slow", 1, 2)),
NativeAgentSummary::of(&view("ready", 1, 2)),
]);
tokio::time::timeout(
Duration::from_secs(1),
loader.next_with(|summary| async move {
if summary.agent.session_id == "slow" {
std::future::pending::<()>().await;
}
Ok(view(&summary.agent.session_id, 1, 2))
}),
)
.await
.unwrap();
assert_eq!(loader.views()[0].agent.session_id, "ready");
}
#[tokio::test]
async fn read_panics_are_reported_and_retryable() {
let mut loader = NativeAgentLoader::default();
let original = view("child", 1, 2);
loader.update(vec![NativeAgentSummary::of(&original)]);
loader
.next_with(|_| async {
panic!("reader panic");
#[allow(unreachable_code)]
Ok(view("child", 1, 2))
})
.await;
assert!(loader.error().unwrap().contains("reader panic"));
assert!(!loader.has_work());
loader.update(vec![NativeAgentSummary::of(&original)]);
loader
.next_with(|_| async { Ok(view("child", 1, 2)) })
.await;
assert!(loader.error().is_none());
}
#[tokio::test(start_paused = true)]
async fn lagging_projection_retries_but_persistent_mismatch_is_reported() {
let expected = NativeAgentSummary::of(&view("child", 1, 3));
let mut reads = 0;
let result = converge_native_projection(&expected, || {
reads += 1;
let ordinal = if reads == 1 { 2 } else { 3 };
async move { Ok(Some(view("child", 1, ordinal))) }
})
.await
.unwrap();
assert_eq!(result.projection.applied_event_ordinal, 3);
assert_eq!(reads, 2);
let error =
converge_native_projection(&expected, || async { Ok(Some(view("child", 1, 2))) })
.await
.unwrap_err();
assert!(error.to_string().contains("did not converge"), "{error}");
let newer =
converge_native_projection(&expected, || async { Ok(Some(view("child", 1, 4))) })
.await
.unwrap();
assert_eq!(newer.projection.applied_event_ordinal, 4);
}
#[tokio::test]
async fn replay_during_an_inflight_read_discards_the_old_generation() {
let mut loader = NativeAgentLoader::default();
loader.update(vec![NativeAgentSummary::of(&view("child", 1, 2))]);
let release = Arc::new(tokio::sync::Semaphore::new(0));
let load = {
let release = release.clone();
move |_: NativeAgentSummary| {
let release = release.clone();
async move {
release.acquire().await.unwrap().forget();
Ok(view("child", 1, 2))
}
}
};
assert!(
tokio::time::timeout(Duration::from_millis(20), loader.next_with(load.clone()))
.await
.is_err()
);
loader.update(vec![NativeAgentSummary::of(&view("child", 4, 5))]);
release.add_permits(1);
loader.next_with(load).await;
assert!(loader.views().is_empty());
loader
.next_with(|_| async { Ok(view("child", 4, 5)) })
.await;
assert_eq!(loader.views()[0].generation_ordinal, 4);
}
}