diff --git a/Cargo.lock b/Cargo.lock index 99f715c8e6..2f1e900a9e 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2635,9 +2635,9 @@ dependencies = [ [[package]] name = "desert_core" -version = "0.2.1" +version = "0.2.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "623ecd066d1ec74037f703e5a5833304a5621a0c0201ad72ea04a80054483fc7" +checksum = "ac6b6856f399753c591dcdb7b57fd651eee3e07787672392cf551b8ba1725bdd" dependencies = [ "bigdecimal", "bit-vec 0.6.3", @@ -2647,6 +2647,7 @@ dependencies = [ "chrono-tz", "flate2", "hashbrown 0.16.1", + "im", "lazy_static", "mac_address", "nonempty-collections", @@ -2658,9 +2659,9 @@ dependencies = [ [[package]] name = "desert_macro" -version = "0.1.10" +version = "0.1.11" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d9ed06e04bc6899b2609b595a931e84d79a5f4187225d6f4cd942b2ebf54db81" +checksum = "4573798d7c0123902a9fc19f43c1c0dd2daf587c121119fd5ebcdab46f4e984a" dependencies = [ "bytes", "desert_core", @@ -2672,12 +2673,13 @@ dependencies = [ [[package]] name = "desert_rust" -version = "0.1.10" +version = "0.1.11" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0fd78c31bbbee1a0cd9203362cb1b00306344759d064fc1e8636d5c8c64bfc61" +checksum = "88084a4ce8019c315bcc0142f2fa97c1d7b389a96cfadea6f07a5607126667a0" dependencies = [ "desert_core", "desert_macro", + "im", ] [[package]] diff --git a/Cargo.toml b/Cargo.toml index 3bb9e141db..f761cf209b 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -102,7 +102,7 @@ crossterm = "0.29" darling = "0.20.11" dashmap = "7.0.0-rc2" derive_more = { version = "2.0.1", features = ["display", "into", "from_str"] } -desert_rust = { version = "0.1.10", features = ["bigdecimal", "uuid", "chrono", "nonempty-collections", "serde-json", "bit-vec", "url", "mac_address"] } +desert_rust = { version = "0.1.11", features = ["bigdecimal", "uuid", "chrono", "nonempty-collections", "serde-json", "bit-vec", "url", "mac_address", "im"] } dir-diff = "0.3.3" dirs = "6.0.0" dotenvy = "0.15.7" diff --git a/golem-common/src/base_model/mod.rs b/golem-common/src/base_model/mod.rs index 1018ce2fed..fe77de16f9 100644 --- a/golem-common/src/base_model/mod.rs +++ b/golem-common/src/base_model/mod.rs @@ -531,7 +531,9 @@ pub fn validate_lower_kebab_case_identifier( Debug, Eq, Hash, + Ord, PartialEq, + PartialOrd, golem_schema_derive::IntoSchema, golem_schema_derive::FromSchema, )] diff --git a/golem-common/src/model/mod.rs b/golem-common/src/model/mod.rs index 7a3ddeab08..73f79520d3 100644 --- a/golem-common/src/model/mod.rs +++ b/golem-common/src/model/mod.rs @@ -82,7 +82,7 @@ use desert_rust::{ SerializationContext, }; use http::Uri; -use im::OrdMap; +use im::{OrdMap, Vector}; use rand::prelude::IteratorRandom; use serde::{Deserialize, Serialize}; use std::borrow::Cow; @@ -692,6 +692,321 @@ impl SafeDisplay for RetryConfig { } } +pub const DEFAULT_RECENT_INVOCATION_RESULTS_CAPACITY: usize = 1024; +// Four probes keep the false-positive rate below 2% through 100 times the default exact capacity. +pub const DEFAULT_INVOCATION_RESULT_BLOOM_BITS: usize = 1 << 20; +pub const DEFAULT_INVOCATION_RESULT_BLOOM_HASHES: u8 = 4; + +/// A fixed-size, persistent Bloom filter used to prove that unseen idempotency keys are new +/// without consulting the physical invocation-result index. False positives are allowed; false +/// negatives are not. +#[derive(Clone, Debug, PartialEq, Eq, BinaryCodec)] +pub struct InvocationResultBloom { + words: Vector, + bit_count: usize, + hash_count: u8, +} + +impl InvocationResultBloom { + pub fn new(bit_count: usize, hash_count: u8) -> Self { + assert!( + bit_count > 0, + "invocation result Bloom filter must not be empty" + ); + assert!( + hash_count > 0, + "invocation result Bloom filter must use a hash" + ); + let word_count = bit_count.div_ceil(u64::BITS as usize); + Self { + words: std::iter::repeat_n(0, word_count).collect(), + bit_count, + hash_count, + } + } + + pub fn insert(&mut self, key: &IdempotencyKey) { + for bit in self.bit_indexes(key) { + let word_index = bit / u64::BITS as usize; + let bit_index = bit % u64::BITS as usize; + let word = self.words[word_index] | (1u64 << bit_index); + self.words.set(word_index, word); + } + } + + pub fn might_contain(&self, key: &IdempotencyKey) -> bool { + self.bit_indexes(key).all(|bit| { + let word_index = bit / u64::BITS as usize; + let bit_index = bit % u64::BITS as usize; + self.words[word_index] & (1u64 << bit_index) != 0 + }) + } + + fn bit_indexes(&self, key: &IdempotencyKey) -> impl Iterator + use<> { + let digest = blake3::hash(key.value.as_bytes()); + let bytes = digest.as_bytes(); + let first = u64::from_le_bytes(bytes[0..8].try_into().unwrap()); + let second = u64::from_le_bytes(bytes[8..16].try_into().unwrap()) | 1; + let bit_count = self.bit_count as u64; + let hash_count = self.hash_count; + (0..hash_count).map(move |index| { + first + .wrapping_add((index as u64).wrapping_mul(second)) + .wrapping_rem(bit_count) as usize + }) + } +} + +impl Default for InvocationResultBloom { + fn default() -> Self { + Self::new( + DEFAULT_INVOCATION_RESULT_BLOOM_BITS, + DEFAULT_INVOCATION_RESULT_BLOOM_HASHES, + ) + } +} + +// A deterministic approximation of the persistent tree node, pointers, and scalar value. Dynamic +// idempotency-key bytes are accounted for separately. This only controls when the representation +// switches; neither correctness nor the eventual memory bound depends on exact allocator sizing. +const INVOCATION_RESULT_MAP_ENTRY_OVERHEAD_BYTES: usize = 64; + +#[derive(Clone, Debug, PartialEq, Eq, BinaryCodec)] +enum InvocationResultMembershipState { + Exact { + by_key: OrdMap, + key_bytes: usize, + }, + Indexed { + recent_by_key: OrdMap, + recent_by_index: OrdMap, + bloom: InvocationResultBloom, + }, +} + +/// An exact projection of completed invocation results that switches to a bounded representation +/// once the estimated exact-map footprint exceeds the Bloom filter plus its recent exact entries. +#[derive(Clone, Debug, PartialEq, Eq, BinaryCodec)] +pub struct InvocationResultMembership { + state: InvocationResultMembershipState, + capacity: usize, + bloom_bits: usize, + bloom_hashes: u8, + change_generation: u64, + /// Number of revert entries folded into this status. It identifies the current oplog branch + /// so physical and hydrated result entries from an earlier branch are never reused. + revert_generation: u64, +} + +impl InvocationResultMembership { + pub fn new(capacity: usize, bloom_bits: usize, bloom_hashes: u8) -> Self { + Self { + state: InvocationResultMembershipState::Exact { + by_key: OrdMap::new(), + key_bytes: 0, + }, + capacity, + bloom_bits, + bloom_hashes, + change_generation: 0, + revert_generation: 0, + } + } + + pub fn get(&self, key: &IdempotencyKey) -> Option<&OplogIndex> { + match &self.state { + InvocationResultMembershipState::Exact { by_key, .. } => by_key.get(key), + InvocationResultMembershipState::Indexed { recent_by_key, .. } => { + recent_by_key.get(key) + } + } + } + + pub fn contains_key(&self, key: &IdempotencyKey) -> bool { + self.get(key).is_some() + } + + pub fn insert(&mut self, key: IdempotencyKey, result_index: OplogIndex) { + self.change_generation = self.change_generation.wrapping_add(1); + match &mut self.state { + InvocationResultMembershipState::Exact { by_key, key_bytes } => { + if !by_key.contains_key(&key) { + *key_bytes = key_bytes.saturating_add(key.value.len()); + } + by_key.insert(key, result_index); + self.promote_if_needed(); + } + InvocationResultMembershipState::Indexed { + recent_by_key, + recent_by_index, + bloom, + } => { + bloom.insert(&key); + if let Some(previous_index) = recent_by_key.remove(&key) { + recent_by_index.remove(&previous_index); + } + recent_by_key.insert(key.clone(), result_index); + recent_by_index.insert(result_index, key); + Self::truncate_recent(self.capacity, recent_by_key, recent_by_index); + } + } + } + + pub fn might_contain(&self, key: &IdempotencyKey) -> bool { + match &self.state { + InvocationResultMembershipState::Exact { by_key, .. } => by_key.contains_key(key), + InvocationResultMembershipState::Indexed { bloom, .. } => bloom.might_contain(key), + } + } + + pub fn is_exact_complete(&self) -> bool { + matches!(self.state, InvocationResultMembershipState::Exact { .. }) + } + + pub fn oldest_retained_index(&self) -> Option { + match &self.state { + InvocationResultMembershipState::Exact { by_key, .. } => by_key.values().copied().min(), + InvocationResultMembershipState::Indexed { + recent_by_index, .. + } => recent_by_index.get_min().map(|(index, _)| *index), + } + } + + /// Changes whenever a result is added to or removed from the exact membership. It allows + /// in-process admission checks to ignore unrelated oplog commits while still detecting that a + /// result may have appeared and subsequently been evicted. + pub fn change_generation(&self) -> u64 { + self.change_generation + } + + /// Returns the number of revert entries folded into this status. A change means cached result + /// entries may belong to an obsolete oplog branch even when their indexes still exist. + pub fn revert_generation(&self) -> u64 { + self.revert_generation + } + + pub fn set_revert_generation(&mut self, generation: u64) { + self.revert_generation = generation; + } + + pub fn iter(&self) -> impl Iterator { + match &self.state { + InvocationResultMembershipState::Exact { by_key, .. } => by_key.iter(), + InvocationResultMembershipState::Indexed { recent_by_key, .. } => recent_by_key.iter(), + } + } + + pub fn keys(&self) -> impl Iterator { + match &self.state { + InvocationResultMembershipState::Exact { by_key, .. } => by_key.keys(), + InvocationResultMembershipState::Indexed { recent_by_key, .. } => recent_by_key.keys(), + } + } + + pub fn len(&self) -> usize { + match &self.state { + InvocationResultMembershipState::Exact { by_key, .. } => by_key.len(), + InvocationResultMembershipState::Indexed { recent_by_key, .. } => recent_by_key.len(), + } + } + + pub fn is_empty(&self) -> bool { + self.len() == 0 + } + + pub fn remove(&mut self, key: &IdempotencyKey) -> Option { + let index = match &mut self.state { + InvocationResultMembershipState::Exact { by_key, key_bytes } => { + let index = by_key.remove(key)?; + *key_bytes = key_bytes.saturating_sub(key.value.len()); + index + } + InvocationResultMembershipState::Indexed { + recent_by_key, + recent_by_index, + .. + } => { + let index = recent_by_key.remove(key)?; + recent_by_index.remove(&index); + index + } + }; + self.change_generation = self.change_generation.wrapping_add(1); + Some(index) + } + + fn promote_if_needed(&mut self) { + let InvocationResultMembershipState::Exact { by_key, key_bytes } = &self.state else { + return; + }; + let count = by_key.len(); + if count == 0 { + return; + } + let exact_bytes = key_bytes + .saturating_add(count.saturating_mul(INVOCATION_RESULT_MAP_ENTRY_OVERHEAD_BYTES)); + let retained = count.min(self.capacity); + let average_key_bytes = key_bytes.div_ceil(count); + let indexed_bytes = + self.bloom_bits + .div_ceil(u8::BITS as usize) + .saturating_add(retained.saturating_mul(2usize.saturating_mul( + INVOCATION_RESULT_MAP_ENTRY_OVERHEAD_BYTES.saturating_add(average_key_bytes), + ))); + if exact_bytes <= indexed_bytes { + return; + } + + let mut bloom = InvocationResultBloom::new(self.bloom_bits, self.bloom_hashes); + let mut by_index: Vec<_> = by_key + .iter() + .map(|(key, index)| { + bloom.insert(key); + (*index, key.clone()) + }) + .collect(); + by_index.sort_unstable_by_key(|(index, _)| *index); + let mut recent_by_key = OrdMap::new(); + let mut recent_by_index = OrdMap::new(); + for (index, key) in by_index.into_iter().rev().take(self.capacity) { + recent_by_key.insert(key.clone(), index); + recent_by_index.insert(index, key); + } + self.state = InvocationResultMembershipState::Indexed { + recent_by_key, + recent_by_index, + bloom, + }; + } + + fn truncate_recent( + capacity: usize, + recent_by_key: &mut OrdMap, + recent_by_index: &mut OrdMap, + ) { + while recent_by_key.len() > capacity { + let Some((oldest_index, oldest_key)) = recent_by_index + .get_min() + .map(|(index, key)| (*index, key.clone())) + else { + break; + }; + recent_by_index.remove(&oldest_index); + recent_by_key.remove(&oldest_key); + } + } +} + +impl Default for InvocationResultMembership { + fn default() -> Self { + Self::new( + DEFAULT_RECENT_INVOCATION_RESULTS_CAPACITY, + DEFAULT_INVOCATION_RESULT_BLOOM_BITS, + DEFAULT_INVOCATION_RESULT_BLOOM_HASHES, + ) + } +} + /// Contains status information about a worker according to a given oplog index. /// /// This status is just cached information, all fields must be computable by the oplog alone. @@ -708,9 +1023,10 @@ pub struct AgentStatusRecord { pub pending_updates: VecDeque, pub failed_updates: Vec, pub successful_updates: Vec, - pub invocation_results: HashMap, + pub invocation_results: InvocationResultMembership, pub received_card_transfers: ReceivedCardTransferIndex, pub current_idempotency_key: Option, + pub cancelled_idempotency_key: Option, pub component_revision: ComponentRevision, pub component_size: u64, pub total_linear_memory_size: u64, @@ -756,9 +1072,10 @@ impl Default for AgentStatusRecord { pending_updates: VecDeque::new(), failed_updates: Vec::new(), successful_updates: Vec::new(), - invocation_results: HashMap::new(), + invocation_results: InvocationResultMembership::default(), received_card_transfers: ReceivedCardTransferIndex::default(), current_idempotency_key: None, + cancelled_idempotency_key: None, component_revision: ComponentRevision::INITIAL, component_size: 0, total_linear_memory_size: 0, diff --git a/golem-common/src/model/tests.rs b/golem-common/src/model/tests.rs index f03d7e9867..68c7d3710a 100644 --- a/golem-common/src/model/tests.rs +++ b/golem-common/src/model/tests.rs @@ -18,7 +18,9 @@ use crate::model::oplog::OplogIndex; use crate::model::worker::TypedAgentConfigEntry; use crate::model::{ AccountEmail, AccountId, AgentFilter, AgentFingerprint, AgentId, AgentMetadata, AgentMode, - AgentStatus, AgentStatusRecord, ComponentId, FilterComparator, IdempotencyKey, + AgentStatus, AgentStatusRecord, ComponentId, DEFAULT_INVOCATION_RESULT_BLOOM_BITS, + DEFAULT_INVOCATION_RESULT_BLOOM_HASHES, DEFAULT_RECENT_INVOCATION_RESULTS_CAPACITY, + FilterComparator, IdempotencyKey, InvocationResultBloom, InvocationResultMembership, PendingInvocationRef, PendingUpdateKind, PendingUpdateRef, ReceivedCardTransferIndex, ReceivedCardTransferState, StringFilterComparator, Timestamp, }; @@ -29,6 +31,130 @@ use std::vec; use test_r::test; use uuid::{Uuid, uuid}; +#[test] +fn invocation_result_membership_bounds_exact_entries_without_false_negatives() { + let mut membership = InvocationResultMembership::new(2, 64, 3); + let first = IdempotencyKey::new("first".to_string()); + let second = IdempotencyKey::new("second".to_string()); + let third = IdempotencyKey::new("third".to_string()); + let fourth = IdempotencyKey::new("fourth".to_string()); + let fifth = IdempotencyKey::new("fifth".to_string()); + + membership.insert(first.clone(), OplogIndex::from_u64(10)); + membership.insert(second.clone(), OplogIndex::from_u64(20)); + assert!(membership.is_exact_complete()); + + membership.insert(third.clone(), OplogIndex::from_u64(30)); + membership.insert(fourth.clone(), OplogIndex::from_u64(40)); + membership.insert(fifth.clone(), OplogIndex::from_u64(50)); + + assert_eq!(membership.len(), 2); + assert!(!membership.is_exact_complete()); + assert_eq!(membership.change_generation(), 5); + assert_eq!( + membership.oldest_retained_index(), + Some(OplogIndex::from_u64(40)) + ); + assert_eq!(membership.get(&first), None); + assert_eq!(membership.get(&second), None); + assert_eq!(membership.get(&third), None); + assert_eq!(membership.get(&fourth), Some(&OplogIndex::from_u64(40))); + assert_eq!(membership.get(&fifth), Some(&OplogIndex::from_u64(50))); + assert!(membership.might_contain(&first)); + assert!(membership.might_contain(&second)); + assert!(membership.might_contain(&third)); + assert!(membership.might_contain(&fourth)); + assert!(membership.might_contain(&fifth)); +} + +#[test] +fn invocation_result_membership_updates_recency_for_repeated_keys() { + let mut membership = InvocationResultMembership::new(2, 64, 3); + let first = IdempotencyKey::new("first".to_string()); + let second = IdempotencyKey::new("second".to_string()); + let third = IdempotencyKey::new("third".to_string()); + let fourth = IdempotencyKey::new("fourth".to_string()); + let fifth = IdempotencyKey::new("fifth".to_string()); + let sixth = IdempotencyKey::new("sixth".to_string()); + + membership.insert(first.clone(), OplogIndex::from_u64(10)); + membership.insert(second.clone(), OplogIndex::from_u64(20)); + membership.insert(third, OplogIndex::from_u64(30)); + membership.insert(fourth, OplogIndex::from_u64(40)); + membership.insert(fifth.clone(), OplogIndex::from_u64(50)); + membership.insert(first.clone(), OplogIndex::from_u64(60)); + membership.insert(sixth.clone(), OplogIndex::from_u64(70)); + + assert_eq!(membership.get(&first), Some(&OplogIndex::from_u64(60))); + assert_eq!(membership.get(&second), None); + assert_eq!(membership.get(&fifth), None); + assert_eq!(membership.get(&sixth), Some(&OplogIndex::from_u64(70))); +} + +#[test] +fn invocation_result_membership_binary_round_trip_preserves_membership() { + use crate::serialization::{deserialize, serialize}; + + let mut membership = InvocationResultMembership::new(2, 64, 3); + let first = IdempotencyKey::new("first".to_string()); + let second = IdempotencyKey::new("second".to_string()); + let third = IdempotencyKey::new("third".to_string()); + let fourth = IdempotencyKey::new("fourth".to_string()); + let fifth = IdempotencyKey::new("fifth".to_string()); + membership.insert(first.clone(), OplogIndex::from_u64(10)); + membership.insert(second, OplogIndex::from_u64(20)); + membership.insert(third, OplogIndex::from_u64(30)); + membership.insert(fourth, OplogIndex::from_u64(40)); + membership.insert(fifth, OplogIndex::from_u64(50)); + membership.set_revert_generation(4); + + let bytes = serialize(&membership).unwrap(); + let recovered: InvocationResultMembership = deserialize(&bytes).unwrap(); + + assert_eq!(recovered, membership); + assert!(recovered.might_contain(&first)); + assert_eq!(recovered.revert_generation(), 4); +} + +#[test] +fn small_invocation_result_membership_serializes_without_a_bloom_filter() { + use crate::serialization::{deserialize, serialize}; + + let mut membership = InvocationResultMembership::default(); + membership.insert( + IdempotencyKey::new("first".to_string()), + OplogIndex::from_u64(10), + ); + + assert!(membership.is_exact_complete()); + let bytes = serialize(&membership).unwrap(); + assert!(bytes.len() < 1024); + let recovered: InvocationResultMembership = deserialize(&bytes).unwrap(); + assert_eq!(recovered, membership); +} + +#[test] +fn default_invocation_result_bloom_keeps_new_invocations_local_at_100x_capacity() { + let history_size = 100 * DEFAULT_RECENT_INVOCATION_RESULTS_CAPACITY; + let mut bloom = InvocationResultBloom::new( + DEFAULT_INVOCATION_RESULT_BLOOM_BITS, + DEFAULT_INVOCATION_RESULT_BLOOM_HASHES, + ); + for index in 0..history_size { + bloom.insert(&IdempotencyKey::new(format!("existing-{index}"))); + } + + let sample_size = 100_000; + let false_positives = (0..sample_size) + .filter(|index| bloom.might_contain(&IdempotencyKey::new(format!("new-{index}")))) + .count(); + + assert!( + false_positives * 100 < sample_size * 2, + "default Bloom filter sent {false_positives}/{sample_size} new invocations to physical lookup" + ); +} + #[test] fn timestamp_conversion() { let ts: Timestamp = Timestamp::now_utc(); diff --git a/golem-debugging-service/src/config.rs b/golem-debugging-service/src/config.rs index 82085703a4..e225aab64b 100644 --- a/golem-debugging-service/src/config.rs +++ b/golem-debugging-service/src/config.rs @@ -85,6 +85,7 @@ impl DebugConfig { active_agents: self.active_agents, agent_status_flush: Default::default(), agent_status_checkpoint: Default::default(), + invocation_results: Default::default(), scheduler: self.scheduler, public_worker_api: self.public_worker_api, memory: self.memory, diff --git a/golem-worker-executor/config/worker-executor.sample.env b/golem-worker-executor/config/worker-executor.sample.env index 36c7a6dc9a..89dd015524 100644 --- a/golem-worker-executor/config/worker-executor.sample.env +++ b/golem-worker-executor/config/worker-executor.sample.env @@ -70,6 +70,11 @@ GOLEM__INDEXED_STORAGE_RETRY__MAX_DELAY="1s" GOLEM__INDEXED_STORAGE_RETRY__MAX_JITTER_FACTOR=0.15 GOLEM__INDEXED_STORAGE_RETRY__MIN_DELAY="100ms" GOLEM__INDEXED_STORAGE_RETRY__MULTIPLIER=3.0 +GOLEM__INVOCATION_RESULTS__BLOOM_BITS=1048576 +GOLEM__INVOCATION_RESULTS__BLOOM_HASHES=4 +GOLEM__INVOCATION_RESULTS__HYDRATED_CACHE_CAPACITY=1024 +GOLEM__INVOCATION_RESULTS__PHYSICAL_INDEX_CATCH_UP_CHUNK_SIZE=1024 +GOLEM__INVOCATION_RESULTS__RECENT_CAPACITY=1024 GOLEM__KEY_VALUE_STORAGE__TYPE="NamespaceRouted" GOLEM__KEY_VALUE_STORAGE__CONFIG__CACHE__TYPE="Redis" GOLEM__KEY_VALUE_STORAGE__CONFIG__CACHE__CONFIG__DATABASE=0 @@ -357,6 +362,11 @@ GOLEM__INDEXED_STORAGE_RETRY__MAX_DELAY="1s" GOLEM__INDEXED_STORAGE_RETRY__MAX_JITTER_FACTOR=0.15 GOLEM__INDEXED_STORAGE_RETRY__MIN_DELAY="100ms" GOLEM__INDEXED_STORAGE_RETRY__MULTIPLIER=3.0 +GOLEM__INVOCATION_RESULTS__BLOOM_BITS=1048576 +GOLEM__INVOCATION_RESULTS__BLOOM_HASHES=4 +GOLEM__INVOCATION_RESULTS__HYDRATED_CACHE_CAPACITY=1024 +GOLEM__INVOCATION_RESULTS__PHYSICAL_INDEX_CATCH_UP_CHUNK_SIZE=1024 +GOLEM__INVOCATION_RESULTS__RECENT_CAPACITY=1024 GOLEM__KEY_VALUE_STORAGE__TYPE="InMemory" GOLEM__KEY_VALUE_STORAGE_RETRY__MAX_ATTEMPTS=3 GOLEM__KEY_VALUE_STORAGE_RETRY__MAX_DELAY="1s" @@ -595,6 +605,11 @@ GOLEM__INDEXED_STORAGE_RETRY__MAX_DELAY="1s" GOLEM__INDEXED_STORAGE_RETRY__MAX_JITTER_FACTOR=0.15 GOLEM__INDEXED_STORAGE_RETRY__MIN_DELAY="100ms" GOLEM__INDEXED_STORAGE_RETRY__MULTIPLIER=3.0 +GOLEM__INVOCATION_RESULTS__BLOOM_BITS=1048576 +GOLEM__INVOCATION_RESULTS__BLOOM_HASHES=4 +GOLEM__INVOCATION_RESULTS__HYDRATED_CACHE_CAPACITY=1024 +GOLEM__INVOCATION_RESULTS__PHYSICAL_INDEX_CATCH_UP_CHUNK_SIZE=1024 +GOLEM__INVOCATION_RESULTS__RECENT_CAPACITY=1024 GOLEM__KEY_VALUE_STORAGE__TYPE="InMemory" GOLEM__KEY_VALUE_STORAGE_RETRY__MAX_ATTEMPTS=3 GOLEM__KEY_VALUE_STORAGE_RETRY__MAX_DELAY="1s" diff --git a/golem-worker-executor/config/worker-executor.toml b/golem-worker-executor/config/worker-executor.toml index c952826739..673369bebb 100644 --- a/golem-worker-executor/config/worker-executor.toml +++ b/golem-worker-executor/config/worker-executor.toml @@ -121,6 +121,13 @@ max_jitter_factor = 0.15 min_delay = "100ms" multiplier = 3.0 +[invocation_results] +bloom_bits = 1048576 +bloom_hashes = 4 +hydrated_cache_capacity = 1024 +physical_index_catch_up_chunk_size = 1024 +recent_capacity = 1024 + [key_value_storage] type = "NamespaceRouted" @@ -542,6 +549,13 @@ without_time = false # min_delay = "100ms" # multiplier = 3.0 # +# [invocation_results] +# bloom_bits = 1048576 +# bloom_hashes = 4 +# hydrated_cache_capacity = 1024 +# physical_index_catch_up_chunk_size = 1024 +# recent_capacity = 1024 +# # [key_value_storage] # type = "InMemory" # @@ -908,6 +922,13 @@ without_time = false # min_delay = "100ms" # multiplier = 3.0 # +# [invocation_results] +# bloom_bits = 1048576 +# bloom_hashes = 4 +# hydrated_cache_capacity = 1024 +# physical_index_catch_up_chunk_size = 1024 +# recent_capacity = 1024 +# # [key_value_storage] # type = "InMemory" # diff --git a/golem-worker-executor/src/bootstrap.rs b/golem-worker-executor/src/bootstrap.rs index bd901c4425..4c668e8f22 100644 --- a/golem-worker-executor/src/bootstrap.rs +++ b/golem-worker-executor/src/bootstrap.rs @@ -45,6 +45,7 @@ pub async fn run( join_set: &mut JoinSet>, ) -> Result { golem_config.durable_stream.validate()?; + golem_config.invocation_results.validate()?; bootstrap_and_run_worker_executor( &ServerBootstrap, golem_config, diff --git a/golem-worker-executor/src/grpc/mod.rs b/golem-worker-executor/src/grpc/mod.rs index 19140abe24..636b56aae9 100644 --- a/golem-worker-executor/src/grpc/mod.rs +++ b/golem-worker-executor/src/grpc/mod.rs @@ -23,7 +23,7 @@ use crate::model::event::InternalWorkerEvent; use crate::model::public_oplog::{ find_component_revision_at, get_public_oplog_chunk, search_public_oplog, }; -use crate::model::{LastError, ReadFileResult}; +use crate::model::{LastError, LookupResult, ReadFileResult}; use crate::services::events::Event; use crate::services::worker_activator::{ DefaultWorkerActivator, LazyWorkerActivator, WorkerActivator, @@ -754,8 +754,8 @@ impl + UsesAllDeps + Send + Sync + owned_agent_id.agent_id(), ))?; - if metadata - .last_known_status + let status = &metadata.last_known_status; + if status .pending_invocations .iter() .any(|invocation| invocation.idempotency_key() == Some(&idempotency_key)) @@ -773,14 +773,44 @@ impl + UsesAllDeps + Send + Sync + .await?; worker.cancel_invocation(idempotency_key).await?; Ok(true) - } else if metadata - .last_known_status - .invocation_results - .contains_key(&idempotency_key) - { + } else if status.invocation_results.contains_key(&idempotency_key) { Ok(false) - } else { + } else if status.current_idempotency_key.as_ref() == Some(&idempotency_key) + || status.invocation_results.is_exact_complete() + || !status.invocation_results.might_contain(&idempotency_key) + { Err(WorkerExecutorError::invalid_request("Invocation not found")) + } else { + let worker = Worker::get_or_create_suspended( + self, + &owned_agent_id, + None, + Vec::new(), + None, + None, + &InvocationContextStack::fresh(), + principal, + ) + .await?; + match worker.lookup_invocation_result(&idempotency_key).await { + LookupResult::Complete(_) | LookupResult::Interrupted => Ok(false), + LookupResult::Pending => { + let status = worker.get_last_known_status().await; + if status + .pending_invocations + .iter() + .any(|invocation| invocation.idempotency_key() == Some(&idempotency_key)) + { + worker.cancel_invocation(idempotency_key).await?; + Ok(true) + } else { + Err(WorkerExecutorError::invalid_request("Invocation not found")) + } + } + LookupResult::New => { + Err(WorkerExecutorError::invalid_request("Invocation not found")) + } + } } } diff --git a/golem-worker-executor/src/metrics.rs b/golem-worker-executor/src/metrics.rs index 208bc2101f..4fc2e951c1 100644 --- a/golem-worker-executor/src/metrics.rs +++ b/golem-worker-executor/src/metrics.rs @@ -447,6 +447,22 @@ pub mod workers { &["reason"] ) .unwrap(); + static ref INVOCATION_RESULT_RESOLUTION_TOTAL: CounterVec = register_counter_vec!( + "invocation_result_resolution_total", + "Invocation-result resolutions by the lookup path and outcome", + &["outcome"] + ) + .unwrap(); + static ref INVOCATION_RESULT_INDEX_CATCH_UP_CHUNKS_TOTAL: Counter = register_counter!( + "invocation_result_index_catch_up_chunks_total", + "Physical invocation-result index catch-up chunks completed" + ) + .unwrap(); + static ref INVOCATION_RESULT_INDEX_CATCH_UP_ENTRIES_TOTAL: Counter = register_counter!( + "invocation_result_index_catch_up_entries_total", + "Oplog entries processed while catching up the physical invocation-result index" + ) + .unwrap(); static ref AGENT_FILESYSTEM_LIFECYCLE_SECONDS: HistogramVec = register_histogram_vec!( "golem_agent_filesystem_lifecycle_seconds", "Time spent creating or deleting an agent runtime filesystem, labelled by operation and outcome", @@ -517,6 +533,17 @@ pub mod workers { .inc(); } + pub fn record_invocation_result_resolution(outcome: &'static str) { + INVOCATION_RESULT_RESOLUTION_TOTAL + .with_label_values(&[outcome]) + .inc(); + } + + pub fn record_invocation_result_index_catch_up(entries: usize) { + INVOCATION_RESULT_INDEX_CATCH_UP_CHUNKS_TOTAL.inc(); + INVOCATION_RESULT_INDEX_CATCH_UP_ENTRIES_TOTAL.inc_by(entries as f64); + } + pub fn record_agent_filesystem_lifecycle( operation: &'static str, success: bool, diff --git a/golem-worker-executor/src/server.rs b/golem-worker-executor/src/server.rs index 1a1a09b475..50e868ac86 100644 --- a/golem-worker-executor/src/server.rs +++ b/golem-worker-executor/src/server.rs @@ -28,6 +28,7 @@ fn main() -> Result<(), anyhow::Error> { match make_config_loader().load_or_dump_config() { Some(mut config) => { config.durable_stream.validate()?; + config.invocation_results.validate()?; rustls::crypto::ring::default_provider() .install_default() .expect("Failed to install crypto provider"); diff --git a/golem-worker-executor/src/services/golem_config.rs b/golem-worker-executor/src/services/golem_config.rs index 5a7e3e48ea..ecb2115d81 100644 --- a/golem-worker-executor/src/services/golem_config.rs +++ b/golem-worker-executor/src/services/golem_config.rs @@ -18,8 +18,11 @@ use figment::providers::{Format, Toml}; use golem_common::config::{ ConfigExample, ConfigLoader, DbPostgresConfig, DbSqliteConfig, HasConfigExamples, RedisConfig, }; -use golem_common::model::RetryConfig; use golem_common::model::base64::Base64; +use golem_common::model::{ + DEFAULT_INVOCATION_RESULT_BLOOM_BITS, DEFAULT_INVOCATION_RESULT_BLOOM_HASHES, + DEFAULT_RECENT_INVOCATION_RESULTS_CAPACITY, InvocationResultMembership, RetryConfig, +}; use golem_common::tracing::TracingConfig; use golem_common::{SafeDisplay, grpc_uri}; use golem_service_base::clients::registry::GrpcRegistryServiceConfig; @@ -78,6 +81,8 @@ pub struct GolemConfig { pub agent_status_flush: AgentStatusFlushConfig, #[serde(default)] pub agent_status_checkpoint: AgentStatusCheckpointConfig, + #[serde(default)] + pub invocation_results: InvocationResultsConfig, pub scheduler: SchedulerConfig, pub public_worker_api: WorkerServiceGrpcConfig, pub memory: MemoryConfig, @@ -224,6 +229,12 @@ impl SafeDisplay for GolemConfig { "{}", self.agent_status_checkpoint.to_safe_string_indented() ); + let _ = writeln!(&mut result, "invocation_results:"); + let _ = writeln!( + &mut result, + "{}", + self.invocation_results.to_safe_string_indented() + ); let _ = writeln!(&mut result, "scheduler:"); let _ = writeln!(&mut result, "{}", self.scheduler.to_safe_string_indented()); let _ = writeln!(&mut result, "public worker api:"); @@ -357,6 +368,7 @@ impl Default for GolemConfig { active_agents: ActiveAgentsConfig::default(), agent_status_flush: AgentStatusFlushConfig::default(), agent_status_checkpoint: AgentStatusCheckpointConfig::default(), + invocation_results: InvocationResultsConfig::default(), public_worker_api: WorkerServiceGrpcConfig::default(), memory: MemoryConfig::default(), filesystem_storage: FilesystemStorageConfig::default(), @@ -776,6 +788,75 @@ impl Default for AgentStatusFlushConfig { } } +/// Controls the bounded in-memory and physical invocation-result lookup structures. +#[derive(Clone, Debug, Serialize, Deserialize)] +pub struct InvocationResultsConfig { + /// Maximum number of recent invocation results retained exactly in agent status. + pub recent_capacity: usize, + /// Number of bits in the persistent invocation-result Bloom filter. + pub bloom_bits: usize, + /// Number of hash probes used by the persistent invocation-result Bloom filter. + pub bloom_hashes: u8, + /// Maximum number of oplog entries processed by one physical-index catch-up step. + pub physical_index_catch_up_chunk_size: u64, + /// Maximum number of invocation results hydrated from the oplog and retained in memory. + pub hydrated_cache_capacity: usize, +} + +impl InvocationResultsConfig { + pub fn validate(&self) -> anyhow::Result<()> { + anyhow::ensure!( + self.bloom_bits > 0, + "invocation result Bloom filter must not be empty" + ); + anyhow::ensure!( + self.bloom_hashes > 0, + "invocation result Bloom filter must use at least one hash" + ); + anyhow::ensure!( + self.physical_index_catch_up_chunk_size > 0, + "invocation result physical-index catch-up chunk size must be at least one" + ); + Ok(()) + } + + pub fn membership(&self) -> InvocationResultMembership { + InvocationResultMembership::new(self.recent_capacity, self.bloom_bits, self.bloom_hashes) + } +} + +impl SafeDisplay for InvocationResultsConfig { + fn to_safe_string(&self) -> String { + let mut result = String::new(); + let _ = writeln!(&mut result, "recent capacity: {}", self.recent_capacity); + let _ = writeln!(&mut result, "bloom bits: {}", self.bloom_bits); + let _ = writeln!(&mut result, "bloom hashes: {}", self.bloom_hashes); + let _ = writeln!( + &mut result, + "physical index catch-up chunk size: {}", + self.physical_index_catch_up_chunk_size + ); + let _ = writeln!( + &mut result, + "hydrated cache capacity: {}", + self.hydrated_cache_capacity + ); + result + } +} + +impl Default for InvocationResultsConfig { + fn default() -> Self { + Self { + recent_capacity: DEFAULT_RECENT_INVOCATION_RESULTS_CAPACITY, + bloom_bits: DEFAULT_INVOCATION_RESULT_BLOOM_BITS, + bloom_hashes: DEFAULT_INVOCATION_RESULT_BLOOM_HASHES, + physical_index_catch_up_chunk_size: 1024, + hydrated_cache_capacity: 1024, + } + } +} + /// Controls the *clean* cached `AgentStatusRecord` checkpoint. /// /// The checkpoint is a separate copy of the status written only at structurally clean boundaries @@ -2457,7 +2538,7 @@ pub fn make_config_loader() -> ConfigLoader { #[cfg(test)] mod tests { - use super::{DurableStreamConfig, Limits}; + use super::{DurableStreamConfig, InvocationResultsConfig, Limits}; use golem_common::SafeDisplay; use serde_json::Value; use test_r::test; @@ -2470,6 +2551,43 @@ mod tests { assert!(config.validate().is_err()); } + #[test] + fn invocation_results_config_rejects_invalid_lookup_structures() { + let mut config = InvocationResultsConfig::default(); + assert!(config.validate().is_ok()); + + config.bloom_bits = 0; + assert!(config.validate().is_err()); + + config.bloom_bits = 1; + config.bloom_hashes = 0; + assert!(config.validate().is_err()); + + config.bloom_hashes = 1; + config.physical_index_catch_up_chunk_size = 0; + assert!(config.validate().is_err()); + } + + #[test] + fn invocation_results_config_constructs_configured_membership() { + let config = InvocationResultsConfig { + recent_capacity: 2, + bloom_bits: 128, + bloom_hashes: 3, + ..InvocationResultsConfig::default() + }; + let mut membership = config.membership(); + for index in 1..=6 { + membership.insert( + golem_common::model::IdempotencyKey::fresh(), + golem_common::model::oplog::OplogIndex::from_u64(index), + ); + } + + assert_eq!(membership.len(), 2); + assert!(!membership.is_exact_complete()); + } + #[test] fn live_stream_event_broadcast_capacity_defaults_to_32() { let limits = Limits::default(); diff --git a/golem-worker-executor/src/services/worker.rs b/golem-worker-executor/src/services/worker.rs index 46f87085d3..06a8e667b6 100644 --- a/golem-worker-executor/src/services/worker.rs +++ b/golem-worker-executor/src/services/worker.rs @@ -22,50 +22,77 @@ use crate::storage::keyvalue::{ KeyValueStorage, KeyValueStorageLabelledApi, KeyValueStorageNamespace, }; use crate::worker::status::calculate_last_known_status_with_checkpoint_reader; +use crate::worker::status::fold_invocation_result_entries; use async_trait::async_trait; use golem_common::model::agent::{AgentMode, ParsedAgentId}; use golem_common::model::oplog::{OplogEntry, OplogIndex}; use golem_common::model::regions::DeletedRegions; use golem_common::model::{ AgentFingerprint, AgentId, AgentMetadata, AgentStatus, AgentStatusRecord, FailedUpdateRecord, - IdempotencyKey, OwnedAgentId, ReceivedCardTransferIndex, ReceivedCardTransferState, ShardId, - SuccessfulUpdateRecord, + IdempotencyKey, InvocationResultMembership, OwnedAgentId, ReceivedCardTransferIndex, + ReceivedCardTransferState, ShardId, SuccessfulUpdateRecord, }; use golem_common::serialization::{deserialize, serialize}; use golem_service_base::error::worker_executor::WorkerExecutorError; use std::collections::{HashMap, HashSet}; -use std::sync::Arc; +use std::sync::{Arc, Mutex as StdMutex, Weak}; +use tokio::sync::Mutex as AsyncMutex; use tracing::debug; -/// Hash field holding the bounded part of the cached `AgentStatusRecord` (everything except the -/// unbounded fields that are stored separately). Always present for a cached status; its absence is -/// treated as a cache miss. +/// Hash field holding the small part of the cached `AgentStatusRecord`. Always present for a cached +/// status; its absence is treated as a cache miss. const STATUS_CORE_FIELD: &str = "core"; +/// Hash field holding the bounded invocation-result membership. Written only when it changes. +const STATUS_MEMBERSHIP_FIELD: &str = "membership"; /// Hash field holding `(skipped_regions, deleted_regions)`. Written only when the regions change. const STATUS_REGIONS_FIELD: &str = "regions"; /// Hash field holding `(failed_updates, successful_updates)`. Written only when they change. const STATUS_UPDATES_FIELD: &str = "updates"; -/// Prefix for per-idempotency-key invocation result fields (`ir:{idempotency_key}` -> `OplogIndex`). -const STATUS_INVOCATION_RESULT_PREFIX: &str = "ir:"; /// Prefix for per-transfer target receipt fields (`tr:{transfer_id}` -> receipt identity). const STATUS_RECEIVED_CARD_TRANSFER_PREFIX: &str = "tr:"; - -fn status_invocation_result_field(key: &IdempotencyKey) -> String { - format!("{STATUS_INVOCATION_RESULT_PREFIX}{}", key.value) -} +const INVOCATION_RESULT_INDEX_METADATA_FIELD: &str = "metadata"; +const INVOCATION_RESULT_INDEX_FIELD_PREFIX: &str = "ir:"; fn status_received_card_transfer_field(transfer_id: &uuid::Uuid) -> String { format!("{STATUS_RECEIVED_CARD_TRANSFER_PREFIX}{transfer_id}") } +fn invocation_result_index_field(key: &IdempotencyKey) -> String { + format!("{INVOCATION_RESULT_INDEX_FIELD_PREFIX}{}", key.value) +} + +#[derive(Debug, Clone, PartialEq, Eq, desert_rust::BinaryCodec)] +pub struct InvocationResultIndexMetadata { + pub covered_through: OplogIndex, + pub revert_generation: u64, + pub current_idempotency_key: Option, + pub cancelled_idempotency_key: Option, +} + +#[derive(Debug, Clone, PartialEq, Eq, desert_rust::BinaryCodec)] +struct PersistedInvocationResult { + // Catch-up is serialized per agent and clears the whole hash before advancing to a newer + // revert generation, so a complete same-generation index cannot retain older-generation + // fields. + revert_generation: u64, + oplog_index: OplogIndex, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum InvocationResultIndexLookup { + Found(OplogIndex), + DefinitiveMiss, + Incomplete, +} + /// The result of computing a status cache write: `(fields_to_set, field_names_to_delete)`. type StatusFieldWrites = (Vec<(String, Vec)>, Vec); -/// The unbounded parts of an [`AgentStatusRecord`] that are stored separately from `core`. They are -/// taken out of the record (`mem::take`) before serializing `core`, so this never clones the large -/// fields. +/// The potentially large parts of an [`AgentStatusRecord`] that are stored separately from `core`. +/// They are taken out of the record (`mem::take`) before serializing `core`, so this never clones +/// the large fields. struct SplitStatusParts { - invocation_results: HashMap, + invocation_results: InvocationResultMembership, received_card_transfers: ReceivedCardTransferIndex, failed_updates: Vec, successful_updates: Vec, @@ -73,12 +100,15 @@ struct SplitStatusParts { deleted_regions: DeletedRegions, } -/// Moves the unbounded fields out of `status`, leaving it as the small fixed-size `core` that is +/// Moves the separately persisted fields out of `status`, leaving the small `core` that is /// serialized into the `core` field. Uses `mem::take`/`mem::replace`, so it does not clone the -/// (potentially large) invocation results / updates / regions. +/// potentially large updates, regions, or transfer index. fn split_status(status: &mut AgentStatusRecord) -> SplitStatusParts { SplitStatusParts { - invocation_results: std::mem::take(&mut status.invocation_results), + invocation_results: std::mem::replace( + &mut status.invocation_results, + InvocationResultMembership::new(0, 1, 1), + ), received_card_transfers: std::mem::take(&mut status.received_card_transfers), failed_updates: std::mem::take(&mut status.failed_updates), successful_updates: std::mem::take(&mut status.successful_updates), @@ -91,10 +121,9 @@ fn split_status(status: &mut AgentStatusRecord) -> SplitStatusParts { /// date. /// /// `core` must be the already-split (emptied) record. When `previous` is `Some`, the result is a -/// delta against it (this is the hot path; invocation results only ever grow there, so `dels` is -/// usually empty). When `previous` is `None` (cold path: create / cache-miss recompute / detach -/// reload), every part is written and `existing_split_fields` is used to delete stale `ir:` and -/// `tr:` fields that are no longer present. +/// delta against it (this is the hot path, where `dels` is usually empty). When `previous` is +/// `None` (cold path: create / cache-miss recompute / detach reload), every part is written and +/// `existing_split_fields` is used to delete stale `tr:` fields. /// /// `core` is always part of `sets` (it carries the `oplog_idx` marker), so the marker and every /// written part advance together in one atomic `set_many` by the caller. @@ -109,6 +138,17 @@ fn compute_status_field_writes( sets.push((STATUS_CORE_FIELD.to_string(), serialize(core)?)); + let membership_changed = match previous { + Some(previous) => previous.invocation_results != parts.invocation_results, + None => true, + }; + if membership_changed { + sets.push(( + STATUS_MEMBERSHIP_FIELD.to_string(), + serialize(&parts.invocation_results)?, + )); + } + let regions_changed = match previous { Some(previous) => { previous.skipped_regions != parts.skipped_regions @@ -139,17 +179,6 @@ fn compute_status_field_writes( match previous { Some(previous) => { - for (key, oplog_idx) in &parts.invocation_results { - if previous.invocation_results.get(key) != Some(oplog_idx) { - sets.push((status_invocation_result_field(key), serialize(oplog_idx)?)); - } - } - for key in previous.invocation_results.keys() { - if !parts.invocation_results.contains_key(key) { - dels.push(status_invocation_result_field(key)); - } - } - for (transfer_id, state) in parts .received_card_transfers .changes_from(&previous.received_card_transfers) @@ -162,15 +191,6 @@ fn compute_status_field_writes( } } None => { - let new_fields: HashSet = parts - .invocation_results - .keys() - .map(status_invocation_result_field) - .collect(); - for (key, oplog_idx) in &parts.invocation_results { - sets.push((status_invocation_result_field(key), serialize(oplog_idx)?)); - } - let new_transfer_fields: HashSet = parts .received_card_transfers .iter() @@ -184,10 +204,8 @@ fn compute_status_field_writes( } for field in existing_split_fields { - if (field.starts_with(STATUS_INVOCATION_RESULT_PREFIX) - && !new_fields.contains(field)) - || (field.starts_with(STATUS_RECEIVED_CARD_TRANSFER_PREFIX) - && !new_transfer_fields.contains(field)) + if field.starts_with(STATUS_RECEIVED_CARD_TRANSFER_PREFIX) + && !new_transfer_fields.contains(field) { dels.push(field.clone()); } @@ -206,14 +224,16 @@ fn reassemble_cached_status( fields: impl IntoIterator, ) -> Option { let mut core: Option = None; + let mut invocation_results: Option = None; let mut regions: Option<(DeletedRegions, DeletedRegions)> = None; let mut updates: Option<(Vec, Vec)> = None; - let mut invocation_results: HashMap = HashMap::new(); let mut received_card_transfers = ReceivedCardTransferIndex::default(); for (name, bytes) in fields { if name == STATUS_CORE_FIELD { core = Some(deserialize::(&bytes).ok()?); + } else if name == STATUS_MEMBERSHIP_FIELD { + invocation_results = Some(deserialize::(&bytes).ok()?); } else if name == STATUS_REGIONS_FIELD { regions = Some(deserialize::<(DeletedRegions, DeletedRegions)>(&bytes).ok()?); } else if name == STATUS_UPDATES_FIELD { @@ -221,9 +241,6 @@ fn reassemble_cached_status( deserialize::<(Vec, Vec)>(&bytes) .ok()?, ); - } else if let Some(key) = name.strip_prefix(STATUS_INVOCATION_RESULT_PREFIX) { - let oplog_idx = deserialize::(&bytes).ok()?; - invocation_results.insert(IdempotencyKey::new(key.to_string()), oplog_idx); } else if let Some(transfer_id) = name.strip_prefix(STATUS_RECEIVED_CARD_TRANSFER_PREFIX) { let transfer_id = uuid::Uuid::parse_str(transfer_id).ok()?; let state = deserialize::(&bytes).ok()?; @@ -233,6 +250,7 @@ fn reassemble_cached_status( } let mut status = core?; + status.invocation_results = invocation_results?; if let Some((skipped_regions, deleted_regions)) = regions { status.skipped_regions = skipped_regions; status.deleted_regions = deleted_regions; @@ -241,7 +259,6 @@ fn reassemble_cached_status( status.failed_updates = failed_updates; status.successful_updates = successful_updates; } - status.invocation_results = invocation_results; status.received_card_transfers = received_card_transfers; Some(status) } @@ -265,6 +282,24 @@ pub trait WorkerService: Send + Sync { async fn remove_cached_status(&self, owned_agent_id: &OwnedAgentId); + async fn catch_up_invocation_result_index( + &self, + _owned_agent_id: &OwnedAgentId, + _agent_mode: AgentMode, + _status: &AgentStatusRecord, + ) -> Result<(), String> { + Ok(()) + } + + async fn lookup_invocation_result_index( + &self, + _owned_agent_id: &OwnedAgentId, + _status: &AgentStatusRecord, + _idempotency_key: &IdempotencyKey, + ) -> Result { + Ok(InvocationResultIndexLookup::Incomplete) + } + /// Returns the persisted [`AgentMode`] for the worker, if it exists. /// /// The mode is decided at worker create time and persisted in the `Create` oplog entry, @@ -277,10 +312,10 @@ pub trait WorkerService: Send + Sync { /// Writes the cached status *blob* for the worker (no `RunningWorkers` index maintenance). /// /// The cached `AgentStatusRecord` is stored split across several fields of a per-agent hash - /// (see [`KeyValueStorageNamespace::AgentStatus`]): a small `core`, the `regions`, the - /// `updates`, and one field per idempotency key. Only the fields that actually changed are - /// written, so the unbounded parts (most notably the invocation results) are not re-sent on - /// every flush. + /// (see [`KeyValueStorageNamespace::AgentStatus`]): a small `core`, the bounded `membership`, + /// the `regions`, the `updates`, and one field per received card transfer. Only fields that + /// changed are written. The complete invocation-result index is maintained in its dedicated + /// namespace. /// /// `previous_status` is the status currently held in the cache (i.e. the last value /// successfully written). When provided, the delta of changed fields is computed against it. @@ -369,6 +404,26 @@ pub struct DefaultWorkerService { oplog_service: Arc, component_service: Arc, config: Arc, + invocation_result_index_locks: Arc>>>>, +} + +struct InvocationResultIndexLock { + registry: Arc>>>>, + owned_agent_id: OwnedAgentId, + inner: Arc>, +} + +impl Drop for InvocationResultIndexLock { + fn drop(&mut self) { + let mut locks = self.registry.lock().unwrap(); + if Arc::strong_count(&self.inner) == 1 + && locks + .get(&self.owned_agent_id) + .is_some_and(|registered| registered.ptr_eq(&Arc::downgrade(&self.inner))) + { + locks.remove(&self.owned_agent_id); + } + } } impl DefaultWorkerService { @@ -385,6 +440,26 @@ impl DefaultWorkerService { oplog_service, component_service, config, + invocation_result_index_locks: Arc::new(StdMutex::new(HashMap::new())), + } + } + + fn invocation_result_index_lock( + &self, + owned_agent_id: &OwnedAgentId, + ) -> InvocationResultIndexLock { + let mut locks = self.invocation_result_index_locks.lock().unwrap(); + let inner = if let Some(lock) = locks.get(owned_agent_id).and_then(Weak::upgrade) { + lock + } else { + let lock = Arc::new(AsyncMutex::new(())); + locks.insert(owned_agent_id.clone(), Arc::downgrade(&lock)); + lock + }; + InvocationResultIndexLock { + registry: self.invocation_result_index_locks.clone(), + owned_agent_id: owned_agent_id.clone(), + inner, } } @@ -412,7 +487,7 @@ impl DefaultWorkerService { } /// Namespace holding the agent's split cached status (one per-agent hash whose fields are - /// `core`, `regions`, `updates`, and `ir:{idempotency_key}`). + /// `core`, `membership`, `regions`, `updates`, and `tr:{transfer_id}`). fn status_namespace(agent_id: &AgentId) -> KeyValueStorageNamespace { KeyValueStorageNamespace::AgentStatus { agent_id: agent_id.clone(), @@ -429,6 +504,12 @@ impl DefaultWorkerService { } } + fn invocation_result_index_namespace(agent_id: &AgentId) -> KeyValueStorageNamespace { + KeyValueStorageNamespace::AgentInvocationResultIndex { + agent_id: agent_id.clone(), + } + } + /// Key holding only the worker's immutable `AgentMode`, stored separately from the status /// so `get_agent_mode` can resolve the oplog namespace without reading the whole /// `AgentStatusRecord`. Populated lazily on a `get_agent_mode` cache miss (durable workers @@ -444,7 +525,7 @@ impl DefaultWorkerService { } /// Reads the cached `AgentStatusRecord` for `owned_agent_id`, if any, reassembling it from the - /// split hash fields (`core`, `regions`, `updates`, `ir:{key}`, `tr:{transfer_id}`). Returns + /// split hash fields (`core`, `membership`, `regions`, `updates`, `tr:{transfer_id}`). Returns /// `None` if the `core` field is missing (cache miss) or any field cannot be deserialized in /// the current format (treated as a cache miss). /// @@ -459,7 +540,7 @@ impl DefaultWorkerService { } /// Reads a split status record (live cache or checkpoint) from `namespace`, reassembling it - /// from the `core` / `regions` / `updates` / `ir:{key}` / `tr:{transfer_id}` fields. Returns + /// from the `core` / `membership` / `regions` / `updates` / `tr:{transfer_id}` fields. Returns /// `None` if `core` is missing (cache miss / torn write) or any field cannot be deserialized in /// the current format. /// @@ -470,8 +551,8 @@ impl DefaultWorkerService { owned_agent_id: &OwnedAgentId, namespace: KeyValueStorageNamespace, ) -> Option { - // Single atomic read of every field of the per-agent status hash (`core`, `regions`, - // `updates`, `ir:{key}`, `tr:{transfer_id}`). This is one round-trip (Redis `HGETALL`, a single + // Single atomic read of every field of the per-agent status hash (`core`, `membership`, + // `regions`, `updates`, `tr:{transfer_id}`). This is one round-trip (Redis `HGETALL`, a single // `SELECT ... WHERE namespace`, or one locked scan in memory) that observes a consistent // snapshot, so it cannot reassemble a torn, mixed-generation record. (A naive `keys` + // `get_many` would be two round-trips, leaving a window where a concurrent writer — the @@ -498,14 +579,14 @@ impl DefaultWorkerService { /// Writes the split status fields for an agent, sending only the parts that changed. /// /// `core` is always written (it carries the `oplog_idx` marker that versions the whole record). - /// `regions`/`updates` are written only when they differ from `previous_status`, and invocation - /// results and received card transfers are written per key (only newly added/changed keys). + /// `membership`/`regions`/`updates` are written only when they differ from `previous_status`, + /// and received card transfers are written per key (only newly added/changed keys). /// /// Atomicity: the marker (in `core`) and every field written in the same call advance together /// in a single atomic `set_many` (one `HMSET` on Redis, one transaction on SQL). This preserves /// the invariant that each persisted field's content matches `core.oplog_idx`, which the oplog - /// fold relies on. When stale fields must be removed (e.g. invocation results dropped by a - /// revert), they are deleted *together with* `core` before the `set_many`. Dropping `core` + /// fold relies on. When stale fields must be removed, they are deleted *together with* `core` + /// before the `set_many`. Dropping `core` /// first is what makes the two-step delete-then-write crash-safe: with `core` absent, any crash /// or read in the gap before the final write is treated as a cache miss and recomputed from the /// oplog, rather than reassembling a torn record whose remaining fields no longer match the @@ -553,7 +634,7 @@ impl DefaultWorkerService { })?; } - // Single atomic write: core + changed parts + new/updated invocation results. + // Single atomic write: core + changed parts. let pairs: Vec<(&str, &[u8])> = sets .iter() .map(|(field, bytes)| (field.as_str(), bytes.as_slice())) @@ -583,7 +664,7 @@ impl DefaultWorkerService { return Ok(status_value); } - // Split the record: take the unbounded fields out so `core` stays small and fixed-size. + // Split the record: take the potentially large fields out so `core` stays small. // `split_status` moves the large fields out of `core` into `parts` (no clone). let mut core = status_value; let parts = split_status(&mut core); @@ -594,11 +675,11 @@ impl DefaultWorkerService { // Reassemble the record (moving the parts back into `core`, no clone) so the caller gets // back a complete baseline for computing the next delta. let mut reassembled = core; + reassembled.invocation_results = parts.invocation_results; reassembled.skipped_regions = parts.skipped_regions; reassembled.deleted_regions = parts.deleted_regions; reassembled.failed_updates = parts.failed_updates; reassembled.successful_updates = parts.successful_updates; - reassembled.invocation_results = parts.invocation_results; reassembled.received_card_transfers = parts.received_card_transfers; Ok(reassembled) } @@ -630,6 +711,31 @@ impl DefaultWorkerService { } } + async fn clear_invocation_result_index( + &self, + owned_agent_id: &OwnedAgentId, + namespace: KeyValueStorageNamespace, + ) -> Result<(), String> { + let fields = self + .key_value_storage + .with("worker", "clear_invocation_result_index") + .keys(namespace.clone()) + .await + .map_err(|err| { + format!("failed to list invocation result index fields for {owned_agent_id}: {err}") + })?; + if !fields.is_empty() { + self.key_value_storage + .with("worker", "clear_invocation_result_index") + .del_many(namespace, fields) + .await + .map_err(|err| { + format!("failed to clear invocation result index for {owned_agent_id}: {err}") + })?; + } + Ok(()) + } + /// Reads the dedicated `agent_mode` key, if present. Returns `None` on a cache miss or if the /// stored value cannot be deserialized in the current format (treated as a miss). async fn read_cached_agent_mode(&self, owned_agent_id: &OwnedAgentId) -> Option { @@ -758,6 +864,7 @@ impl WorkerService for DefaultWorkerService { component_size, total_linear_memory_size: initial_total_linear_memory_size, active_plugins: initial_active_plugins, + invocation_results: self.config.invocation_results.membership(), agent_mode, ..AgentStatusRecord::default() }, @@ -856,6 +963,11 @@ impl WorkerService for DefaultWorkerService { .await; self.remove_split_status(owned_agent_id, Self::checkpoint_namespace(agent_id)) .await; + self.remove_split_status( + owned_agent_id, + Self::invocation_result_index_namespace(agent_id), + ) + .await; // The `agent_mode` key has its own lifecycle and lives in the `Worker` namespace. self.key_value_storage @@ -872,6 +984,181 @@ impl WorkerService for DefaultWorkerService { }); } + async fn catch_up_invocation_result_index( + &self, + owned_agent_id: &OwnedAgentId, + agent_mode: AgentMode, + status: &AgentStatusRecord, + ) -> Result<(), String> { + if agent_mode == AgentMode::Ephemeral { + return Ok(()); + } + + let lock = self.invocation_result_index_lock(owned_agent_id); + let _guard = lock.inner.lock().await; + async { + let namespace = Self::invocation_result_index_namespace(&owned_agent_id.agent_id); + let persisted: Option> = self + .key_value_storage + .with_entity("worker", "read_invocation_result_index", "metadata") + .get_attempt_deserialize( + namespace.clone(), + INVOCATION_RESULT_INDEX_METADATA_FIELD, + ) + .await?; + let status_generation = status.invocation_results.revert_generation(); + let mut metadata = match persisted { + Some(Ok(metadata)) if metadata.revert_generation > status_generation => { + return Ok(()); + } + Some(Ok(metadata)) if metadata.revert_generation == status_generation => { + if metadata.covered_through >= status.oplog_idx { + return Ok(()); + } + metadata + } + Some(Ok(_)) | Some(Err(_)) => { + self.clear_invocation_result_index(owned_agent_id, namespace.clone()) + .await?; + InvocationResultIndexMetadata { + covered_through: OplogIndex::NONE, + revert_generation: status_generation, + current_idempotency_key: None, + cancelled_idempotency_key: None, + } + } + None => InvocationResultIndexMetadata { + covered_through: OplogIndex::NONE, + revert_generation: status_generation, + current_idempotency_key: None, + cancelled_idempotency_key: None, + }, + }; + + while metadata.covered_through < status.oplog_idx { + let remaining = status.oplog_idx.as_u64() - metadata.covered_through.as_u64(); + let count = remaining + .min( + self.config + .invocation_results + .physical_index_catch_up_chunk_size, + ) + .max(1); + let entries = self + .oplog_service + .read_exact( + owned_agent_id, + agent_mode, + metadata.covered_through.next(), + count, + ) + .await; + if entries.is_empty() { + return Err(format!( + "failed to advance invocation result index for {owned_agent_id}: oplog range starting at {} was empty", + metadata.covered_through.next() + )); + } + + let mut mappings = HashMap::new(); + fold_invocation_result_entries( + &mut metadata.current_idempotency_key, + &mut metadata.cancelled_idempotency_key, + &status.deleted_regions, + &entries, + |key, index| { + mappings.insert(key.clone(), index); + }, + ); + metadata.covered_through = *entries.keys().max().unwrap(); + + let mut fields = Vec::with_capacity(mappings.len() + 1); + for (key, oplog_index) in mappings { + fields.push(( + invocation_result_index_field(&key), + serialize(&PersistedInvocationResult { + revert_generation: metadata.revert_generation, + oplog_index, + })?, + )); + } + fields.push(( + INVOCATION_RESULT_INDEX_METADATA_FIELD.to_string(), + serialize(&metadata)?, + )); + let pairs: Vec<(&str, &[u8])> = fields + .iter() + .map(|(field, value)| (field.as_str(), value.as_slice())) + .collect(); + self.key_value_storage + .with_entity( + "worker", + "advance_invocation_result_index", + "invocation_result", + ) + .set_many_raw(namespace.clone(), &pairs) + .await?; + crate::metrics::workers::record_invocation_result_index_catch_up(entries.len()); + } + + Ok(()) + } + .await + } + + async fn lookup_invocation_result_index( + &self, + owned_agent_id: &OwnedAgentId, + status: &AgentStatusRecord, + idempotency_key: &IdempotencyKey, + ) -> Result { + let values = self + .key_value_storage + .with_entity( + "worker", + "lookup_invocation_result_index", + "invocation_result", + ) + .get_many_raw( + Self::invocation_result_index_namespace(&owned_agent_id.agent_id), + vec![ + INVOCATION_RESULT_INDEX_METADATA_FIELD.to_string(), + invocation_result_index_field(idempotency_key), + ], + ) + .await?; + let Some(metadata) = values.first().and_then(Option::as_ref) else { + return Ok(InvocationResultIndexLookup::Incomplete); + }; + let metadata: InvocationResultIndexMetadata = deserialize(metadata)?; + if metadata.revert_generation != status.invocation_results.revert_generation() { + return Ok(InvocationResultIndexLookup::Incomplete); + } + + let complete = metadata.covered_through >= status.oplog_idx + || status + .invocation_results + .oldest_retained_index() + .is_some_and(|oldest| oldest <= metadata.covered_through); + if !complete { + return Ok(InvocationResultIndexLookup::Incomplete); + } + + let Some(value) = values.get(1).and_then(Option::as_ref) else { + return Ok(InvocationResultIndexLookup::DefinitiveMiss); + }; + let value: PersistedInvocationResult = deserialize(value)?; + if value.revert_generation != metadata.revert_generation + || status + .deleted_regions + .is_in_deleted_region(value.oplog_index) + { + return Ok(InvocationResultIndexLookup::DefinitiveMiss); + } + + Ok(InvocationResultIndexLookup::Found(value.oplog_index)) + } + async fn get_agent_mode(&self, owned_agent_id: &OwnedAgentId) -> Option { record_worker_call("get_agent_mode"); @@ -914,6 +1201,13 @@ impl WorkerService for DefaultWorkerService { debug!("Writing cached agent status for {owned_agent_id} to {status_value:?}"); + self.catch_up_invocation_result_index( + owned_agent_id, + status_value.agent_mode, + &status_value, + ) + .await?; + self.write_split_status( owned_agent_id, Self::status_namespace(&owned_agent_id.agent_id), @@ -1034,14 +1328,314 @@ impl HasComponentService for DefaultWorkerService { #[cfg(test)] mod tests { use super::*; + use crate::model::ExecutionStatus; + use crate::services::shard::ShardServiceDefault; + use crate::storage::keyvalue::memory::InMemoryKeyValueStorage; + use async_trait::async_trait; use bytes::Bytes; use golem_common::model::Timestamp; + use golem_common::model::account::AccountId; + use golem_common::model::application::ApplicationId; use golem_common::model::card::{Card, CardId, StoredCard}; - use golem_common::model::component::ComponentRevision; + use golem_common::model::component::{ComponentId, ComponentRevision}; + use golem_common::model::environment::EnvironmentId; + use golem_common::model::invocation_context::TraceId; + use golem_common::model::oplog::{OplogPayload, PayloadId, RawOplogPayload}; use golem_common::model::regions::{DeletedRegions, OplogRegion}; - use golem_common::model::{PendingInvocationRef, PendingUpdateKind, PendingUpdateRef}; - use std::collections::VecDeque; + use golem_common::model::{ + AgentInvocationPayload, AgentInvocationResult, AgentMetadata, PendingInvocationRef, + PendingUpdateKind, PendingUpdateRef, ScanCursor, + }; + use golem_common::read_only_lock; + use golem_service_base::model::component::Component; + use std::collections::{BTreeMap, VecDeque}; + use std::sync::atomic::{AtomicBool, Ordering}; use test_r::test; + use tokio::sync::Notify; + + #[derive(Debug)] + struct IndexTestOplogService { + entries: BTreeMap, + reads: StdMutex>, + pause_next_read: AtomicBool, + read_started: Notify, + resume_read: Notify, + } + + impl IndexTestOplogService { + fn new(entries: BTreeMap) -> Self { + Self { + entries, + reads: StdMutex::new(Vec::new()), + pause_next_read: AtomicBool::new(false), + read_started: Notify::new(), + resume_read: Notify::new(), + } + } + + fn pause_next_read(&self) { + self.pause_next_read.store(true, Ordering::Release); + } + + fn read_starts(&self) -> Vec { + self.reads + .lock() + .unwrap() + .iter() + .map(|(start, _)| *start) + .collect() + } + + fn clear_reads(&self) { + self.reads.lock().unwrap().clear(); + } + } + + #[async_trait] + impl OplogService for IndexTestOplogService { + async fn create( + &self, + _owned_agent_id: &OwnedAgentId, + _agent_mode: AgentMode, + _initial_entry: OplogEntry, + _initial_worker_metadata: AgentMetadata, + _last_known_status: read_only_lock::arc_swap::ReadOnlyView, + _execution_status: read_only_lock::std::ReadOnlyLock, + ) -> Arc { + unreachable!() + } + + async fn create_fresh( + &self, + _owned_agent_id: &OwnedAgentId, + _agent_mode: AgentMode, + _initial_entry: OplogEntry, + _initial_worker_metadata: AgentMetadata, + _last_known_status: read_only_lock::arc_swap::ReadOnlyView, + _execution_status: read_only_lock::std::ReadOnlyLock, + ) -> Arc { + unreachable!() + } + + async fn open( + &self, + _owned_agent_id: &OwnedAgentId, + _agent_mode: AgentMode, + _last_oplog_index: Option, + _initial_worker_metadata: AgentMetadata, + _last_known_status: read_only_lock::arc_swap::ReadOnlyView, + _execution_status: read_only_lock::std::ReadOnlyLock, + ) -> Arc { + unreachable!() + } + + async fn get_last_index( + &self, + _owned_agent_id: &OwnedAgentId, + _agent_mode: AgentMode, + ) -> OplogIndex { + self.entries + .keys() + .next_back() + .copied() + .unwrap_or(OplogIndex::NONE) + } + + async fn delete(&self, _owned_agent_id: &OwnedAgentId, _agent_mode: AgentMode) { + unreachable!() + } + + async fn read_exact( + &self, + _owned_agent_id: &OwnedAgentId, + _agent_mode: AgentMode, + idx: OplogIndex, + n: u64, + ) -> BTreeMap { + self.reads.lock().unwrap().push((idx, n)); + if self.pause_next_read.swap(false, Ordering::AcqRel) { + self.read_started.notify_one(); + self.resume_read.notified().await; + } + let end = idx.as_u64().saturating_add(n.saturating_sub(1)); + self.entries + .range(idx..=OplogIndex::from_u64(end)) + .map(|(index, entry)| (*index, entry.clone())) + .collect() + } + + async fn exists(&self, _owned_agent_id: &OwnedAgentId, _agent_mode: AgentMode) -> bool { + true + } + + async fn scan_for_component( + &self, + _environment_id: &EnvironmentId, + _component_id: &ComponentId, + _modes: Option, + _cursor: ScanCursor, + _count: u64, + ) -> Result<(ScanCursor, Vec), WorkerExecutorError> { + unreachable!() + } + + async fn upload_raw_payload( + &self, + _owned_agent_id: &OwnedAgentId, + _agent_mode: AgentMode, + _data: Vec, + ) -> Result { + unreachable!() + } + + async fn download_raw_payload( + &self, + _owned_agent_id: &OwnedAgentId, + _agent_mode: AgentMode, + _payload_id: PayloadId, + _md5_hash: Vec, + ) -> Result, String> { + unreachable!() + } + } + + struct IndexTestComponentService; + + #[async_trait] + impl ComponentService for IndexTestComponentService { + async fn get( + &self, + _engine: &wasmtime::Engine, + _component_id: ComponentId, + _component_revision: ComponentRevision, + ) -> Result<(wasmtime::component::Component, Component), WorkerExecutorError> { + unreachable!() + } + + async fn get_metadata( + &self, + _component_id: ComponentId, + _forced_revision: Option, + ) -> Result { + unreachable!() + } + + async fn resolve_component( + &self, + _component_reference: String, + _resolving_environment: EnvironmentId, + _resolving_application: ApplicationId, + _resolving_account: AccountId, + ) -> Result, WorkerExecutorError> { + unreachable!() + } + + async fn all_cached_metadata(&self) -> Vec { + Vec::new() + } + + async fn invalidate_all_metadata_for_environment(&self, _environment_id: EnvironmentId) {} + } + + fn invocation_pair( + entries: &mut BTreeMap, + started_at: u64, + key: &IdempotencyKey, + ) { + entries.insert( + OplogIndex::from_u64(started_at), + OplogEntry::AgentInvocationStarted { + timestamp: Timestamp::now_utc(), + idempotency_key: key.clone(), + payload: OplogPayload::Inline(Box::new(AgentInvocationPayload::ManualUpdate { + target_revision: ComponentRevision::INITIAL, + })), + trace_id: TraceId::generate(), + trace_states: Vec::new(), + invocation_context: Vec::new(), + wallet_pin: None, + }, + ); + entries.insert( + OplogIndex::from_u64(started_at + 1), + OplogEntry::AgentInvocationFinished { + timestamp: Timestamp::now_utc(), + result: OplogPayload::Inline(Box::new(AgentInvocationResult::AgentInitialization)), + method_name: None, + consumed_fuel: 0, + component_revision: ComponentRevision::INITIAL, + }, + ); + } + + fn invocation_entries(keys: &[IdempotencyKey]) -> BTreeMap { + let mut entries = BTreeMap::from([(OplogIndex::INITIAL, OplogEntry::no_op())]); + for (offset, key) in keys.iter().enumerate() { + invocation_pair(&mut entries, 2 + offset as u64 * 2, key); + } + entries + } + + fn invocation_status( + oplog_idx: u64, + capacity: usize, + results: &[(&IdempotencyKey, u64)], + revert_generation: u64, + deleted_regions: DeletedRegions, + ) -> AgentStatusRecord { + let mut invocation_results = InvocationResultMembership::new(capacity, 128, 3); + for (key, index) in results { + invocation_results.insert((*key).clone(), OplogIndex::from_u64(*index)); + } + invocation_results.set_revert_generation(revert_generation); + AgentStatusRecord { + oplog_idx: OplogIndex::from_u64(oplog_idx), + invocation_results, + deleted_regions, + ..AgentStatusRecord::default() + } + } + + fn index_test_service( + entries: BTreeMap, + ) -> ( + Arc, + Arc, + OwnedAgentId, + ) { + let oplog = Arc::new(IndexTestOplogService::new(entries)); + let mut config = GolemConfig::default(); + config.invocation_results.physical_index_catch_up_chunk_size = 2; + let service = Arc::new(DefaultWorkerService::new( + Arc::new(InMemoryKeyValueStorage::new()), + Arc::new(ShardServiceDefault::new()), + oplog.clone(), + Arc::new(IndexTestComponentService), + Arc::new(config), + )); + let agent_id = AgentId { + component_id: ComponentId::new(), + agent_id: "invocation-index-test".to_string(), + }; + let owned_agent_id = OwnedAgentId::new(EnvironmentId::new(), &agent_id); + (service, oplog, owned_agent_id) + } + + async fn invocation_index_metadata( + service: &DefaultWorkerService, + owned_agent_id: &OwnedAgentId, + ) -> InvocationResultIndexMetadata { + let value: Option> = service + .key_value_storage + .with_entity("test", "read_invocation_result_index", "metadata") + .get_attempt_deserialize( + DefaultWorkerService::invocation_result_index_namespace(&owned_agent_id.agent_id), + INVOCATION_RESULT_INDEX_METADATA_FIELD, + ) + .await + .unwrap(); + value.unwrap().unwrap() + } fn idempotency_key(value: &str) -> IdempotencyKey { IdempotencyKey::new(value.to_string()) @@ -1157,6 +1751,20 @@ mod tests { assert!(reassemble_cached_status(without_core).is_none()); } + #[test] + fn missing_membership_field_is_a_cache_miss() { + let full = sample_status(); + let mut core = full.clone(); + let parts = split_status(&mut core); + let (sets, _) = compute_status_field_writes(None, &[], &core, &parts).unwrap(); + + let without_membership = sets + .into_iter() + .filter(|(name, _)| name != STATUS_MEMBERSHIP_FIELD) + .map(|(name, bytes)| (name, Bytes::from(bytes))); + assert!(reassemble_cached_status(without_membership).is_none()); + } + #[test] fn hot_delta_only_writes_changed_fields() { let previous = sample_status(); @@ -1172,8 +1780,8 @@ mod tests { } } - // New status: a new invocation result + advanced marker, but identical regions/updates and - // unchanged existing invocation results. + // New status: a new bounded invocation result + advanced marker, but identical split + // regions/updates. let mut new = previous.clone(); new.oplog_idx = OplogIndex::from_u64(50); new.invocation_results @@ -1186,12 +1794,12 @@ mod tests { let written: HashSet<&str> = sets.iter().map(|(f, _)| f.as_str()).collect(); assert!(written.contains(STATUS_CORE_FIELD)); - assert!(written.contains(status_invocation_result_field(&idempotency_key("k3")).as_str())); + assert!(written.contains(STATUS_MEMBERSHIP_FIELD)); // Unchanged parts are NOT re-sent. assert!(!written.contains(STATUS_REGIONS_FIELD)); assert!(!written.contains(STATUS_UPDATES_FIELD)); - assert!(!written.contains(status_invocation_result_field(&idempotency_key("k1")).as_str())); assert!(!written.contains(status_received_card_transfer_field(&transfer_id(1)).as_str())); + assert_eq!(written.len(), 2); assert!(dels.is_empty()); let reassembled = apply_and_reassemble(&mut store, sets, dels).unwrap(); @@ -1224,7 +1832,7 @@ mod tests { } #[test] - fn delta_deletes_removed_invocation_results() { + fn bounded_invocation_result_changes_are_stored_in_membership() { let previous = sample_status(); let mut store = HashMap::new(); @@ -1246,38 +1854,17 @@ mod tests { let (sets, dels) = compute_status_field_writes(Some(&previous), &[], &core, &parts).unwrap(); + assert!(dels.is_empty()); + let written: HashSet<&str> = sets.iter().map(|(field, _)| field.as_str()).collect(); assert_eq!( - dels, - vec![status_invocation_result_field(&idempotency_key("k2"))] + written, + HashSet::from([STATUS_CORE_FIELD, STATUS_MEMBERSHIP_FIELD]) ); let reassembled = apply_and_reassemble(&mut store, sets, dels).unwrap(); assert_eq!(reassembled, new); } - #[test] - fn cold_reconcile_deletes_stale_invocation_results() { - // Store already holds ir:k1 and ir:k2 from a previous state. - let existing_fields = vec![ - STATUS_CORE_FIELD.to_string(), - status_invocation_result_field(&idempotency_key("k1")), - status_invocation_result_field(&idempotency_key("k2")), - ]; - - // New status only has k1. - let mut new = sample_status(); - new.invocation_results.remove(&idempotency_key("k2")); - - let mut core = new.clone(); - let parts = split_status(&mut core); - let (_, dels) = compute_status_field_writes(None, &existing_fields, &core, &parts).unwrap(); - - assert_eq!( - dels, - vec![status_invocation_result_field(&idempotency_key("k2"))] - ); - } - #[test] fn cold_reconcile_deletes_stale_transfer_fields() { let stale_transfer_id = transfer_id(2); @@ -1298,6 +1885,282 @@ mod tests { ); } + #[test] + async fn invocation_result_index_resumes_across_calls_and_chunks() { + let first = idempotency_key("first"); + let second = idempotency_key("second"); + let (service, oplog, owned_agent_id) = + index_test_service(invocation_entries(&[first.clone(), second.clone()])); + let partial = invocation_status(2, 2, &[], 0, DeletedRegions::new()); + let complete = + invocation_status(5, 2, &[(&first, 3), (&second, 5)], 0, DeletedRegions::new()); + + service + .catch_up_invocation_result_index(&owned_agent_id, AgentMode::Durable, &partial) + .await + .unwrap(); + service + .catch_up_invocation_result_index(&owned_agent_id, AgentMode::Durable, &complete) + .await + .unwrap(); + + assert_eq!( + oplog.read_starts(), + vec![ + OplogIndex::INITIAL, + OplogIndex::from_u64(3), + OplogIndex::from_u64(5) + ] + ); + assert_eq!( + service + .lookup_invocation_result_index(&owned_agent_id, &complete, &first) + .await + .unwrap(), + InvocationResultIndexLookup::Found(OplogIndex::from_u64(3)) + ); + assert_eq!( + service + .lookup_invocation_result_index(&owned_agent_id, &complete, &second) + .await + .unwrap(), + InvocationResultIndexLookup::Found(OplogIndex::from_u64(5)) + ); + } + + #[test] + async fn invocation_result_index_ahead_of_status_is_not_reset() { + let key = idempotency_key("completed"); + let (service, oplog, owned_agent_id) = + index_test_service(invocation_entries(std::slice::from_ref(&key))); + let complete = invocation_status(3, 2, &[(&key, 3)], 0, DeletedRegions::new()); + service + .catch_up_invocation_result_index(&owned_agent_id, AgentMode::Durable, &complete) + .await + .unwrap(); + oplog.clear_reads(); + + let stale = invocation_status(2, 2, &[], 0, DeletedRegions::new()); + service + .catch_up_invocation_result_index(&owned_agent_id, AgentMode::Durable, &stale) + .await + .unwrap(); + + assert!(oplog.read_starts().is_empty()); + assert_eq!( + invocation_index_metadata(&service, &owned_agent_id) + .await + .covered_through, + OplogIndex::from_u64(3) + ); + } + + #[test] + async fn invocation_result_index_and_exact_membership_are_jointly_complete() { + let first = idempotency_key("first"); + let second = idempotency_key("second"); + let third = idempotency_key("third"); + let fourth = idempotency_key("fourth"); + let fifth = idempotency_key("fifth"); + let missing = idempotency_key("missing"); + let (service, _oplog, owned_agent_id) = index_test_service(invocation_entries(&[ + first.clone(), + second.clone(), + third.clone(), + fourth.clone(), + fifth.clone(), + ])); + let indexed = invocation_status( + 9, + 2, + &[(&first, 3), (&second, 5), (&third, 7), (&fourth, 9)], + 0, + DeletedRegions::new(), + ); + service + .catch_up_invocation_result_index(&owned_agent_id, AgentMode::Durable, &indexed) + .await + .unwrap(); + + let current = invocation_status( + 11, + 2, + &[ + (&first, 3), + (&second, 5), + (&third, 7), + (&fourth, 9), + (&fifth, 11), + ], + 0, + DeletedRegions::new(), + ); + assert_eq!( + current.invocation_results.oldest_retained_index(), + Some(OplogIndex::from_u64(9)) + ); + assert_eq!( + service + .lookup_invocation_result_index(&owned_agent_id, ¤t, &missing) + .await + .unwrap(), + InvocationResultIndexLookup::DefinitiveMiss + ); + } + + #[test] + async fn incomplete_invocation_result_index_does_not_return_an_obsolete_result() { + let repeated = idempotency_key("repeated"); + let (service, _oplog, owned_agent_id) = + index_test_service(invocation_entries(&[repeated.clone(), repeated.clone()])); + let partial = invocation_status(3, 0, &[(&repeated, 3)], 0, DeletedRegions::new()); + service + .catch_up_invocation_result_index(&owned_agent_id, AgentMode::Durable, &partial) + .await + .unwrap(); + + let current = invocation_status( + 5, + 0, + &[(&repeated, 3), (&repeated, 5)], + 0, + DeletedRegions::new(), + ); + assert_eq!( + service + .lookup_invocation_result_index(&owned_agent_id, ¤t, &repeated) + .await + .unwrap(), + InvocationResultIndexLookup::Incomplete + ); + + service + .catch_up_invocation_result_index(&owned_agent_id, AgentMode::Durable, ¤t) + .await + .unwrap(); + assert_eq!( + service + .lookup_invocation_result_index(&owned_agent_id, ¤t, &repeated) + .await + .unwrap(), + InvocationResultIndexLookup::Found(OplogIndex::from_u64(5)) + ); + } + + #[test] + async fn concurrent_generation_catch_up_clears_reverted_results() { + let reverted = idempotency_key("reverted"); + let current = idempotency_key("current"); + let mut entries = invocation_entries(std::slice::from_ref(&reverted)); + entries.insert( + OplogIndex::from_u64(4), + OplogEntry::revert(OplogRegion::from_index_range( + OplogIndex::from_u64(2)..=OplogIndex::from_u64(3), + )), + ); + invocation_pair(&mut entries, 5, ¤t); + let (service, oplog, owned_agent_id) = index_test_service(entries); + let old_status = invocation_status(3, 2, &[(&reverted, 3)], 0, DeletedRegions::new()); + let deleted_regions = DeletedRegions::from_regions([OplogRegion::from_index_range( + OplogIndex::from_u64(2)..=OplogIndex::from_u64(3), + )]); + let new_status = invocation_status(6, 2, &[(¤t, 6)], 1, deleted_regions); + + oplog.pause_next_read(); + let old_task = tokio::spawn({ + let service = service.clone(); + let owned_agent_id = owned_agent_id.clone(); + async move { + service + .catch_up_invocation_result_index( + &owned_agent_id, + AgentMode::Durable, + &old_status, + ) + .await + } + }); + oplog.read_started.notified().await; + let new_task = tokio::spawn({ + let service = service.clone(); + let owned_agent_id = owned_agent_id.clone(); + let new_status = new_status.clone(); + async move { + service + .catch_up_invocation_result_index( + &owned_agent_id, + AgentMode::Durable, + &new_status, + ) + .await + } + }); + tokio::task::yield_now().await; + assert!(!new_task.is_finished()); + oplog.resume_read.notify_one(); + old_task.await.unwrap().unwrap(); + new_task.await.unwrap().unwrap(); + + let metadata = invocation_index_metadata(&service, &owned_agent_id).await; + assert_eq!(metadata.revert_generation, 1); + assert_eq!(metadata.covered_through, OplogIndex::from_u64(6)); + assert_eq!( + service + .lookup_invocation_result_index(&owned_agent_id, &new_status, &reverted) + .await + .unwrap(), + InvocationResultIndexLookup::DefinitiveMiss + ); + assert_eq!( + service + .lookup_invocation_result_index(&owned_agent_id, &new_status, ¤t) + .await + .unwrap(), + InvocationResultIndexLookup::Found(OplogIndex::from_u64(6)) + ); + let fields = service + .key_value_storage + .with("test", "list_invocation_result_index") + .keys(DefaultWorkerService::invocation_result_index_namespace( + &owned_agent_id.agent_id, + )) + .await + .unwrap(); + assert!(!fields.contains(&invocation_result_index_field(&reverted))); + } + + #[test] + async fn cancelled_catch_up_releases_invocation_result_index_lock_registration() { + let key = idempotency_key("completed"); + let (service, oplog, owned_agent_id) = + index_test_service(invocation_entries(std::slice::from_ref(&key))); + let status = invocation_status(3, 1, &[(&key, 3)], 0, DeletedRegions::new()); + + oplog.pause_next_read(); + let catch_up = tokio::spawn({ + let service = service.clone(); + let owned_agent_id = owned_agent_id.clone(); + async move { + service + .catch_up_invocation_result_index(&owned_agent_id, AgentMode::Durable, &status) + .await + } + }); + oplog.read_started.notified().await; + + catch_up.abort(); + assert!(catch_up.await.unwrap_err().is_cancelled()); + + assert!( + !service + .invocation_result_index_locks + .lock() + .unwrap() + .contains_key(&owned_agent_id), + "a cancelled catch-up must not retain one dead lock registration per agent" + ); + } + #[test] fn tracks_workers_with_pending_invocations_for_assignment_recovery() { let mut status = AgentStatusRecord { diff --git a/golem-worker-executor/src/storage/keyvalue/mod.rs b/golem-worker-executor/src/storage/keyvalue/mod.rs index 0a7d08122d..3ac7c0fa45 100644 --- a/golem-worker-executor/src/storage/keyvalue/mod.rs +++ b/golem-worker-executor/src/storage/keyvalue/mod.rs @@ -697,19 +697,25 @@ pub enum KeyValueStorageNamespace { }, /// Per-agent cached status. Unlike `Worker` (a flat key space), this namespace is stored as /// one structure-per-agent (a Redis hash) so the cached `AgentStatusRecord` can be split into - /// independently written fields: a small fixed-size `core`, the `regions`, the `updates`, and - /// one field per idempotency key (`ir:{key}`). This keeps the per-commit write small and - /// decoupled from the unbounded parts of the status. The `agent_id` is part of the namespace so - /// each agent gets its own isolated key space (enabling per-agent `keys`/`del_many`). + /// independently written fields: a small `core`, the bounded invocation-result `membership`, + /// the `regions`, the `updates`, and one field per received card transfer. The `agent_id` is + /// part of the namespace so each agent gets its own isolated key space (enabling per-agent + /// `keys`/`del_many`). AgentStatus { agent_id: AgentId, }, + /// Per-agent invocation result index. Uses the same hash-style layout and cache routing as + /// [`Self::AgentStatus`], but has an independent physical namespace. + AgentInvocationResultIndex { + agent_id: AgentId, + }, /// Per-agent *clean* cached status checkpoint. Same physical layout as [`Self::AgentStatus`] - /// (one structure-per-agent split into `core` / `regions` / `updates` / `ir:{key}`), but - /// written only at structurally clean boundaries (snapshot save, throttled idle) where no - /// jumpable oplog region is open. It is never advanced by the background status flusher, so it - /// always holds a baseline before any later jump region and lets the status recompute fold - /// forward from it instead of re-reading the whole oplog from index 1. + /// (one structure-per-agent split into `core` / `membership` / `regions` / `updates` and + /// per-transfer fields), but written only at structurally clean boundaries (snapshot save, + /// throttled idle) where no jumpable oplog region is open. It is never advanced by the + /// background status flusher, so it always holds a baseline before any later jump region and + /// lets the status recompute fold forward from it instead of re-reading the whole oplog from + /// index 1. AgentStatusCheckpoint { agent_id: AgentId, }, diff --git a/golem-worker-executor/src/storage/keyvalue/multi_sqlite.rs b/golem-worker-executor/src/storage/keyvalue/multi_sqlite.rs index 5907a5700f..9b3bc9f508 100644 --- a/golem-worker-executor/src/storage/keyvalue/multi_sqlite.rs +++ b/golem-worker-executor/src/storage/keyvalue/multi_sqlite.rs @@ -114,6 +114,9 @@ impl MultiSqliteKeyValueStorage { KeyValueStorageNamespace::AgentStatus { agent_id } => { format!("kv-worker-{}.db", self.agent_id_hash(agent_id).await) } + KeyValueStorageNamespace::AgentInvocationResultIndex { agent_id } => { + format!("kv-worker-{}.db", self.agent_id_hash(agent_id).await) + } KeyValueStorageNamespace::AgentStatusCheckpoint { agent_id } => { format!("kv-worker-{}.db", self.agent_id_hash(agent_id).await) } diff --git a/golem-worker-executor/src/storage/keyvalue/namespace_routed.rs b/golem-worker-executor/src/storage/keyvalue/namespace_routed.rs index 17afd49840..073e7db59a 100644 --- a/golem-worker-executor/src/storage/keyvalue/namespace_routed.rs +++ b/golem-worker-executor/src/storage/keyvalue/namespace_routed.rs @@ -38,6 +38,7 @@ impl NamespaceRoutedKeyValueStorage { match namespace { KeyValueStorageNamespace::Worker { .. } => &self.cache, KeyValueStorageNamespace::AgentStatus { .. } => &self.cache, + KeyValueStorageNamespace::AgentInvocationResultIndex { .. } => &self.cache, KeyValueStorageNamespace::AgentStatusCheckpoint { .. } => &self.cache, _ => &self.persistent, } diff --git a/golem-worker-executor/src/storage/keyvalue/postgres.rs b/golem-worker-executor/src/storage/keyvalue/postgres.rs index 388fa9b943..79368966ca 100644 --- a/golem-worker-executor/src/storage/keyvalue/postgres.rs +++ b/golem-worker-executor/src/storage/keyvalue/postgres.rs @@ -78,6 +78,9 @@ impl PostgresKeyValueStorage { KeyValueStorageNamespace::AgentStatus { agent_id } => { format!("agent-status:{}", agent_id.to_redis_key()) } + KeyValueStorageNamespace::AgentInvocationResultIndex { agent_id } => { + format!("agent-invocation-result-index:{}", agent_id.to_redis_key()) + } KeyValueStorageNamespace::AgentStatusCheckpoint { agent_id } => { format!("agent-status-checkpoint:{}", agent_id.to_redis_key()) } diff --git a/golem-worker-executor/src/storage/keyvalue/redis.rs b/golem-worker-executor/src/storage/keyvalue/redis.rs index 46c150cb27..92bd62787d 100644 --- a/golem-worker-executor/src/storage/keyvalue/redis.rs +++ b/golem-worker-executor/src/storage/keyvalue/redis.rs @@ -39,6 +39,10 @@ impl RedisKeyValueStorage { KeyValueStorageNamespace::AgentStatus { agent_id } => { Some(format!("agent-status:{}", agent_id.to_redis_key())) } + KeyValueStorageNamespace::AgentInvocationResultIndex { agent_id } => Some(format!( + "agent-invocation-result-index:{}", + agent_id.to_redis_key() + )), // Per-agent clean checkpoint hash; same per-agent isolation as `AgentStatus`. KeyValueStorageNamespace::AgentStatusCheckpoint { agent_id } => Some(format!( "agent-status-checkpoint:{}", diff --git a/golem-worker-executor/src/storage/keyvalue/sqlite.rs b/golem-worker-executor/src/storage/keyvalue/sqlite.rs index 6ed3a85c11..27e63510a5 100644 --- a/golem-worker-executor/src/storage/keyvalue/sqlite.rs +++ b/golem-worker-executor/src/storage/keyvalue/sqlite.rs @@ -73,6 +73,9 @@ impl SqliteKeyValueStorage { KeyValueStorageNamespace::AgentStatus { agent_id } => { format!("agent-status:{}", agent_id.to_redis_key()) } + KeyValueStorageNamespace::AgentInvocationResultIndex { agent_id } => { + format!("agent-invocation-result-index:{}", agent_id.to_redis_key()) + } KeyValueStorageNamespace::AgentStatusCheckpoint { agent_id } => { format!("agent-status-checkpoint:{}", agent_id.to_redis_key()) } diff --git a/golem-worker-executor/src/worker/invocation_loop.rs b/golem-worker-executor/src/worker/invocation_loop.rs index 2be4c02473..59ac90ec49 100644 --- a/golem-worker-executor/src/worker/invocation_loop.rs +++ b/golem-worker-executor/src/worker/invocation_loop.rs @@ -13,7 +13,7 @@ // limitations under the License. use crate::durable_host::tool::operation::OwnerFailureWinner; -use crate::model::{ReadFileResult, TrapType}; +use crate::model::{LookupResult, ReadFileResult, TrapType}; use crate::sandbox_filesystem::{SandboxFilesystem, SandboxFilesystemAdapter}; use crate::services::agent_filesystem::{ LimitTransition, ResidentFilesystem, ResidentFilesystemActivity, SealedFilesystem, @@ -1999,10 +1999,10 @@ impl Invocation<'_, Ctx> { } invocation => { if let Some(idempotency_key) = invocation.idempotency_key() { - let has_result = { - let invocation_results = self.parent.invocation_results.read().await; - invocation_results.contains_key(idempotency_key) - }; + let has_result = matches!( + self.parent.lookup_invocation_result(idempotency_key).await, + LookupResult::Complete(_) | LookupResult::Interrupted + ); if !has_result { self.invoke_agent(invocation).await } else { diff --git a/golem-worker-executor/src/worker/mod.rs b/golem-worker-executor/src/worker/mod.rs index 23207f6422..2f72671d00 100644 --- a/golem-worker-executor/src/worker/mod.rs +++ b/golem-worker-executor/src/worker/mod.rs @@ -71,7 +71,7 @@ use crate::services::oplog::plugin::ForwardingOplog; use crate::services::oplog::{CommitLevel, Oplog, OplogOps, downcast_oplog}; use crate::services::resource_limits::AtomicResourceEntry; use crate::services::resource_usage_metering::ResourceUsageAccount; -use crate::services::worker::GetWorkerMetadataResult; +use crate::services::worker::{GetWorkerMetadataResult, InvocationResultIndexLookup}; use crate::services::worker_event::{WorkerEventService, WorkerEventServiceDefault}; use crate::services::{ All, HasActiveAgents, HasAgentTypesService, HasAgentWebhooksService, HasAll, @@ -86,7 +86,9 @@ use crate::worker::instance::{OwnerExecution, OwnerRuntimeResources}; use crate::worker::invocation_loop::{ ConcurrentAgentPermitState, InvocationLoop, run_invocation_loop_task, }; -use crate::worker::status::calculate_last_known_status_with_checkpoint; +use crate::worker::status::{ + calculate_last_known_status_with_checkpoint, fold_invocation_result_entries, +}; use crate::workerctx::{WorkerCtx, WorkerFilesystemContext}; use futures::channel::oneshot; use golem_common::base_model::agent::CachePolicy; @@ -430,7 +432,7 @@ pub struct Worker { /// [`TraceOrigin`]. external_invocation_origins: Arc>>, - invocation_results: Arc>>, + hydrated_invocation_results: Arc>, ephemeral_invocation: StdMutex, initial_worker_metadata: AgentMetadata, resource_entry: Arc, @@ -875,7 +877,6 @@ impl Worker { let current_status_snapshot = current_status.load_full(); let metrics_status = Arc::new(WorkerStatusMetric::new(current_status_snapshot.status)); - let initial_invocation_results = current_status_snapshot.invocation_results.clone(); let last_oplog_idx = current_status_snapshot.oplog_idx; drop(current_status_snapshot); @@ -887,16 +888,10 @@ impl Worker { let queue = Arc::new(RwLock::new(VecDeque::new())); let external_invocation_origins = Arc::new(RwLock::new(HashMap::new())); - let invocation_results = Arc::new(RwLock::new(HashMap::from_iter( - initial_invocation_results.iter().map(|(key, oplog_idx)| { - ( - key.clone(), - InvocationResult::Lazy { - oplog_idx: *oplog_idx, - }, - ) - }), - ))); + let hydrated_invocation_results = + Arc::new(RwLock::new(HydratedInvocationResultCache::new( + deps.config().invocation_results.hydrated_cache_capacity, + ))); let instance = Arc::new(Mutex::new(WorkerInstance::Unloaded { startup_failure: reconstructed_ephemeral.then(inactive_ephemeral_agent_error), @@ -990,7 +985,7 @@ impl Worker { deps: all_deps, queue, external_invocation_origins, - invocation_results, + hydrated_invocation_results, ephemeral_invocation: StdMutex::new(if reconstructed_ephemeral { EphemeralInvocationState::Accepted(None) } else { @@ -2455,25 +2450,28 @@ impl Worker { } } - pub async fn invocation_results(&self) -> HashMap { - self.last_known_status.load().invocation_results.clone() - } - // should only be called from invocation loop pub async fn store_invocation_success( &self, key: &IdempotencyKey, output: AgentInvocationOutput, ) { - let mut map = self.invocation_results.write().await; + let mut map = self.hydrated_invocation_results.write().await; map.insert( key.clone(), InvocationResult::Cached { result: Ok(output.clone()), }, + self.last_known_status + .load() + .invocation_results + .revert_generation(), + output + .oplog_index + .unwrap_or_else(|| self.last_known_status.load().oplog_idx), ); // `drop` before taking `origins`: `fail_pending_invocations` locks - // origins -> invocation_results, so holding `map` here would invert that + // origins -> hydrated_invocation_results, so holding `map` here would invert that // order and can deadlock. Not a scope tidy-up. drop(map); self.external_invocation_origins.write().await.remove(key); @@ -2500,7 +2498,9 @@ impl Worker { invocation_keys_to_fail(&status, Some(key), !trap_type.is_invocation_rejection()); let stderr = self.worker_event_service.get_last_invocation_errors(); let golem_error = trap_type.as_golem_error(&stderr); - let mut map = self.invocation_results.write().await; + let mut map = self.hydrated_invocation_results.write().await; + // Co-pending fail-fast results exist only in this bounded warm cache. Once evicted, a + // poller sees the same `Pending` state that reconstructing this status from the oplog does. for key in &keys_to_fail { map.insert( key.clone(), @@ -2510,13 +2510,15 @@ impl Worker { stderr: stderr.clone(), }), }, + status.invocation_results.revert_generation(), + status.oplog_idx, ); if let Some(golem_error) = &golem_error { self.publish_completion(key, Err(golem_error.clone())); } } // See `store_invocation_success`: origins must not be taken while - // `invocation_results` is held. + // `hydrated_invocation_results` is held. drop(map); let mut origins = self.external_invocation_origins.write().await; for key in &keys_to_fail { @@ -2525,7 +2527,7 @@ impl Worker { } pub(super) async fn store_invocation_resuming(&self, key: &IdempotencyKey) { - let mut map = self.invocation_results.write().await; + let mut map = self.hydrated_invocation_results.write().await; map.remove(key); } @@ -3162,10 +3164,9 @@ impl Worker { /// Enqueue invocation, classified by the caller. Passing `ReadOnly` for a /// mutating method would skip cache invalidation and produce stale reads. /// - /// For `ReadOnly`, returns the epoch captured under the same instance lock - /// that commits the pending entry. Populating the cache later must use - /// this captured epoch, not the current one, to avoid storing a stale - /// result under a post-mutation epoch. + /// For `ReadOnly`, returns the epoch captured before admission. Populating + /// the cache later must use this captured epoch, not the current one, to + /// avoid storing a stale result under a post-mutation epoch. pub(crate) async fn enqueue_worker_invocation_with_effect( &self, invocation: AgentInvocation, @@ -3206,21 +3207,10 @@ impl Worker { return Err(err.clone()); } - if let Some(idempotency_key) = invocation.idempotency_key() { - let has_result = self - .invocation_results - .read() - .await - .contains_key(idempotency_key); - let status = self.last_known_status.load(); - let is_pending = status - .pending_invocations - .iter() - .any(|entry| entry.has_idempotency_key(idempotency_key)); - let is_current = status.current_idempotency_key.as_ref() == Some(idempotency_key); - if has_result || is_pending || is_current { - return Ok(None); - } + if let Some(idempotency_key) = invocation.idempotency_key() + && self.lookup_invocation_result(idempotency_key).await != LookupResult::New + { + return Ok(None); } let ( @@ -3249,15 +3239,13 @@ impl Worker { invocation_context_spans, ); - // Snapshot the epoch under the instance lock that commits the - // pending entry. Read-only captures the current epoch for later - // cache fill. Mutating invocations no longer bump here — the bump - // happens on *successful completion* in + // Snapshot the epoch for a later read-only cache fill. Keyed admission releases and + // reacquires the instance lock below; that staleness is safe because + // `populate_read_only_cache` rechecks the epoch before publishing the result. Mutating + // invocations no longer bump here — the bump happens on *successful completion* in // `DurableWorkerCtx::on_agent_invocation_success`, so a cached // read-only result stays serviceable while the mutation is queued - // / running. The populate-time recheck in - // `populate_read_only_cache` covers the race where the mutation - // completes before the read-only observer fills the cache. + // or running. let read_only_epoch_snapshot = match read_only_cache_effect { read_only_cache::InvocationEffect::ReadOnly => { Some(self.read_only_cache_epoch.load(Ordering::SeqCst)) @@ -3266,8 +3254,51 @@ impl Worker { | read_only_cache::InvocationEffect::UnknownAssumeMutating => None, }; - self.add_and_commit_oplog_internal(&instance_guard, entry, None) + let mut caller_instance_guard = Some(instance_guard); + if let Some(idempotency_key) = semantic_idempotency_key.as_ref() { + drop(caller_instance_guard.take()); + loop { + let status = self.last_known_status.load_full(); + if self.lookup_invocation_result(idempotency_key).await != LookupResult::New { + return Ok(None); + } + let current = self.last_known_status.load(); + if current.invocation_results.change_generation() + != status.invocation_results.change_generation() + || current.invocation_results.revert_generation() + != status.invocation_results.revert_generation() + { + continue; + } + let instance_guard = self.lock_non_stopping_worker_owned().await; + if instance_guard.is_deleting() { + return Err(WorkerExecutorError::invalid_request( + "Cannot enqueue invocation to a deleting worker", + )); + } + if !self + .state_actor + .append_invocation_if_version( + entry.clone(), + idempotency_key.clone(), + status.invocation_results.change_generation(), + status.invocation_results.revert_generation(), + instance_guard, + ) + .await + { + continue; + } + break; + } + } else { + self.add_and_commit_oplog_internal( + caller_instance_guard.as_ref().unwrap(), + entry, + None, + ) .await; + } if let Some(idempotency_key) = semantic_idempotency_key { // Captured here, inside the producer span, because a consumer links @@ -3285,11 +3316,13 @@ impl Worker { .insert(idempotency_key, origin); } - if let WorkerInstance::Running(running) = &*instance_guard { + if let Some(instance_guard) = caller_instance_guard.as_ref() + && let WorkerInstance::Running(running) = &**instance_guard + { running.sender.send(WorkerCommand::WorkAvailable).unwrap(); }; - drop(instance_guard); + drop(caller_instance_guard); Ok(read_only_epoch_snapshot) } @@ -4395,31 +4428,8 @@ impl Worker { } } - let status = self.get_last_known_status().await; for idempotency_key in unfinished { - let mut invocation_result = { - self.invocation_results - .read() - .await - .get(&idempotency_key) - .cloned() - }; - let Some(invocation_result) = invocation_result.as_mut() else { - continue; - }; - invocation_result - .cache( - &self.owned_agent_id, - self.agent_mode(), - self.initial_worker_metadata.fingerprint, - self, - ) - .await; - match lookup_result_from_cached_result( - &status, - &idempotency_key, - invocation_result.clone(), - ) { + match self.lookup_invocation_result(&idempotency_key).await { LookupResult::Complete(Ok(_)) => { self.complete_durable_streaming_session(&idempotency_key) .await?; @@ -5413,21 +5423,35 @@ impl Worker { pub async fn lookup_invocation_result(&self, key: &IdempotencyKey) -> LookupResult { let status = self.last_known_status.load_full().as_ref().clone(); - let maybe_result = self - .invocation_results + let cached = self + .hydrated_invocation_results .read() .await - .get(key) - .cloned() - .or_else(|| { - status - .invocation_results - .get(key) - .map(|oplog_idx| InvocationResult::Lazy { - oplog_idx: *oplog_idx, - }) - }); - if let Some(mut result) = maybe_result { + .get_valid(key, &status) + .map(|(result, oplog_idx)| (result.clone(), oplog_idx)); + let maybe_result = if cached.is_some() { + crate::metrics::workers::record_invocation_result_resolution("memory_exact"); + cached + } else if let Some(oplog_idx) = status.invocation_results.get(key) { + crate::metrics::workers::record_invocation_result_resolution("memory_exact"); + Some(( + InvocationResult::Lazy { + oplog_idx: *oplog_idx, + }, + *oplog_idx, + )) + } else if status.invocation_results.is_exact_complete() { + crate::metrics::workers::record_invocation_result_resolution("memory_exact_miss"); + None + } else if !status.invocation_results.might_contain(key) { + crate::metrics::workers::record_invocation_result_resolution("bloom_negative"); + None + } else { + self.resolve_old_invocation_result_index(&status, key) + .await + .map(|oplog_idx| (InvocationResult::Lazy { oplog_idx }, oplog_idx)) + }; + if let Some((mut result, result_oplog_idx)) = maybe_result { result .cache( &self.owned_agent_id, @@ -5436,6 +5460,12 @@ impl Worker { self, ) .await; + self.hydrated_invocation_results.write().await.insert( + key.clone(), + result.clone(), + status.invocation_results.revert_generation(), + result_oplog_idx, + ); lookup_result_from_cached_result(&status, key, result) } else { let is_pending = status @@ -5451,6 +5481,87 @@ impl Worker { } } + async fn resolve_old_invocation_result_index( + &self, + status: &AgentStatusRecord, + key: &IdempotencyKey, + ) -> Option { + let worker_service = self.deps.worker_service(); + let lookup = worker_service + .lookup_invocation_result_index(&self.owned_agent_id, status, key) + .await; + match lookup { + Ok(InvocationResultIndexLookup::Found(index)) => { + crate::metrics::workers::record_invocation_result_resolution("physical_hit"); + return Some(index); + } + Ok(InvocationResultIndexLookup::DefinitiveMiss) => { + crate::metrics::workers::record_invocation_result_resolution("physical_miss"); + return None; + } + Ok(InvocationResultIndexLookup::Incomplete) | Err(_) => { + crate::metrics::workers::record_invocation_result_resolution("physical_incomplete"); + } + } + + if worker_service + .catch_up_invocation_result_index(&self.owned_agent_id, self.agent_mode(), status) + .await + .is_ok() + { + match worker_service + .lookup_invocation_result_index(&self.owned_agent_id, status, key) + .await + { + Ok(InvocationResultIndexLookup::Found(index)) => { + crate::metrics::workers::record_invocation_result_resolution("physical_hit"); + return Some(index); + } + Ok(InvocationResultIndexLookup::DefinitiveMiss) => { + crate::metrics::workers::record_invocation_result_resolution("physical_miss"); + return None; + } + Ok(InvocationResultIndexLookup::Incomplete) | Err(_) => {} + } + } + + crate::metrics::workers::record_invocation_result_resolution("oplog_fallback"); + let mut current_idempotency_key = None; + let mut cancelled_idempotency_key = None; + let mut result = None; + let mut first = OplogIndex::INITIAL; + let chunk_size = self + .deps + .config() + .invocation_results + .physical_index_catch_up_chunk_size + .max(1); + while first <= status.oplog_idx { + let count = (status.oplog_idx.as_u64() - first.as_u64() + 1).min(chunk_size); + let entries = self + .deps + .oplog_service() + .read_exact(&self.owned_agent_id, self.agent_mode(), first, count) + .await; + if entries.is_empty() { + break; + } + fold_invocation_result_entries( + &mut current_idempotency_key, + &mut cancelled_idempotency_key, + &status.deleted_regions, + &entries, + |candidate, index| { + if candidate == key { + result = Some(index); + } + }, + ); + first = entries.keys().max().unwrap().next(); + } + result + } + async fn stop_internal( &self, called_from_invocation_loop: bool, @@ -5808,9 +5919,12 @@ impl Worker { let status = self.last_known_status.load_full().as_ref().clone(); let keys_to_fail = invocation_keys_to_fail(&status, None, true); - let mut invocation_results = self.invocation_results.write().await; + let mut invocation_results = self.hydrated_invocation_results.write().await; for idempotency_key in &keys_to_fail { - if invocation_results.contains_key(idempotency_key) { + if invocation_results + .get_valid(idempotency_key, &status) + .is_some() + { continue; } invocation_results.insert( @@ -5829,6 +5943,8 @@ impl Worker { stderr: String::new(), }), }, + status.invocation_results.revert_generation(), + status.oplog_idx, ); self.publish_completion(idempotency_key, Err(error.clone())); origins.remove(idempotency_key); @@ -6091,6 +6207,7 @@ impl Worker { .iter() .map(|i| i.environment_plugin_grant_id) .collect(), + invocation_results: this.config().invocation_results.membership(), agent_mode, ..Default::default() }; @@ -6210,6 +6327,7 @@ impl Worker { // TODO: should be private, exposed for the invocation loop for now. pub async fn reattach_worker_status(&self) { + self.hydrated_invocation_results.write().await.clear(); self.state_actor.reattach_worker_status().await; } @@ -7667,6 +7785,86 @@ enum InvocationResult { }, } +/// Bounded cache of invocation result payloads loaded from the oplog. The status membership only +/// stores result oplog indexes; this cache avoids repeatedly loading and decoding those entries. +struct HydratedInvocationResultCache { + values: HashMap, + insertion_order: VecDeque, + capacity: usize, +} + +struct HydratedInvocationResultCacheEntry { + result: InvocationResult, + /// Oplog branch generation in which this result was produced. + revert_generation: u64, + oplog_idx: OplogIndex, +} + +impl HydratedInvocationResultCache { + fn new(capacity: usize) -> Self { + Self { + values: HashMap::new(), + insertion_order: VecDeque::new(), + capacity, + } + } + + fn get_valid( + &self, + key: &IdempotencyKey, + status: &AgentStatusRecord, + ) -> Option<(&InvocationResult, OplogIndex)> { + self.values.get(key).and_then(|entry| { + (entry.revert_generation == status.invocation_results.revert_generation() + && !status.deleted_regions.is_in_deleted_region(entry.oplog_idx)) + .then_some((&entry.result, entry.oplog_idx)) + }) + } + + #[cfg(test)] + fn contains_key(&self, key: &IdempotencyKey) -> bool { + self.values.contains_key(key) + } + + fn insert( + &mut self, + key: IdempotencyKey, + value: InvocationResult, + revert_generation: u64, + oplog_idx: OplogIndex, + ) { + if !self.values.contains_key(&key) { + self.insertion_order.push_back(key.clone()); + } + self.values.insert( + key, + HydratedInvocationResultCacheEntry { + result: value, + revert_generation, + oplog_idx, + }, + ); + while self.values.len() > self.capacity { + if let Some(oldest) = self.insertion_order.pop_front() { + self.values.remove(&oldest); + } + } + } + + fn remove(&mut self, key: &IdempotencyKey) -> Option { + let result = self.values.remove(key); + if result.is_some() { + self.insertion_order.retain(|entry| entry != key); + } + result + } + + fn clear(&mut self) { + self.values.clear(); + self.insertion_order.clear(); + } +} + impl InvocationResult { pub async fn cache( &mut self, @@ -8200,6 +8398,33 @@ mod tests { assert!(reconstructed.accept(&first).is_err()); } + #[test] + fn hydrated_invocation_result_cache_is_bounded() { + let first = IdempotencyKey::new("first".to_string()); + let second = IdempotencyKey::new("second".to_string()); + let mut cache = HydratedInvocationResultCache::new(1); + cache.insert( + first.clone(), + InvocationResult::Lazy { + oplog_idx: OplogIndex::from_u64(2), + }, + 0, + OplogIndex::from_u64(2), + ); + + cache.insert( + second.clone(), + InvocationResult::Lazy { + oplog_idx: OplogIndex::from_u64(4), + }, + 0, + OplogIndex::from_u64(4), + ); + + assert!(!cache.contains_key(&first)); + assert!(cache.contains_key(&second)); + } + #[test] fn ephemeral_invocation_requires_a_final_phantom_id_even_when_it_may_exist() { let result = validate_resolved_invocation_identity( diff --git a/golem-worker-executor/src/worker/state_actor.rs b/golem-worker-executor/src/worker/state_actor.rs index 3dec7c891c..787294f563 100644 --- a/golem-worker-executor/src/worker/state_actor.rs +++ b/golem-worker-executor/src/worker/state_actor.rs @@ -64,7 +64,7 @@ use golem_common::model::account::AccountId; use golem_common::model::agent::AgentMode; use golem_common::model::oplog::{OplogEntry, OplogIndex}; use golem_common::model::{ - AgentStatus, AgentStatusRecord, OwnedAgentId, ScheduledAction, Timestamp, + AgentStatus, AgentStatusRecord, IdempotencyKey, OwnedAgentId, ScheduledAction, Timestamp, }; use golem_service_base::error::worker_executor::InterruptKind; use std::any::Any; @@ -118,6 +118,14 @@ enum StatusJob { _card_event_boundary_guard: OwnedMutexGuard<()>, done: oneshot::Sender<()>, }, + AppendInvocationIfVersion { + entry: Box, + idempotency_key: IdempotencyKey, + expected_result_generation: u64, + expected_revert_generation: u64, + instance_guard: OwnedMutexGuard, + done: oneshot::Sender, + }, /// Returns the published status after reattaching it when a jump or revert detached it. /// Serialization on the status queue prevents observing an in-flight status transition. AttachedStatus { @@ -256,6 +264,40 @@ impl WorkerStateActor { ) .await; } + StatusJob::AppendInvocationIfVersion { + entry, + idempotency_key, + expected_result_generation, + expected_revert_generation, + instance_guard, + done, + } => { + complete_status_job( + async { + state.ensure_status_attached().await; + let status = state.last_known_status.load(); + if !can_append_invocation( + &status, + &idempotency_key, + expected_result_generation, + expected_revert_generation, + ) { + return false; + } + drop(status); + state.oplog.add(*entry).await; + state + .commit_and_update_state(CommitLevel::Always, None) + .await; + if let WorkerInstance::Running(running) = &*instance_guard { + running.sender.send(WorkerCommand::WorkAvailable).unwrap(); + } + true + }, + done, + ) + .await; + } StatusJob::AttachedStatus { done } => { if state.detached.load(Ordering::Acquire) { state.reattach().await; @@ -380,6 +422,26 @@ impl WorkerStateActor { .await } + pub async fn append_invocation_if_version( + &self, + entry: OplogEntry, + idempotency_key: IdempotencyKey, + expected_result_generation: u64, + expected_revert_generation: u64, + instance_guard: OwnedMutexGuard, + ) -> bool { + self.commit + .run_status_job(|done| StatusJob::AppendInvocationIfVersion { + entry: Box::new(entry), + idempotency_key, + expected_result_generation, + expected_revert_generation, + instance_guard, + done, + }) + .await + } + pub async fn attached_status(&self) -> Arc { self.commit .run_status_job(|done| StatusJob::AttachedStatus { done }) @@ -697,14 +759,126 @@ fn is_authority_state_entry(entry: &OplogEntry) -> bool { ) } +fn can_append_invocation( + status: &AgentStatusRecord, + idempotency_key: &IdempotencyKey, + expected_result_generation: u64, + expected_revert_generation: u64, +) -> bool { + status.invocation_results.change_generation() == expected_result_generation + && status.invocation_results.revert_generation() == expected_revert_generation + && !status.invocation_results.contains_key(idempotency_key) + && status.current_idempotency_key.as_ref() != Some(idempotency_key) + && !status + .pending_invocations + .iter() + .any(|invocation| invocation.has_idempotency_key(idempotency_key)) +} + #[cfg(test)] mod tests { - use super::complete_status_job; + use super::{can_append_invocation, complete_status_job}; + use golem_common::model::oplog::OplogIndex; + use golem_common::model::{AgentStatusRecord, IdempotencyKey, PendingInvocationRef, Timestamp}; use std::sync::Arc; use std::sync::atomic::{AtomicU64, Ordering}; use test_r::test; use tokio::sync::{Notify, oneshot}; + fn pending_invocation(key: IdempotencyKey) -> PendingInvocationRef { + PendingInvocationRef { + timestamp: Timestamp::now_utc(), + oplog_index: OplogIndex::INITIAL, + idempotency_key: Some(key), + manual_update_target_revision: None, + } + } + + #[test] + fn invocation_admission_ignores_unrelated_oplog_and_pending_changes() { + let key = IdempotencyKey::fresh(); + let mut status = AgentStatusRecord { + oplog_idx: OplogIndex::from_u64(100), + pending_invocations: vec![pending_invocation(IdempotencyKey::fresh())], + ..AgentStatusRecord::default() + }; + let result_generation = status.invocation_results.change_generation(); + let revert_generation = status.invocation_results.revert_generation(); + + assert!(can_append_invocation( + &status, + &key, + result_generation, + revert_generation + )); + + status.oplog_idx = OplogIndex::from_u64(1_000); + status + .pending_invocations + .push(pending_invocation(IdempotencyKey::fresh())); + assert!(can_append_invocation( + &status, + &key, + result_generation, + revert_generation + )); + } + + #[test] + fn invocation_admission_rejects_same_key_or_result_branch_changes() { + let key = IdempotencyKey::fresh(); + let mut status = AgentStatusRecord::default(); + let result_generation = status.invocation_results.change_generation(); + let revert_generation = status.invocation_results.revert_generation(); + + status.pending_invocations = vec![pending_invocation(key.clone())]; + assert!(!can_append_invocation( + &status, + &key, + result_generation, + revert_generation + )); + + status.pending_invocations.clear(); + status.current_idempotency_key = Some(key.clone()); + assert!(!can_append_invocation( + &status, + &key, + result_generation, + revert_generation + )); + + status.current_idempotency_key = None; + status + .invocation_results + .insert(IdempotencyKey::fresh(), OplogIndex::INITIAL); + assert!(!can_append_invocation( + &status, + &key, + result_generation, + revert_generation + )); + + status + .invocation_results + .insert(key.clone(), OplogIndex::from_u64(2)); + let key_result_generation = status.invocation_results.change_generation(); + assert!(!can_append_invocation( + &status, + &key, + key_result_generation, + revert_generation + )); + + status.invocation_results.set_revert_generation(1); + assert!(!can_append_invocation( + &status, + &key, + key_result_generation, + revert_generation + )); + } + #[test] async fn authority_publication_survives_producer_cancellation_before_actor_reply() { let generation = Arc::new(AtomicU64::new(0)); diff --git a/golem-worker-executor/src/worker/status.rs b/golem-worker-executor/src/worker/status.rs index 3416d8cb81..facba842fc 100644 --- a/golem-worker-executor/src/worker/status.rs +++ b/golem-worker-executor/src/worker/status.rs @@ -10,9 +10,9 @@ use golem_common::model::oplog::{ use golem_common::model::regions::{DeletedRegions, DeletedRegionsBuilder, OplogRegion}; use golem_common::model::{ AgentResourceDescription, AgentStatus, AgentStatusRecord, FailedUpdateRecord, IdempotencyKey, - OplogProcessorCheckpointState, OwnedAgentId, PendingCardEventRef, PendingInvocationRef, - PendingUpdateKind, PendingUpdateRef, ReceivedCardTransferIndex, ReceivedCardTransferState, - RetryConfig, RetryPolicyState, SuccessfulUpdateRecord, Timestamp, + InvocationResultMembership, OplogProcessorCheckpointState, OwnedAgentId, PendingCardEventRef, + PendingInvocationRef, PendingUpdateKind, PendingUpdateRef, ReceivedCardTransferIndex, + ReceivedCardTransferState, RetryConfig, RetryPolicyState, SuccessfulUpdateRecord, Timestamp, }; use golem_common::serialization::deserialize; use std::collections::{BTreeMap, HashMap, HashSet, VecDeque}; @@ -169,11 +169,16 @@ pub async fn try_fold_status_from( this: &T, owned_agent_id: &OwnedAgentId, agent_mode: AgentMode, - baseline: AgentStatusRecord, + mut baseline: AgentStatusRecord, ) -> Option where T: HasOplogService + HasConfig + HasComponentService + Sync, { + let full_rebuild = baseline.oplog_idx == OplogIndex::NONE; + if full_rebuild && baseline.invocation_results.is_empty() { + baseline.invocation_results = this.config().invocation_results.membership(); + } + let last_oplog_index = this .oplog_service() .get_last_index(owned_agent_id, agent_mode) @@ -195,17 +200,143 @@ where return Some(baseline); } - let new_entries: BTreeMap = this - .oplog_service() - .read_exact( + let chunk_size = this + .config() + .invocation_results + .physical_index_catch_up_chunk_size + .max(1); + + if full_rebuild { + return fold_status_with_precomputed_regions( + this, owned_agent_id, agent_mode, - baseline.oplog_idx.next(), - last_oplog_index.as_u64() - baseline.oplog_idx.as_u64(), + baseline, + last_oplog_index, + chunk_size, ) .await; + } - update_status_with_new_entries(agent_mode, baseline, new_entries, &this.config().retry) + let original_baseline = baseline.clone(); + let mut first = baseline.oplog_idx.next(); + while first <= last_oplog_index { + let count = (last_oplog_index.as_u64() - first.as_u64() + 1).min(chunk_size); + let entries = this + .oplog_service() + .read_exact(owned_agent_id, agent_mode, first, count) + .await; + if entries.is_empty() { + return None; + } + let finalize_oplog_processor_checkpoints = + entries.keys().next_back() == Some(&last_oplog_index); + baseline = match update_status_with_new_entries_internal( + agent_mode, + baseline, + entries, + &this.config().retry, + true, + finalize_oplog_processor_checkpoints, + ) { + Some(status) => status, + None => { + return fold_status_with_precomputed_regions( + this, + owned_agent_id, + agent_mode, + original_baseline, + last_oplog_index, + chunk_size, + ) + .await; + } + }; + first = baseline.oplog_idx.next(); + } + Some(baseline) +} + +async fn fold_status_with_precomputed_regions( + this: &T, + owned_agent_id: &OwnedAgentId, + agent_mode: AgentMode, + mut baseline: AgentStatusRecord, + last_oplog_index: OplogIndex, + chunk_size: u64, +) -> Option +where + T: HasOplogService + HasConfig + HasComponentService + Sync, +{ + let start = baseline.oplog_idx.next(); + let mut deleted_regions = baseline.deleted_regions.clone(); + let mut region_entries = BTreeMap::new(); + let mut first = start; + while first <= last_oplog_index { + let count = (last_oplog_index.as_u64() - first.as_u64() + 1).min(chunk_size); + let entries = this + .oplog_service() + .read_exact(owned_agent_id, agent_mode, first, count) + .await; + if entries.is_empty() { + return None; + } + deleted_regions = calculate_deleted_regions(deleted_regions, &entries); + region_entries.extend(entries.iter().filter_map(|(index, entry)| { + let relevant = matches!( + entry, + OplogEntry::Jump { .. } + | OplogEntry::Revert { .. } + | OplogEntry::PendingUpdate { + description: UpdateDescription::SnapshotBased { .. }, + .. + } + | OplogEntry::SuccessfulUpdate { .. } + | OplogEntry::FailedUpdate { .. } + ); + relevant.then(|| (*index, entry.clone())) + })); + first = entries.keys().max().unwrap().next(); + } + + let skipped_regions = calculate_skipped_regions( + baseline.skipped_regions.clone(), + &deleted_regions, + ®ion_entries, + ); + + if baseline_is_invalidated(&baseline, &skipped_regions) { + return None; + } + baseline.deleted_regions = deleted_regions; + baseline.skipped_regions = skipped_regions; + + first = start; + while first <= last_oplog_index { + let count = (last_oplog_index.as_u64() - first.as_u64() + 1).min(chunk_size); + let entries = this + .oplog_service() + .read_exact(owned_agent_id, agent_mode, first, count) + .await; + if entries.is_empty() { + return None; + } + let finalize_oplog_processor_checkpoints = + entries.keys().next_back() == Some(&last_oplog_index); + let deleted_regions = baseline.deleted_regions.clone(); + let skipped_regions = baseline.skipped_regions.clone(); + baseline = update_status_with_precomputed_regions( + agent_mode, + baseline, + entries, + &this.config().retry, + deleted_regions, + skipped_regions, + finalize_oplog_processor_checkpoints, + ); + first = baseline.oplog_idx.next(); + } + Some(baseline) } // update a worker status with new entries. Returns None if the status cannot be calculated from the new entries alone and needs to be recalculated from the beginning. @@ -215,6 +346,24 @@ pub fn update_status_with_new_entries( new_entries: BTreeMap, // TODO: changing the retry policy will cause inconsistencies when reading existing oplogs. default_retry_policy: &RetryConfig, +) -> Option { + update_status_with_new_entries_internal( + agent_mode, + last_known, + new_entries, + default_retry_policy, + true, + true, + ) +} + +fn update_status_with_new_entries_internal( + agent_mode: AgentMode, + last_known: AgentStatusRecord, + new_entries: BTreeMap, + default_retry_policy: &RetryConfig, + validate_baseline: bool, + finalize_oplog_processor_checkpoints: bool, ) -> Option { let deleted_regions = calculate_deleted_regions(last_known.deleted_regions.clone(), &new_entries); @@ -230,33 +379,52 @@ pub fn update_status_with_new_entries( // (Note that this is a rare case - for Jumps, this is not happening if the executor successfully writes out // the new status before performing the jump; for Reverts, the status is recalculated anyway, but only once, when // the revert is applied) - if skipped_regions.is_in_deleted_region(last_known.oplog_idx) { - let last_known_skipped_regions_without_overrides = - if last_known.skipped_regions.is_overridden() { - let mut cloned = last_known.skipped_regions.clone(); - cloned.merge_override(); - cloned - } else { - last_known.skipped_regions.clone() - }; + if validate_baseline && baseline_is_invalidated(&last_known, &skipped_regions) { + return None; + } - let new_skipped_regions_without_overrides = if skipped_regions.is_overridden() { - let mut cloned = skipped_regions.clone(); - cloned.merge_override(); - cloned - } else { - skipped_regions.clone() - }; + Some(update_status_with_precomputed_regions( + agent_mode, + last_known, + new_entries, + default_retry_policy, + deleted_regions, + skipped_regions, + finalize_oplog_processor_checkpoints, + )) +} - let effective_skipped_regions_changed = - new_skipped_regions_without_overrides != last_known_skipped_regions_without_overrides; - // We might have already calculated the status with these skipped regions as an override during a snapshot update. - // No need to recompute in this case, we are already up to date. - if effective_skipped_regions_changed { - return None; - } +fn baseline_is_invalidated(baseline: &AgentStatusRecord, skipped_regions: &DeletedRegions) -> bool { + if !skipped_regions.is_in_deleted_region(baseline.oplog_idx) { + return false; } + let baseline_without_overrides = if baseline.skipped_regions.is_overridden() { + let mut cloned = baseline.skipped_regions.clone(); + cloned.merge_override(); + cloned + } else { + baseline.skipped_regions.clone() + }; + let new_without_overrides = if skipped_regions.is_overridden() { + let mut cloned = skipped_regions.clone(); + cloned.merge_override(); + cloned + } else { + skipped_regions.clone() + }; + new_without_overrides != baseline_without_overrides +} + +fn update_status_with_precomputed_regions( + agent_mode: AgentMode, + last_known: AgentStatusRecord, + new_entries: BTreeMap, + default_retry_policy: &RetryConfig, + deleted_regions: DeletedRegions, + skipped_regions: DeletedRegions, + finalize_oplog_processor_checkpoints: bool, +) -> AgentStatusRecord { let active_plugins = last_known.active_plugins.clone(); let (status, current_retry_state, overridden_retry_config) = calculate_latest_worker_status( @@ -301,12 +469,14 @@ pub fn update_status_with_new_entries( &new_entries, ); - let (invocation_results, current_idempotency_key) = calculate_invocation_results( - last_known.invocation_results, - last_known.current_idempotency_key, - &deleted_regions, - &new_entries, - ); + let (invocation_results, current_idempotency_key, cancelled_idempotency_key) = + calculate_invocation_results( + last_known.invocation_results, + last_known.current_idempotency_key, + last_known.cancelled_idempotency_key, + &deleted_regions, + &new_entries, + ); let total_linear_memory_size = calculate_total_linear_memory_size( last_known.total_linear_memory_size, @@ -326,9 +496,10 @@ pub fn update_status_with_new_entries( &active_plugins, &deleted_regions, &new_entries, + finalize_oplog_processor_checkpoints, ); - let result = AgentStatusRecord { + AgentStatusRecord { oplog_idx: new_entries .keys() .max() @@ -345,6 +516,7 @@ pub fn update_status_with_new_entries( invocation_results, received_card_transfers, current_idempotency_key, + cancelled_idempotency_key, component_revision, component_size, owned_resources, @@ -360,9 +532,7 @@ pub fn update_status_with_new_entries( last_automatic_snapshot_timestamp, last_automatic_snapshot_component_revision, agent_mode, - }; - - Some(result) + } } fn calculate_latest_worker_status( @@ -1055,19 +1225,55 @@ fn calculate_update_fields( } fn calculate_invocation_results( - invocation_results: HashMap, + invocation_results: InvocationResultMembership, current_idempotency_key: Option, + cancelled_idempotency_key: Option, deleted_regions: &DeletedRegions, entries: &BTreeMap, -) -> (HashMap, Option) { +) -> ( + InvocationResultMembership, + Option, + Option, +) { let mut invocation_results = invocation_results; + let revert_count = entries + .values() + .filter(|entry| matches!(entry, OplogEntry::Revert { .. })) + .count() as u64; + invocation_results.set_revert_generation( + invocation_results + .revert_generation() + .wrapping_add(revert_count), + ); let mut current_idempotency_key = current_idempotency_key; - let mut cancelled_idempotency_key = None; + let mut cancelled_idempotency_key = cancelled_idempotency_key; + + fold_invocation_result_entries( + &mut current_idempotency_key, + &mut cancelled_idempotency_key, + deleted_regions, + entries, + |key, index| invocation_results.insert(key.clone(), index), + ); + + ( + invocation_results, + current_idempotency_key, + cancelled_idempotency_key, + ) +} +pub(crate) fn fold_invocation_result_entries( + current_idempotency_key: &mut Option, + cancelled_idempotency_key: &mut Option, + deleted_regions: &DeletedRegions, + entries: &BTreeMap, + mut observe_result: impl FnMut(&IdempotencyKey, OplogIndex), +) { for (oplog_idx, entry) in entries { // Skipping entries in deleted regions (by revert) if deleted_regions.is_in_deleted_region(*oplog_idx) { - cancelled_idempotency_key = None; + *cancelled_idempotency_key = None; continue; } @@ -1075,50 +1281,48 @@ fn calculate_invocation_results( OplogEntry::AgentInvocationStarted { idempotency_key, .. } => { - cancelled_idempotency_key = None; - current_idempotency_key = Some(idempotency_key.clone()); + *cancelled_idempotency_key = None; + *current_idempotency_key = Some(idempotency_key.clone()); } OplogEntry::AgentInvocationFinished { .. } => { - cancelled_idempotency_key = None; - if let Some(idempotency_key) = ¤t_idempotency_key { - invocation_results.insert(idempotency_key.clone(), *oplog_idx); + *cancelled_idempotency_key = None; + if let Some(idempotency_key) = &*current_idempotency_key { + observe_result(idempotency_key, *oplog_idx); } - current_idempotency_key = None; + *current_idempotency_key = None; } OplogEntry::CancelPendingInvocation { idempotency_key, .. } => { - cancelled_idempotency_key = Some(idempotency_key.clone()); + *cancelled_idempotency_key = Some(idempotency_key.clone()); } OplogEntry::Error { error: AgentError::PermissionDenied(_), .. } => { if let Some(idempotency_key) = cancelled_idempotency_key.take() { - invocation_results.insert(idempotency_key, *oplog_idx); - } else if let Some(idempotency_key) = ¤t_idempotency_key { - invocation_results.insert(idempotency_key.clone(), *oplog_idx); + observe_result(&idempotency_key, *oplog_idx); + } else if let Some(idempotency_key) = &*current_idempotency_key { + observe_result(idempotency_key, *oplog_idx); } } OplogEntry::Error { .. } => { - cancelled_idempotency_key = None; - if let Some(idempotency_key) = ¤t_idempotency_key { - invocation_results.insert(idempotency_key.clone(), *oplog_idx); + *cancelled_idempotency_key = None; + if let Some(idempotency_key) = &*current_idempotency_key { + observe_result(idempotency_key, *oplog_idx); } } OplogEntry::Exited { .. } => { - cancelled_idempotency_key = None; - if let Some(idempotency_key) = ¤t_idempotency_key { - invocation_results.insert(idempotency_key.clone(), *oplog_idx); + *cancelled_idempotency_key = None; + if let Some(idempotency_key) = &*current_idempotency_key { + observe_result(idempotency_key, *oplog_idx); } } _ => { - cancelled_idempotency_key = None; + *cancelled_idempotency_key = None; } } } - - (invocation_results, current_idempotency_key) } fn calculate_total_linear_memory_size( @@ -1237,6 +1441,7 @@ fn calculate_oplog_processor_checkpoints( active_plugins: &HashSet, deleted_regions: &DeletedRegions, entries: &BTreeMap, + finalize: bool, ) -> HashMap { for (idx, entry) in entries { if deleted_regions.is_in_deleted_region(*idx) { @@ -1308,9 +1513,11 @@ fn calculate_oplog_processor_checkpoints( } } - result.retain(|grant_id, state| { - active_plugins.contains(grant_id) || state.sending_up_to > state.confirmed_up_to - }); + if finalize { + result.retain(|grant_id, state| { + active_plugins.contains(grant_id) || state.sending_up_to > state.confirmed_up_to + }); + } result } @@ -2234,6 +2441,25 @@ mod test { ); } + #[test] + async fn full_recompute_reads_each_oplog_chunk_twice() { + let (test_case, _checkpoint, _stale_live, final_expected) = jump_repair_fixture(); + + let result = calculate_last_known_status( + &test_case, + &test_case.owned_agent_id, + AgentMode::Durable, + None, + ) + .await; + + assert_eq!(result, Some(final_expected)); + assert_eq!( + *test_case.read_starts.lock().unwrap(), + vec![1, 3, 5, 7, 1, 3, 5, 7] + ); + } + #[test] async fn checkpoint_repair_falls_back_to_full_recompute_when_checkpoint_unusable() { let (test_case, _checkpoint, stale_live, final_expected) = jump_repair_fixture(); @@ -2475,6 +2701,10 @@ mod test { .expected_status .clone(); self.add(OplogEntry::revert(region.clone()), move |mut status| { + let revert_generation = status + .invocation_results + .revert_generation() + .wrapping_add(1); status.active_plugins = old_status.active_plugins; status.skipped_regions = old_status.skipped_regions; @@ -2490,6 +2720,9 @@ mod test { status.successful_updates = old_status.successful_updates; status.failed_updates = old_status.failed_updates; status.invocation_results = old_status.invocation_results; + status + .invocation_results + .set_revert_generation(revert_generation); status.component_revision_for_replay = old_status.component_revision_for_replay; status.last_manual_update_snapshot_index = old_status.last_manual_update_snapshot_index; @@ -2537,6 +2770,7 @@ mod test { status .pending_invocations .retain(|ti| ti.idempotency_key() != Some(&idempotency_key)); + status.cancelled_idempotency_key = Some(idempotency_key); status }) } @@ -2550,6 +2784,7 @@ mod test { None, ), move |mut status| { + status.cancelled_idempotency_key = None; status .invocation_results .insert(idempotency_key, status.oplog_idx); @@ -2854,10 +3089,12 @@ mod test { impl HasConfig for TestCase { fn config(&self) -> Arc { - Arc::new(GolemConfig { + let mut config = GolemConfig { retry: RetryConfig::default(), ..Default::default() - }) + }; + config.invocation_results.physical_index_catch_up_chunk_size = 2; + Arc::new(config) } } @@ -2973,6 +3210,7 @@ mod test { &active_plugins, &deleted_regions, &entries, + true, ); assert_eq!(result.len(), 1); @@ -3027,6 +3265,7 @@ mod test { &active_plugins, &deleted_regions, &entries, + true, ); assert_eq!(result.len(), 2); @@ -3067,6 +3306,7 @@ mod test { &active_plugins, &deleted_regions, &entries, + true, ); let state = result.get(&grant_id).unwrap(); @@ -3104,6 +3344,7 @@ mod test { &active_plugins, &deleted_regions, &entries, + true, ); assert!( @@ -3139,6 +3380,7 @@ mod test { &active_plugins, &deleted_regions, &entries, + true, ); assert!( @@ -3204,6 +3446,7 @@ mod test { &active_plugins, &deleted_regions, &entries, + true, ); assert!( @@ -3240,6 +3483,7 @@ mod test { &active_plugins, &deleted_regions, &entries, + true, ); assert_eq!(result.len(), 1); @@ -3292,6 +3536,7 @@ mod test { &active_plugins, &deleted_regions, &entries, + true, ); let state = result.get(&grant_id).unwrap(); @@ -3304,6 +3549,101 @@ mod test { assert_eq!(state.sending_up_to, OplogIndex::from_u64(8)); } + #[test] + fn oplog_processor_checkpoint_fold_is_chunk_composable() { + fn fold( + entries: &BTreeMap, + chunk_size: usize, + ) -> AgentStatusRecord { + let mut status = AgentStatusRecord::default(); + let last_index = *entries.keys().next_back().unwrap(); + let all_entries: Vec<_> = entries.iter().collect(); + for chunk in all_entries.chunks(chunk_size) { + let chunk: BTreeMap<_, _> = chunk + .iter() + .map(|(index, entry)| (**index, (**entry).clone())) + .collect(); + let finalize = chunk.keys().next_back() == Some(&last_index); + status = super::update_status_with_precomputed_regions( + AgentMode::Durable, + status, + chunk, + &RetryConfig::default(), + DeletedRegions::new(), + DeletedRegions::new(), + finalize, + ); + } + status + } + + let grant_id = EnvironmentPluginGrantId::new(); + let target = AgentId { + component_id: ComponentId::new(), + agent_id: "checkpoint-target".to_string(), + }; + let test_case = TestCase::builder(0).build(); + let mut create = test_case.entries[0].oplog_entry.clone(); + let OplogEntry::Create { + initial_active_plugins, + .. + } = &mut create + else { + unreachable!() + }; + initial_active_plugins.insert(grant_id); + + let entries = BTreeMap::from([ + (OplogIndex::INITIAL, create), + ( + OplogIndex::from_u64(2), + OplogEntry::OplogProcessorCheckpoint { + timestamp: Timestamp::now_utc(), + plugin_grant_id: grant_id, + target_agent_id: target.clone(), + confirmed_up_to: OplogIndex::INITIAL, + sending_up_to: OplogIndex::from_u64(4), + last_batch_start: OplogIndex::INITIAL, + }, + ), + ( + OplogIndex::from_u64(3), + OplogEntry::DeactivatePlugin { + timestamp: Timestamp::now_utc(), + plugin_grant_id: grant_id, + }, + ), + ( + OplogIndex::from_u64(4), + OplogEntry::OplogProcessorCheckpoint { + timestamp: Timestamp::now_utc(), + plugin_grant_id: grant_id, + target_agent_id: target.clone(), + confirmed_up_to: OplogIndex::from_u64(4), + sending_up_to: OplogIndex::from_u64(4), + last_batch_start: OplogIndex::INITIAL, + }, + ), + ( + OplogIndex::from_u64(5), + OplogEntry::ActivatePlugin { + timestamp: Timestamp::now_utc(), + plugin_grant_id: grant_id, + }, + ), + ]); + + let unchunked = fold(&entries, entries.len()); + assert_eq!(fold(&entries, 1), unchunked); + assert_eq!(fold(&entries, 2), unchunked); + let checkpoint = unchunked + .oplog_processor_checkpoints + .get(&grant_id) + .unwrap(); + assert_eq!(checkpoint.target_agent_id, Some(target)); + assert_eq!(checkpoint.confirmed_up_to, OplogIndex::from_u64(4)); + } + #[test] fn deactivate_then_reactivate_seeds_new_checkpoint() { let grant_id = EnvironmentPluginGrantId::new(); @@ -3339,6 +3679,7 @@ mod test { &active_plugins, &deleted_regions, &entries, + true, ); let state = result.get(&grant_id).unwrap(); diff --git a/golem-worker-executor/tests/api.rs b/golem-worker-executor/tests/api.rs index fbf03d200f..a85b581df1 100644 --- a/golem-worker-executor/tests/api.rs +++ b/golem-worker-executor/tests/api.rs @@ -2626,62 +2626,86 @@ async fn invoking_with_same_idempotency_key_is_idempotent_after_restart( #[tagged_as("agent_counters")] agent_counters: &PrecompiledComponent, ) -> anyhow::Result<()> { let context = TestContext::new(last_unique_id); - let executor = start(deps, &context).await?; + let overrides = TestExecutorOverrides { + configure: Some(Arc::new(|config| { + config.invocation_results.recent_capacity = 2; + config.invocation_results.bloom_bits = 128; + config.invocation_results.bloom_hashes = 2; + config.invocation_results.physical_index_catch_up_chunk_size = 2; + })), + ..TestExecutorOverrides::default() + }; + let executor = start_with_overrides(deps, &context, overrides.clone()).await?; let component = executor .component_dep(&context.default_environment_id, agent_counters) .store() .await?; - let repo_id = agent_id!("Repository", "test-repo-3"); - let worker_id = executor.start_agent(&component.id, repo_id.clone()).await?; - - let idempotency_key = IdempotencyKey::fresh(); - executor - .invoke_and_await_agent_with_key( - &component, - &repo_id, - &idempotency_key, - "add", - data_value!("G1000", "Golem T-Shirt M"), - ) + let counter_id = agent_id!("Counter", "idempotency-index-after-restart"); + let worker_id = executor + .start_agent(&component.id, counter_id.clone()) .await?; + let mut oldest_key = None; + for expected in 1..=32 { + let idempotency_key = IdempotencyKey::fresh(); + let result = executor + .invoke_and_await_agent_with_key( + &component, + &counter_id, + &idempotency_key, + "increment", + data_value!(), + ) + .await? + .into_typed::()?; + assert_eq!(result, expected); + oldest_key.get_or_insert(idempotency_key); + } + drop(executor); - let executor = start(deps, &context).await?; + let executor = start_with_overrides(deps, &context, overrides).await?; + assert_eq!( + executor + .start_agent(&component.id, counter_id.clone()) + .await?, + worker_id + ); - executor + let oldest_key = oldest_key.unwrap(); + let read_exact_before = executor.oplog_service_call_count(&worker_id, "read_exact"); + let duplicate_result = executor .invoke_and_await_agent_with_key( &component, - &repo_id, - &idempotency_key, - "add", - data_value!("G1000", "Golem T-Shirt M"), + &counter_id, + &oldest_key, + "increment", + data_value!(), ) - .await?; + .await? + .into_typed::()?; + assert_eq!(duplicate_result, 1); + assert_eq!( + executor.oplog_service_call_count(&worker_id, "read_exact"), + read_exact_before, + "an evicted idempotency result must resolve from the hot physical index without scanning the oplog" + ); - let contents = executor - .invoke_and_await_agent(&component, &repo_id, "list", data_value!()) - .await?; + let next_key = IdempotencyKey::fresh(); + let next_result = executor + .invoke_and_await_agent_with_key( + &component, + &counter_id, + &next_key, + "increment", + data_value!(), + ) + .await? + .into_typed::()?; + assert_eq!(next_result, 33); executor.check_oplog_is_queryable(&worker_id).await?; - - let contents_value = contents - .into_return_value() - .expect("Expected a single return value"); - - assert_eq!( - contents_value, - SchemaValue::List { - elements: vec![SchemaValue::Record { - fields: vec![ - SchemaValue::String("G1000".to_string()), - SchemaValue::String("Golem T-Shirt M".to_string()), - SchemaValue::U64(1), - ], - }], - } - ); Ok(()) } diff --git a/golem-worker-executor/tests/key_value_storage.rs b/golem-worker-executor/tests/key_value_storage.rs index c20f47886c..f5ac9b73dc 100644 --- a/golem-worker-executor/tests/key_value_storage.rs +++ b/golem-worker-executor/tests/key_value_storage.rs @@ -401,10 +401,26 @@ fn ns3() -> Namespaces { } } +#[test_dep(scope = PerWorker, tagged_as = "ns4")] +fn ns4() -> Namespaces { + Namespaces { + ns: KeyValueStorageNamespace::AgentInvocationResultIndex { + agent_id: AgentId { + component_id: ComponentId::new(), + agent_id: "test".to_string(), + }, + }, + ns2: KeyValueStorageNamespace::UserDefined { + environment_id: EnvironmentId(uuid!("296aa41a-ff44-4882-8f34-08b7fe431aa4")), + bucket: "test-bucket-3".to_string(), + }, + } +} + inherit_test_dep!(WorkerExecutorTestDependencies); define_matrix_dimension!(kvs: Arc -> "in_memory", "redis", "sqlite", "multi_sqlite", "postgres", "namespace_routed"); -define_matrix_dimension!(nss: Namespaces -> "ns1", "ns2", "ns3"); +define_matrix_dimension!(nss: Namespaces -> "ns1", "ns2", "ns3", "ns4"); #[test] #[tracing::instrument] @@ -565,6 +581,57 @@ async fn get_all_returns_namespace_snapshot( ); } +#[test] +#[tracing::instrument] +async fn agent_invocation_result_index_is_separate_from_agent_status( + _deps: &WorkerExecutorTestDependencies, + #[dimension(kvs)] kvs: &Arc, +) { + let kvs = kvs.get_key_value_storage().await; + let agent_id = AgentId { + component_id: ComponentId::new(), + agent_id: "test".to_string(), + }; + let status_ns = KeyValueStorageNamespace::AgentStatus { + agent_id: agent_id.clone(), + }; + let result_index_ns = KeyValueStorageNamespace::AgentInvocationResultIndex { agent_id }; + + kvs.set( + "test", + "api", + "entity", + status_ns.clone(), + "same-key", + b"status-value", + ) + .await + .unwrap(); + kvs.set( + "test", + "api", + "entity", + result_index_ns.clone(), + "same-key", + b"result-index-value", + ) + .await + .unwrap(); + + assert_eq!( + kvs.get("test", "api", "entity", status_ns, "same-key") + .await + .unwrap(), + Some(bytes::Bytes::from_static(b"status-value")) + ); + assert_eq!( + kvs.get("test", "api", "entity", result_index_ns, "same-key",) + .await + .unwrap(), + Some(bytes::Bytes::from_static(b"result-index-value")) + ); +} + #[test] #[tracing::instrument] async fn set_if_not_exists( diff --git a/golem-worker-executor/tests/namespace_routed_key_value_storage.rs b/golem-worker-executor/tests/namespace_routed_key_value_storage.rs index ef7bf4101d..4f14d46bfd 100644 --- a/golem-worker-executor/tests/namespace_routed_key_value_storage.rs +++ b/golem-worker-executor/tests/namespace_routed_key_value_storage.rs @@ -100,33 +100,45 @@ async fn build_namespace_routed_kvs( } #[test] -async fn routes_worker_namespace_to_redis(deps: &WorkerExecutorTestDependencies) { +async fn routes_agent_namespaces_to_redis(deps: &WorkerExecutorTestDependencies) { let (kvs, redis, postgres, _postgres_container) = build_namespace_routed_kvs(deps).await; - let ns = KeyValueStorageNamespace::Worker { - agent_id: AgentId { - component_id: ComponentId::new(), - agent_id: "route-test-agent".to_string(), - }, + let agent_id = AgentId { + component_id: ComponentId::new(), + agent_id: "route-test-agent".to_string(), }; - let key = "worker-route-key"; - let value = b"worker-route-value"; - - kvs.set("test", "api", "entity", ns.clone(), key, value) - .await - .unwrap(); - - let redis_read = redis - .get("test", "api", "entity", ns.clone(), key) - .await - .unwrap(); - let postgres_read = postgres - .get("test", "api", "entity", ns, key) - .await - .unwrap(); - - assert_eq!(redis_read, Some(value.as_slice().into())); - assert_eq!(postgres_read, None); + let cases = [ + ( + KeyValueStorageNamespace::Worker { + agent_id: agent_id.clone(), + }, + "worker-route-key", + b"worker-route-value".as_slice(), + ), + ( + KeyValueStorageNamespace::AgentInvocationResultIndex { agent_id }, + "result-index-route-key", + b"result-index-route-value".as_slice(), + ), + ]; + + for (namespace, key, value) in cases { + kvs.set("test", "api", "entity", namespace.clone(), key, value) + .await + .unwrap(); + + let redis_read = redis + .get("test", "api", "entity", namespace.clone(), key) + .await + .unwrap(); + let postgres_read = postgres + .get("test", "api", "entity", namespace, key) + .await + .unwrap(); + + assert_eq!(redis_read, Some(value.into())); + assert_eq!(postgres_read, None); + } } #[test] diff --git a/golem-worker-executor/tests/rpc.rs b/golem-worker-executor/tests/rpc.rs index 54b58b5ed5..8b7b858bb9 100644 --- a/golem-worker-executor/tests/rpc.rs +++ b/golem-worker-executor/tests/rpc.rs @@ -367,13 +367,25 @@ async fn durable_streaming_output_recovers_after_executor_restart( _tracing: &Tracing, ) -> anyhow::Result<()> { let context = TestContext::new(last_unique_id); - let executor = start(deps, &context).await?; + let overrides = TestExecutorOverrides { + configure: Some(Arc::new(|config| { + config.invocation_results.recent_capacity = 0; + config.invocation_results.bloom_bits = 1; + config.invocation_results.bloom_hashes = 1; + })), + ..Default::default() + }; + let executor = start_with_overrides(deps, &context, overrides.clone()).await?; let component = executor .component_dep(&context.default_environment_id, agent_rpc_rust) .store() .await?; - let agent_id = agent_id!("StreamingRpcTarget", "output-restart"); - let worker_agent_id = executor.start_agent(&component.id, agent_id).await?; + let worker_agent_id = executor + .start_agent( + &component.id, + agent_id!("StreamingRpcTarget", "output-restart"), + ) + .await?; let metadata = executor.get_worker_metadata(&worker_agent_id).await?; let (_, input) = data_value!().into_parts(); let start_request = InvocationRequest { @@ -445,7 +457,7 @@ async fn durable_streaming_output_recovers_after_executor_restart( drop(responses); drop(executor); - let executor = start(deps, &context).await?; + let executor = start_with_overrides(deps, &context, overrides).await?; let Some(invocation_request::Request::Start(start)) = start_request.request.as_ref() else { anyhow::bail!("durable output restart request is not Start"); }; diff --git a/golem-worker-executor/tests/tool_streaming.rs b/golem-worker-executor/tests/tool_streaming.rs index 191632bffa..122a6fed58 100644 --- a/golem-worker-executor/tests/tool_streaming.rs +++ b/golem-worker-executor/tests/tool_streaming.rs @@ -1650,11 +1650,29 @@ async fn guest_trap_fences_a_blocked_sibling_and_drains_the_owner_group( "the original guest-trap provenance must survive owner-group fencing: {error:?}" ); - if let Some(active) = executor.active_entity_metadata(&owned_agent_id).await { - assert!(active.tool_operations.operations.is_empty()); - assert!(active.lane.holder.is_none()); - assert_eq!(active.lane.active_invocation_count, 0); - assert!(active.slots.iter().all(|slot| slot.invocations.is_empty())); + let cleanup = tokio::time::timeout(std::time::Duration::from_secs(30), async { + loop { + if executor + .active_entity_metadata(&owned_agent_id) + .await + .is_none_or(|active| { + active.tool_operations.operations.is_empty() + && active.lane.holder.is_none() + && active.lane.active_invocation_count == 0 + && active.slots.iter().all(|slot| slot.invocations.is_empty()) + }) + { + break; + } + tokio::task::yield_now().await; + } + }) + .await; + if cleanup.is_err() { + let active = executor.active_entity_metadata(&owned_agent_id).await; + anyhow::bail!( + "timed out waiting for guest-trap owner cleanup; active metadata: {active:#?}" + ); } let oplog = executor.get_oplog(&worker_id, OplogIndex::INITIAL).await?; let invocation_start = oplog diff --git a/integration-tests/src/benchmarks/all.rs b/integration-tests/src/benchmarks/all.rs index d8c69e08a5..ae44b66f81 100644 --- a/integration-tests/src/benchmarks/all.rs +++ b/integration-tests/src/benchmarks/all.rs @@ -109,6 +109,14 @@ async fn main() { >(mode, verbosity, item, primary_only, otlp)) }), ); + benchmarks_by_name.insert( + "idempotency-key-lookup", + Box::new(|mode, verbosity, item, primary_only, otlp| { + Box::pin(run_benchmark::< + integration_tests::benchmarks::idempotency_key::IdempotencyKeyLookup, + >(mode, verbosity, item, primary_only, otlp)) + }), + ); benchmarks_by_name.insert( "throughput-echo", Box::new(|mode, verbosity, item, primary_only, otlp| { diff --git a/integration-tests/src/benchmarks/idempotency_key.rs b/integration-tests/src/benchmarks/idempotency_key.rs new file mode 100644 index 0000000000..9cedb0ab6c --- /dev/null +++ b/integration-tests/src/benchmarks/idempotency_key.rs @@ -0,0 +1,171 @@ +// Copyright 2024-2026 Golem Cloud +// +// Licensed under the Golem Source License v1.1 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://license.golem.cloud/LICENSE +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use crate::benchmarks::{cleanup_user_state, delete_workers}; +use async_trait::async_trait; +use golem_common::model::agent::ParsedAgentId; +use golem_common::model::component::ComponentDto; +use golem_common::model::environment::EnvironmentId; +use golem_common::model::{AgentId, IdempotencyKey}; +use golem_common::{agent_id, data_value}; +use golem_test_framework::benchmark::{Benchmark, BenchmarkRecorder, RunConfig}; +use golem_test_framework::config::benchmark::TestMode; +use golem_test_framework::config::dsl_impl::TestUserContext; +use golem_test_framework::config::{BenchmarkTestDependencies, TestDependencies}; +use golem_test_framework::dsl::{TestDsl, TestDslExtended}; +use indoc::indoc; +use std::time::Instant; +use tracing::Level; + +pub struct IdempotencyKeyLookup { + config: RunConfig, +} + +pub struct BenchmarkContext { + deps: BenchmarkTestDependencies, +} + +pub struct IterationContext { + user: TestUserContext, + component: ComponentDto, + agent_id: ParsedAgentId, + env_id: EnvironmentId, +} + +#[async_trait] +impl Benchmark for IdempotencyKeyLookup { + type BenchmarkContext = BenchmarkContext; + type IterationContext = IterationContext; + + fn name() -> &'static str { + "idempotency-key-lookup" + } + + fn description() -> &'static str { + indoc! { + "Invokes one long-lived durable agent with `size` unique idempotency keys, then performs + `length` duplicate lookups each for the newest and oldest keys. Records the latency and + throughput of unique invocations, recent duplicates, and old duplicates separately." + } + } + + async fn create_benchmark_context( + mode: &TestMode, + verbosity: Level, + cluster_size: usize, + disable_compilation_cache: bool, + otlp: bool, + ) -> Self::BenchmarkContext { + BenchmarkContext { + deps: BenchmarkTestDependencies::new( + mode, + verbosity, + cluster_size, + disable_compilation_cache, + otlp, + ) + .await, + } + } + + async fn cleanup(context: Self::BenchmarkContext) { + context.deps.kill_all().await; + } + + async fn create(_mode: &TestMode, config: RunConfig) -> Self { + Self { config } + } + + async fn setup_iteration(&self, context: &Self::BenchmarkContext) -> Self::IterationContext { + let user = context.deps.user().await.unwrap(); + let (_, env) = user.app_and_env().await.unwrap(); + let component = user + .component(&env.id, "benchmark_agent_rust_release") + .name("benchmark:agent-rust") + .store() + .await + .unwrap(); + + IterationContext { + user, + component, + agent_id: agent_id!("RustBenchmarkAgent", "idempotency-key-lookup"), + env_id: env.id, + } + } + + async fn warmup( + &self, + _benchmark_context: &Self::BenchmarkContext, + context: &Self::IterationContext, + ) { + invoke(context, &IdempotencyKey::fresh()).await; + } + + async fn run( + &self, + _benchmark_context: &Self::BenchmarkContext, + context: &Self::IterationContext, + recorder: BenchmarkRecorder, + ) { + assert!(self.config.size > 0, "size must be at least one"); + + let mut keys = Vec::with_capacity(self.config.size); + for _ in 0..self.config.size { + let key = IdempotencyKey::fresh(); + let started = Instant::now(); + invoke(context, &key).await; + recorder.duration(&"unique-invocation".into(), started.elapsed()); + keys.push(key); + } + + let old = keys.first().unwrap(); + let recent = keys.last().unwrap(); + for _ in 0..self.config.length { + let started = Instant::now(); + invoke(context, recent).await; + recorder.duration(&"recent-duplicate".into(), started.elapsed()); + } + for _ in 0..self.config.length { + let started = Instant::now(); + invoke(context, old).await; + recorder.duration(&"old-duplicate".into(), started.elapsed()); + } + } + + async fn cleanup_iteration( + &self, + _benchmark_context: &Self::BenchmarkContext, + context: Self::IterationContext, + ) { + if let Ok(agent_id) = AgentId::from_agent_id(context.component.id, &context.agent_id) { + delete_workers(&context.user, &[agent_id]).await; + } + cleanup_user_state(&context.user, &context.env_id).await; + } +} + +async fn invoke(context: &IterationContext, key: &IdempotencyKey) { + context + .user + .invoke_and_await_agent_with_key( + &context.component, + &context.agent_id, + key, + "echo", + data_value!("benchmark"), + ) + .await + .expect("agent invocation failed"); +} diff --git a/integration-tests/src/benchmarks/mod.rs b/integration-tests/src/benchmarks/mod.rs index e43792c7dc..2783bcd922 100644 --- a/integration-tests/src/benchmarks/mod.rs +++ b/integration-tests/src/benchmarks/mod.rs @@ -32,6 +32,7 @@ use tracing_opentelemetry::OpenTelemetrySpanExt; pub mod cleanup; pub mod cold_start_unknown; pub mod durability_overhead; +pub mod idempotency_key; pub mod latency; pub mod sleep; pub mod throughput; diff --git a/test-components/benchmarks/Cargo.lock b/test-components/benchmarks/Cargo.lock index 4ac6d818bd..059e91efe5 100644 --- a/test-components/benchmarks/Cargo.lock +++ b/test-components/benchmarks/Cargo.lock @@ -26,6 +26,12 @@ version = "1.0.104" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "330a5ed07fa54e4702c9d6c4174f74427fc0ef6e214bbd677ae50a5099946470" +[[package]] +name = "arrayvec" +version = "0.7.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d3fb67a6e08acf24fdeccbac2cb6ac4305825bd1f117462e0e6f2f193345ad56" + [[package]] name = "async-trait" version = "0.1.91" @@ -89,6 +95,20 @@ version = "2.13.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b588b76d00fde79687d7646a9b5bdf3cc0f655e0bbd080335a95d7e96f3587da" +[[package]] +name = "blake3" +version = "1.8.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6d9e454fc11f76977dc803893aff6304ed33d6a26efae8696573bea74baa27ae" +dependencies = [ + "arrayvec", + "cc", + "cfg-if", + "constant_time_eq", + "cpufeatures", + "rayon-core", +] + [[package]] name = "bumpalo" version = "3.20.3" @@ -148,12 +168,52 @@ dependencies = [ "golem-rust", ] +[[package]] +name = "constant_time_eq" +version = "0.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3d52eff69cd5e647efe296129160853a42795992097e8af39800e1060caeea9b" + [[package]] name = "core-foundation-sys" version = "0.8.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "773648b94d0e5d620f64f280777445740e61fe701025087ec8b57f45c791888b" +[[package]] +name = "cpufeatures" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5ca28b0ae3115b884660db4118d803791fd6756b6e88f39c0f3f7859060d7566" +dependencies = [ + "libc", +] + +[[package]] +name = "crossbeam-deque" +version = "0.8.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5181e0de7b61eb03a81e347d6dd8797bae9da5146707b51077e2d71a54ec0ceb" +dependencies = [ + "crossbeam-epoch", + "crossbeam-utils", +] + +[[package]] +name = "crossbeam-epoch" +version = "0.9.20" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2d6914041f254d6e9176c01941b21115dcfb7089e55135a35411081bd106ef3f" +dependencies = [ + "crossbeam-utils", +] + +[[package]] +name = "crossbeam-utils" +version = "0.8.22" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "61803da095bee82a81bb1a452ecc25d3b2f1416d1897eb86430c6159ef717c17" + [[package]] name = "ctor" version = "0.4.3" @@ -417,6 +477,7 @@ dependencies = [ "base64", "bigdecimal", "bit-vec", + "blake3", "chrono", "combine", "golem-schema-derive", @@ -824,6 +885,16 @@ dependencies = [ "num-traits", ] +[[package]] +name = "rayon-core" +version = "1.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "22e18b0f0062d30d4230b2e85ff77fdfe4326feb054b9783a3460d8435c8ab91" +dependencies = [ + "crossbeam-deque", + "crossbeam-utils", +] + [[package]] name = "regex" version = "1.13.1" diff --git a/test-components/benchmarks/package-lock.json b/test-components/benchmarks/package-lock.json index 74e91fe2f2..9c9b6e7959 100644 --- a/test-components/benchmarks/package-lock.json +++ b/test-components/benchmarks/package-lock.json @@ -24,6 +24,9 @@ "name": "@golemcloud/golem-ts-sdk", "version": "0.0.0", "license": "SEE LICENSE IN LICENSE", + "dependencies": { + "@noble/hashes": "^1.8.0" + }, "devDependencies": { "@eslint/js": "^9.33.0", "@rollup/plugin-commonjs": "^28.0.6",