diff --git a/crates/trusted-server-adapter-fastly/src/app.rs b/crates/trusted-server-adapter-fastly/src/app.rs index 123288da5..530dea9d8 100644 --- a/crates/trusted-server-adapter-fastly/src/app.rs +++ b/crates/trusted-server-adapter-fastly/src/app.rs @@ -601,19 +601,22 @@ async fn execute_named( handler, NamedRouteHandler::AdminEcLookup | NamedRouteHandler::AdminEidsLookup ) { - let response = PartnerRegistry::from_config(&state.settings.ec.partners) - .and_then(|registry| match handler { + let response = async { + let registry = PartnerRegistry::from_config(&state.settings.ec.partners)?; + match handler { NamedRouteHandler::AdminEcLookup => { // Deliberately do not use an EC request-state graph: that // copy is bot-gated, while operators use curl for this // authenticated diagnostic. let kv = crate::maybe_identity_graph(&state.settings); - handle_admin_ec_lookup(kv.as_ref(), ®istry, &req) + handle_admin_ec_lookup(kv.as_ref(), ®istry, &req).await } NamedRouteHandler::AdminEidsLookup => handle_admin_eids_lookup(®istry, &req), _ => unreachable!("admin diagnostics should use early dispatch"), - }) - .unwrap_or_else(|error| http_error(&error)); + } + } + .await + .unwrap_or_else(|error| http_error(&error)); return Ok(response); } @@ -692,6 +695,7 @@ async fn run_named_route( &req, &ec.ec_context, ) + .await } } NamedRouteHandler::SetTester => handle_set_tester(&state.settings), @@ -881,6 +885,7 @@ async fn dispatch_fallback( && let Err(err) = ec .ec_context .generate_if_needed(&state.settings, ec.kv_graph.as_ref()) + .await { log::warn!("EC generation failed for publisher proxy: {err:?}"); } diff --git a/crates/trusted-server-adapter-fastly/src/ec_kv.rs b/crates/trusted-server-adapter-fastly/src/ec_kv.rs index c427ede96..4914d323a 100644 --- a/crates/trusted-server-adapter-fastly/src/ec_kv.rs +++ b/crates/trusted-server-adapter-fastly/src/ec_kv.rs @@ -2,7 +2,9 @@ //! //! Maps the platform-neutral identity-graph store operations onto the //! Fastly KV Store API, including generation markers for compare-and-swap -//! writes (`if_generation_match`). +//! writes (`if_generation_match`). The SDK exposes only blocking completion +//! hostcalls for KV operations, so these local async methods retain that Fastly +//! behavior. Native backends must yield during network I/O instead. use error_stack::{Report, ResultExt}; use fastly::kv_store::{InsertMode, KVStore, KVStoreError, ListPage}; @@ -86,12 +88,13 @@ impl FastlyEcKvStore { } } +#[async_trait::async_trait(?Send)] impl EcKvStore for FastlyEcKvStore { fn store_name(&self) -> &str { &self.store_name } - fn lookup(&self, key: &str) -> Result, Report> { + async fn lookup(&self, key: &str) -> Result, Report> { let store = self.open_store()?; let mut response = match store.lookup(key) { Ok(resp) => resp, @@ -117,7 +120,7 @@ impl EcKvStore for FastlyEcKvStore { })) } - fn key_exists(&self, key: &str) -> Result> { + async fn key_exists(&self, key: &str) -> Result> { let store = self.open_store()?; contains_exact_key(&self.store_name, key, |cursor| { // build_list defaults to ListMode::Strong, reading primary state. @@ -129,7 +132,7 @@ impl EcKvStore for FastlyEcKvStore { }) } - fn insert( + async fn insert( &self, key: &str, write: EcKvWrite<'_>, @@ -160,7 +163,7 @@ impl EcKvStore for FastlyEcKvStore { } } - fn list_keys_with_prefix( + async fn list_keys_with_prefix( &self, prefix: &str, limit: u32, @@ -178,7 +181,7 @@ impl EcKvStore for FastlyEcKvStore { }) } - fn delete(&self, key: &str) -> Result<(), Report> { + async fn delete(&self, key: &str) -> Result<(), Report> { let store = self.open_store()?; store .delete(key) @@ -205,7 +208,7 @@ mod tests { FastlyEcKvStore::new(TEST_STORE) } - fn write(key: &str, body: &str, mode: EcKvWriteMode) -> EcKvWriteOutcome { + async fn write(key: &str, body: &str, mode: EcKvWriteMode) -> EcKvWriteOutcome { store() .insert( key, @@ -216,6 +219,7 @@ mod tests { mode, }, ) + .await .expect("should reach the store") } @@ -334,166 +338,201 @@ mod tests { #[test] fn strong_existence_checks_real_store_keys_exactly() { - let backend = store(); - let key = "existence-test-key"; - let longer = "existence-test-key-longer"; - write(longer, "body", EcKvWriteMode::Overwrite); - assert!( - !backend.key_exists(key).expect("should list the prefix"), - "should not confuse a prefix with an exact key" - ); - write(key, "body", EcKvWriteMode::Overwrite); - assert!( - backend.key_exists(key).expect("should list the issued key"), - "should see a completed issuance" - ); - backend.delete(key).expect("should delete exact key"); - backend.delete(longer).expect("should delete longer key"); - assert!( - !backend.key_exists(key).expect("should list after deletion"), - "should report absence after deletion" - ); + futures::executor::block_on(async { + let backend = store(); + let key = "existence-test-key"; + let longer = "existence-test-key-longer"; + write(longer, "body", EcKvWriteMode::Overwrite).await; + assert!( + !backend + .key_exists(key) + .await + .expect("should list the prefix"), + "should not confuse a prefix with an exact key" + ); + write(key, "body", EcKvWriteMode::Overwrite).await; + assert!( + backend + .key_exists(key) + .await + .expect("should list the issued key"), + "should see a completed issuance" + ); + backend.delete(key).await.expect("should delete exact key"); + backend + .delete(longer) + .await + .expect("should delete longer key"); + assert!( + !backend + .key_exists(key) + .await + .expect("should list after deletion"), + "should report absence after deletion" + ); + }); } #[test] fn opening_a_store_this_service_does_not_have_is_an_error() { - let error = FastlyEcKvStore::new("no_such_store") - .lookup("any-key") - .expect_err("should not resolve against a store that is not linked"); - - assert!( - matches!( - error.current_context(), - TrustedServerError::KvStore { store_name, .. } if store_name == "no_such_store" - ), - "should name the store it could not open: {error:?}" - ); + futures::executor::block_on(async { + let error = FastlyEcKvStore::new("no_such_store") + .lookup("any-key") + .await + .expect_err("should not resolve against a store that is not linked"); + + assert!( + matches!( + error.current_context(), + TrustedServerError::KvStore { store_name, .. } if store_name == "no_such_store" + ), + "should name the store it could not open: {error:?}" + ); + }); } #[test] fn a_missing_key_is_absent_rather_than_an_error() { - let absent = format!("{}.ABC123", "1".repeat(64)); - - assert!( - store() - .lookup(&absent) - .expect("should reach the store") - .is_none(), - "a key the store does not hold is absent, not a failure" - ); + futures::executor::block_on(async { + let absent = format!("{}.ABC123", "1".repeat(64)); + + assert!( + store() + .lookup(&absent) + .await + .expect("should reach the store") + .is_none(), + "a key the store does not hold is absent, not a failure" + ); + }); } #[test] fn an_entry_round_trips_through_insert_lookup_and_delete() { - let key = format!("{}.ABC123", "2".repeat(64)); - let backend = store(); + futures::executor::block_on(async { + let key = format!("{}.ABC123", "2".repeat(64)); + let backend = store(); - assert_eq!( - write(&key, "entry-body-1", EcKvWriteMode::Overwrite), - EcKvWriteOutcome::Written, - "should write the entry" - ); + assert_eq!( + write(&key, "entry-body-1", EcKvWriteMode::Overwrite).await, + EcKvWriteOutcome::Written, + "should write the entry" + ); - let found = backend - .lookup(&key) - .expect("should reach the store") - .expect("should hold the entry just written"); - assert_eq!(found.body, b"entry-body-1", "should read back the body"); - assert_eq!( - found.metadata.as_deref(), - Some(METADATA.as_bytes()), - "should read back the metadata" - ); - - backend.delete(&key).expect("should delete the entry"); - assert!( - backend + let found = backend .lookup(&key) + .await .expect("should reach the store") - .is_none(), - "a deleted key is absent" - ); + .expect("should hold the entry just written"); + assert_eq!(found.body, b"entry-body-1", "should read back the body"); + assert_eq!( + found.metadata.as_deref(), + Some(METADATA.as_bytes()), + "should read back the metadata" + ); + + backend.delete(&key).await.expect("should delete the entry"); + assert!( + backend + .lookup(&key) + .await + .expect("should reach the store") + .is_none(), + "a deleted key is absent" + ); + }); } #[test] fn add_mode_refuses_a_key_that_already_exists() { - let key = format!("{}.ABC123", "3".repeat(64)); - let backend = store(); - - assert_eq!( - write(&key, "entry-body-1", EcKvWriteMode::Add), - EcKvWriteOutcome::Written, - "should create a key nothing holds" - ); - assert_eq!( - write(&key, "entry-body-2", EcKvWriteMode::Add), - EcKvWriteOutcome::PreconditionFailed, - "a precondition failure is control flow, not an error" - ); - - backend.delete(&key).expect("should delete the entry"); + futures::executor::block_on(async { + let key = format!("{}.ABC123", "3".repeat(64)); + let backend = store(); + + assert_eq!( + write(&key, "entry-body-1", EcKvWriteMode::Add).await, + EcKvWriteOutcome::Written, + "should create a key nothing holds" + ); + assert_eq!( + write(&key, "entry-body-2", EcKvWriteMode::Add).await, + EcKvWriteOutcome::PreconditionFailed, + "a precondition failure is control flow, not an error" + ); + + backend.delete(&key).await.expect("should delete the entry"); + }); } #[test] fn a_generation_mismatch_is_reported_as_a_precondition_failure() { - let key = format!("{}.ABC123", "4".repeat(64)); - let backend = store(); - - write(&key, "entry-body-1", EcKvWriteMode::Overwrite); - let generation = backend - .lookup(&key) - .expect("should reach the store") - .expect("should hold the entry just written") - .generation; - - assert_eq!( - write( - &key, - "entry-body-2", - EcKvWriteMode::IfGenerationMatch(generation) - ), - EcKvWriteOutcome::Written, - "should write when the generation still matches" - ); - assert_eq!( - write( - &key, - "entry-body-3", - EcKvWriteMode::IfGenerationMatch(generation) - ), - EcKvWriteOutcome::PreconditionFailed, - "the generation moved on with the previous write" - ); + futures::executor::block_on(async { + let key = format!("{}.ABC123", "4".repeat(64)); + let backend = store(); - backend.delete(&key).expect("should delete the entry"); + write(&key, "entry-body-1", EcKvWriteMode::Overwrite).await; + let generation = backend + .lookup(&key) + .await + .expect("should reach the store") + .expect("should hold the entry just written") + .generation; + + assert_eq!( + write( + &key, + "entry-body-2", + EcKvWriteMode::IfGenerationMatch(generation) + ) + .await, + EcKvWriteOutcome::Written, + "should write when the generation still matches" + ); + assert_eq!( + write( + &key, + "entry-body-3", + EcKvWriteMode::IfGenerationMatch(generation) + ) + .await, + EcKvWriteOutcome::PreconditionFailed, + "the generation moved on with the previous write" + ); + + backend.delete(&key).await.expect("should delete the entry"); + }); } #[test] fn counting_a_prefix_counts_only_the_keys_under_it() { - let hash = "5".repeat(64); - let backend = store(); - let keys = [format!("{hash}.AAA111"), format!("{hash}.BBB222")]; - for key in &keys { - write(key, "entry-body-1", EcKvWriteMode::Overwrite); - } - - assert_eq!( - backend - .count_keys_with_prefix(&hash, 100) - .expect("should list the prefix"), - 2, - "should count both keys issued under this hash" - ); - assert_eq!( - backend - .count_keys_with_prefix(&"6".repeat(64), 100) - .expect("should list the prefix"), - 0, - "should count nothing under a hash nothing was issued for" - ); + futures::executor::block_on(async { + let hash = "5".repeat(64); + let backend = store(); + let keys = [format!("{hash}.AAA111"), format!("{hash}.BBB222")]; + for key in &keys { + write(key, "entry-body-1", EcKvWriteMode::Overwrite).await; + } - for key in &keys { - backend.delete(key).expect("should delete the entry"); - } + assert_eq!( + backend + .count_keys_with_prefix(&hash, 100) + .await + .expect("should list the prefix"), + 2, + "should count both keys issued under this hash" + ); + assert_eq!( + backend + .count_keys_with_prefix(&"6".repeat(64), 100) + .await + .expect("should list the prefix"), + 0, + "should count nothing under a hash nothing was issued for" + ); + + for key in &keys { + backend.delete(key).await.expect("should delete the entry"); + } + }); } } diff --git a/crates/trusted-server-adapter-fastly/src/main.rs b/crates/trusted-server-adapter-fastly/src/main.rs index b332b87f8..d0a575c16 100644 --- a/crates/trusted-server-adapter-fastly/src/main.rs +++ b/crates/trusted-server-adapter-fastly/src/main.rs @@ -456,7 +456,11 @@ fn finalize_response( } let state = ec_state.and_then(|mut ec_state| { - match apply_edgezero_ec_finalize(inputs.settings, &mut ec_state, response) { + match futures::executor::block_on(apply_edgezero_ec_finalize( + inputs.settings, + &mut ec_state, + response, + )) { Ok(partner_registry) => Some(FinalizeState { partner_registry, ec_state, @@ -587,7 +591,7 @@ fn take_finalize_sentinel(response: &mut HttpResponse) -> bool { .is_some() } -fn apply_edgezero_ec_finalize( +async fn apply_edgezero_ec_finalize( settings: &Settings, ec_state: &mut EcFinalizeState, response: &mut HttpResponse, @@ -606,7 +610,8 @@ fn apply_edgezero_ec_finalize( ec_state.eids_cookie.as_deref(), ec_state.sharedid_cookie.as_deref(), response, - ); + ) + .await; Ok(partner_registry) } diff --git a/crates/trusted-server-adapter-fastly/src/post_delivery.rs b/crates/trusted-server-adapter-fastly/src/post_delivery.rs index 6082af538..ce54fccbc 100644 --- a/crates/trusted-server-adapter-fastly/src/post_delivery.rs +++ b/crates/trusted-server-adapter-fastly/src/post_delivery.rs @@ -344,7 +344,7 @@ mod tests { } /// Finalize state for a returning browser whose EC row the local store holds. - fn returning_reader_state(settings: &Settings) -> FinalizeState { + async fn returning_reader_state(settings: &Settings) -> FinalizeState { let request = edgezero_core::http::request_builder() .uri("https://test-publisher.example.com/article") .header("cookie", format!("ts-ec={}", seeded_ec_id())) @@ -366,7 +366,7 @@ mod tests { EcContext::read_from_request_with_geo(settings, &request, &test_services(), Some(&geo)) .expect("should read the reader's EC context"); let kv = require_identity_graph(settings).expect("should open the local identity store"); - ec_context.set_kv_snapshot(kv.load_snapshot(&seeded_ec_id())); + ec_context.set_kv_snapshot(kv.load_snapshot(&seeded_ec_id()).await); FinalizeState { partner_registry: PartnerRegistry::from_config(&settings.ec.partners) .expect("should build the partner registry"), @@ -397,7 +397,7 @@ mod tests { ); let settings = pull_sync_settings(); let (queue, _body) = queue_with_one_auction(); - let state = returning_reader_state(&settings); + let state = futures::executor::block_on(returning_reader_state(&settings)); // The origin aborted after the page head was committed. handle_delivery( @@ -435,7 +435,7 @@ mod tests { let client = ScriptedHttpClient::forbidding(); let clock = TestClock::new(); let settings = pull_sync_settings(); - let state = returning_reader_state(&settings); + let state = futures::executor::block_on(returning_reader_state(&settings)); let mut state = Some(state); handle_delivery( diff --git a/crates/trusted-server-core/src/auction/endpoints.rs b/crates/trusted-server-core/src/auction/endpoints.rs index 93dd78d32..b79d38552 100644 --- a/crates/trusted-server-core/src/auction/endpoints.rs +++ b/crates/trusted-server-core/src/auction/endpoints.rs @@ -344,7 +344,7 @@ pub async fn handle_auction( // present registry: without one, `resolve_auction_eids` yields no // server-side EIDs, so the snapshot would be an unused billable KV read. let auction_kv_snapshot = match (kv, ec_id.as_deref(), registry) { - (Some(graph), Some(ec_id), Some(_)) => graph.load_snapshot(ec_id), + (Some(graph), Some(ec_id), Some(_)) => graph.load_snapshot(ec_id).await, _ => EcKvSnapshot::NotRead, }; // Hand the loaded row to the request context so response finalization — @@ -819,7 +819,10 @@ mod tests { let ec_id = format!("{}.ABC123", "a".repeat(64)); let mut live = crate::ec::kv_types::KvEntry::tombstone(1000); live.consent.ok = true; - graph.create(&ec_id, &live).expect("should seed live row"); + graph + .create(&ec_id, &live) + .await + .expect("should seed live row"); lookups.store(0, std::sync::atomic::Ordering::Relaxed); let mut ec_context = make_ec_context(Jurisdiction::NonRegulated, Some(&ec_id)); @@ -872,7 +875,8 @@ mod tests { None, Some("shared-cookie-id"), &mut response, - ); + ) + .await; assert_eq!( lookups.load(std::sync::atomic::Ordering::Relaxed), @@ -881,6 +885,7 @@ mod tests { ); let (stored, _) = graph .get(&ec_id) + .await .expect("should read store") .expect("row should exist"); assert_eq!( @@ -913,7 +918,10 @@ mod tests { let ec_id = format!("{}.eidbdy", "a".repeat(64)); let mut live = crate::ec::kv_types::KvEntry::tombstone(1000); live.consent.ok = true; - graph.create(&ec_id, &live).expect("should seed live row"); + graph + .create(&ec_id, &live) + .await + .expect("should seed live row"); let mut ec_context = make_ec_context(Jurisdiction::NonRegulated, Some(&ec_id)); ec_context.set_eid_sync_source(crate::ec::EidSyncSource::Auction); @@ -967,10 +975,12 @@ mod tests { None, // No `ts-eids` cookie on this request at all. None, &mut response, - ); + ) + .await; let (stored, _) = graph .get(&ec_id) + .await .expect("should read store") .expect("row should exist"); assert_eq!( diff --git a/crates/trusted-server-core/src/ec/admin.rs b/crates/trusted-server-core/src/ec/admin.rs index e21fab702..2bd6f3a7b 100644 --- a/crates/trusted-server-core/src/ec/admin.rs +++ b/crates/trusted-server-core/src/ec/admin.rs @@ -273,7 +273,7 @@ struct SkippedPartnerId { /// /// Returns [`TrustedServerError::KvStore`] when the store open or read /// fails. -pub fn handle_admin_ec_lookup( +pub async fn handle_admin_ec_lookup( kv: Option<&KvIdentityGraph>, registry: &PartnerRegistry, req: &Request, @@ -287,7 +287,7 @@ pub fn handle_admin_ec_lookup( Err(response) => return Ok(*response), }; - let Some(lookup) = kv.lookup_raw(&ec_id)? else { + let Some(lookup) = kv.lookup_raw(&ec_id).await? else { log::info!("Admin EC lookup: no entry for '{}'", log_id(&ec_id)); return Ok(json_error( StatusCode::NOT_FOUND, @@ -742,18 +742,22 @@ mod tests { .expect("should build test request") } - fn kv_with_entry(ec_id: &str, entry: &KvEntry) -> KvIdentityGraph { + async fn kv_with_entry(ec_id: &str, entry: &KvEntry) -> KvIdentityGraph { let kv = KvIdentityGraph::in_memory("test-store"); - kv.create(ec_id, entry).expect("should seed KV entry"); + kv.create(ec_id, entry).await.expect("should seed KV entry"); kv } - fn kv_with_raw_body(ec_id: &str, body: &str) -> KvIdentityGraph { + async fn kv_with_raw_body(ec_id: &str, body: &str) -> KvIdentityGraph { let metadata = serde_json::json!({ "ok": true, "country": "US", "v": 1 }).to_string(); - kv_with_raw_body_and_metadata(ec_id, body, &metadata) + kv_with_raw_body_and_metadata(ec_id, body, &metadata).await } - fn kv_with_raw_body_and_metadata(ec_id: &str, body: &str, metadata: &str) -> KvIdentityGraph { + async fn kv_with_raw_body_and_metadata( + ec_id: &str, + body: &str, + metadata: &str, + ) -> KvIdentityGraph { let store = InMemoryEcKv::new("test-store"); store .insert( @@ -765,6 +769,7 @@ mod tests { mode: EcKvWriteMode::Add, }, ) + .await .expect("should seed raw KV body"); KvIdentityGraph::new(store) } @@ -971,13 +976,14 @@ mod tests { } } - #[test] - fn returns_entry_with_auction_view() { + #[tokio::test(flavor = "current_thread")] + async fn returns_entry_with_auction_view() { let ec_id = test_ec_id(); - let kv = kv_with_entry(&ec_id, &sample_entry()); + let kv = kv_with_entry(&ec_id, &sample_entry()).await; let req = get_request(&format!("/_ts/admin/ec/{ec_id}")); let response = handle_admin_ec_lookup(Some(&kv), &test_registry(), &req) + .await .expect("should handle lookup"); assert_eq!(response.status(), StatusCode::OK); @@ -1037,8 +1043,8 @@ mod tests { ); } - #[test] - fn preserves_raw_entry_and_metadata_shapes() { + #[tokio::test(flavor = "current_thread")] + async fn preserves_raw_entry_and_metadata_shapes() { let ec_id = test_ec_id(); let body = serde_json::json!({ "v": 1, @@ -1070,10 +1076,11 @@ mod tests { "future_metadata": { "source": "edge" } }) .to_string(); - let kv = kv_with_raw_body_and_metadata(&ec_id, &body, &metadata); + let kv = kv_with_raw_body_and_metadata(&ec_id, &body, &metadata).await; let req = get_request(&format!("/_ts/admin/ec/{ec_id}")); let response = handle_admin_ec_lookup(Some(&kv), &test_registry(), &req) + .await .expect("should handle lookup"); let json = response_json(response); @@ -1093,14 +1100,15 @@ mod tests { assert_eq!(json["auction"]["eids"][0]["source"], "bidstream.example"); } - #[test] - fn valid_json_with_invalid_entry_schema_remains_visible() { + #[tokio::test(flavor = "current_thread")] + async fn valid_json_with_invalid_entry_schema_remains_visible() { let ec_id = test_ec_id(); let body = serde_json::json!({ "future": "value" }).to_string(); - let kv = kv_with_raw_body(&ec_id, &body); + let kv = kv_with_raw_body(&ec_id, &body).await; let req = get_request(&format!("/_ts/admin/ec/{ec_id}")); let response = handle_admin_ec_lookup(Some(&kv), &test_registry(), &req) + .await .expect("should handle lookup"); let json = response_json(response); @@ -1115,17 +1123,19 @@ mod tests { assert!(json.get("auction").is_none()); } - #[test] - fn reports_tombstone_entries() { + #[tokio::test(flavor = "current_thread")] + async fn reports_tombstone_entries() { let ec_id = test_ec_id(); // Only an identity the store already holds can be tombstoned, so seed // the live entry the withdrawal replaces. let kv = kv_with_entry( &ec_id, &KvEntry::minimal("bidstream.example", "uid-live", 1_741_824_000), - ); + ) + .await; assert_eq!( kv.write_withdrawal_tombstone(&ec_id, drop) + .await .expect("should write tombstone"), TombstoneOutcome::Written, "should tombstone the seeded identity" @@ -1133,6 +1143,7 @@ mod tests { let req = get_request(&format!("/_ts/admin/ec/{ec_id}")); let response = handle_admin_ec_lookup(Some(&kv), &test_registry(), &req) + .await .expect("should handle lookup"); assert_eq!(response.status(), StatusCode::OK); @@ -1147,35 +1158,38 @@ mod tests { ); } - #[test] - fn missing_entry_returns_404() { + #[tokio::test(flavor = "current_thread")] + async fn missing_entry_returns_404() { let kv = KvIdentityGraph::in_memory("test-store"); let req = get_request(&format!("/_ts/admin/ec/{}", test_ec_id())); let response = handle_admin_ec_lookup(Some(&kv), &test_registry(), &req) + .await .expect("should handle lookup"); assert_eq!(response.status(), StatusCode::NOT_FOUND); } - #[test] - fn invalid_id_returns_400() { + #[tokio::test(flavor = "current_thread")] + async fn invalid_id_returns_400() { let kv = KvIdentityGraph::in_memory("test-store"); let req = get_request("/_ts/admin/ec/not-a-valid-id"); let response = handle_admin_ec_lookup(Some(&kv), &test_registry(), &req) + .await .expect("should handle lookup"); assert_eq!(response.status(), StatusCode::BAD_REQUEST); } - #[test] - fn corrupt_entry_returns_parse_error_and_raw_body() { + #[tokio::test(flavor = "current_thread")] + async fn corrupt_entry_returns_parse_error_and_raw_body() { let ec_id = test_ec_id(); - let kv = kv_with_raw_body(&ec_id, "not json at all"); + let kv = kv_with_raw_body(&ec_id, "not json at all").await; let req = get_request(&format!("/_ts/admin/ec/{ec_id}")); let response = handle_admin_ec_lookup(Some(&kv), &test_registry(), &req) + .await .expect("should handle lookup"); assert_eq!( @@ -1203,8 +1217,8 @@ mod tests { ); } - #[test] - fn invalid_schema_version_reports_validation_error() { + #[tokio::test(flavor = "current_thread")] + async fn invalid_schema_version_reports_validation_error() { let ec_id = test_ec_id(); let body = serde_json::json!({ "v": 99, @@ -1213,10 +1227,11 @@ mod tests { "geo": { "country": "US" } }) .to_string(); - let kv = kv_with_raw_body(&ec_id, &body); + let kv = kv_with_raw_body(&ec_id, &body).await; let req = get_request(&format!("/_ts/admin/ec/{ec_id}")); let response = handle_admin_ec_lookup(Some(&kv), &test_registry(), &req) + .await .expect("should handle lookup"); assert_eq!(response.status(), StatusCode::OK); @@ -1235,13 +1250,14 @@ mod tests { ); } - #[test] - fn bare_route_uses_ts_ec_cookie() { + #[tokio::test(flavor = "current_thread")] + async fn bare_route_uses_ts_ec_cookie() { let ec_id = test_ec_id(); - let kv = kv_with_entry(&ec_id, &sample_entry()); + let kv = kv_with_entry(&ec_id, &sample_entry()).await; let req = get_request_with_cookie("/_ts/admin/ec", &format!("other=1; ts-ec={ec_id}; x=2")); let response = handle_admin_ec_lookup(Some(&kv), &test_registry(), &req) + .await .expect("should handle lookup"); assert_eq!(response.status(), StatusCode::OK); @@ -1253,20 +1269,21 @@ mod tests { ); } - #[test] - fn bare_route_uses_last_duplicate_ts_ec_cookie() { + #[tokio::test(flavor = "current_thread")] + async fn bare_route_uses_last_duplicate_ts_ec_cookie() { // `CookieJar` — the parser the live EC lifecycle uses — keeps the last // duplicate pair. The diagnostic must inspect that same record, not the // first one a header scan would find. let stale_ec_id = format!("{}.stale1", "b".repeat(64)); let live_ec_id = test_ec_id(); - let kv = kv_with_entry(&live_ec_id, &sample_entry()); + let kv = kv_with_entry(&live_ec_id, &sample_entry()).await; let req = get_request_with_cookie( "/_ts/admin/ec", &format!("ts-ec={stale_ec_id}; ts-ec={live_ec_id}"), ); let response = handle_admin_ec_lookup(Some(&kv), &test_registry(), &req) + .await .expect("should handle lookup"); assert_eq!(response.status(), StatusCode::OK); @@ -1278,12 +1295,13 @@ mod tests { ); } - #[test] - fn bare_route_without_cookie_returns_404() { + #[tokio::test(flavor = "current_thread")] + async fn bare_route_without_cookie_returns_404() { let kv = KvIdentityGraph::in_memory("test-store"); let req = get_request("/_ts/admin/ec"); let response = handle_admin_ec_lookup(Some(&kv), &test_registry(), &req) + .await .expect("should handle lookup"); assert_eq!(response.status(), StatusCode::NOT_FOUND); @@ -1297,12 +1315,13 @@ mod tests { ); } - #[test] - fn missing_identity_graph_returns_501() { + #[tokio::test(flavor = "current_thread")] + async fn missing_identity_graph_returns_501() { let req = get_request(&format!("/_ts/admin/ec/{}", test_ec_id())); - let response = - handle_admin_ec_lookup(None, &test_registry(), &req).expect("should handle lookup"); + let response = handle_admin_ec_lookup(None, &test_registry(), &req) + .await + .expect("should handle lookup"); assert_eq!(response.status(), StatusCode::NOT_IMPLEMENTED); } @@ -1331,12 +1350,12 @@ mod tests { ); } - #[test] - fn kv_read_failure_propagates() { + #[tokio::test(flavor = "current_thread")] + async fn kv_read_failure_propagates() { let kv = KvIdentityGraph::failing("broken-store"); let req = get_request(&format!("/_ts/admin/ec/{}", test_ec_id())); - let result = handle_admin_ec_lookup(Some(&kv), &test_registry(), &req); + let result = handle_admin_ec_lookup(Some(&kv), &test_registry(), &req).await; assert!(result.is_err(), "should propagate KV read failures"); } diff --git a/crates/trusted-server-core/src/ec/batch_sync.rs b/crates/trusted-server-core/src/ec/batch_sync.rs index e0a4bb0cb..d314696d9 100644 --- a/crates/trusted-server-core/src/ec/batch_sync.rs +++ b/crates/trusted-server-core/src/ec/batch_sync.rs @@ -43,8 +43,9 @@ const MAX_BODY_SIZE: usize = 2 * 1024 * 1024; use super::kv_types::MAX_UID_LENGTH; +#[async_trait::async_trait(?Send)] trait BatchSyncWriter { - fn upsert_partner_id_if_exists( + async fn upsert_partner_id_if_exists( &self, ec_id: &str, partner_id: &str, @@ -52,14 +53,15 @@ trait BatchSyncWriter { ) -> Result>; } +#[async_trait::async_trait(?Send)] impl BatchSyncWriter for KvIdentityGraph { - fn upsert_partner_id_if_exists( + async fn upsert_partner_id_if_exists( &self, ec_id: &str, partner_id: &str, uid: &str, ) -> Result> { - KvIdentityGraph::upsert_partner_id_if_exists(self, ec_id, partner_id, uid) + KvIdentityGraph::upsert_partner_id_if_exists(self, ec_id, partner_id, uid).await } } @@ -175,7 +177,7 @@ async fn handle_batch_sync_with_writer( } // 4. Process mappings with per-item validation and rejection reasons. - let (accepted, errors) = process_mappings(writer, &partner.source_domain, &body.mappings); + let (accepted, errors) = process_mappings(writer, &partner.source_domain, &body.mappings).await; let rejected = errors.len(); let status = if rejected > 0 { @@ -199,7 +201,7 @@ async fn handle_batch_sync_with_writer( /// group. On infrastructure failure, the failing group and all unprocessed /// valid groups are rejected as unavailable; already validated invalid inputs /// keep their specific errors. Errors are sorted by original input index. -fn process_mappings( +async fn process_mappings( writer: &dyn BatchSyncWriter, partner_id: &str, mappings: &[SyncMapping], @@ -244,7 +246,10 @@ fn process_mappings( let mut accepted = 0; for (group_index, group) in groups.iter().enumerate() { - match writer.upsert_partner_id_if_exists(&group.ec_id, partner_id, &group.partner_uid) { + match writer + .upsert_partner_id_if_exists(&group.ec_id, partner_id, &group.partner_uid) + .await + { Ok(UpsertResult::Written | UpsertResult::Unchanged) => { accepted += group.indexes.len(); } @@ -369,8 +374,9 @@ mod tests { } } + #[async_trait::async_trait(?Send)] impl BatchSyncWriter for MockWriter { - fn upsert_partner_id_if_exists( + async fn upsert_partner_id_if_exists( &self, ec_id: &str, partner_id: &str, @@ -465,8 +471,8 @@ mod tests { ); } - #[test] - fn handle_batch_sync_uses_post_read_limit_for_malformed_content_length() { + #[tokio::test(flavor = "current_thread")] + async fn handle_batch_sync_uses_post_read_limit_for_malformed_content_length() { let writer = MockWriter::new(vec![]); let registry = test_registry(); let limiter = MockRateLimiter { @@ -566,8 +572,8 @@ mod tests { assert!(writer.calls().is_empty(), "should not write any mapping"); } - #[test] - fn process_mappings_returns_multistatus_errors_per_mapping() { + #[tokio::test(flavor = "current_thread")] + async fn process_mappings_returns_multistatus_errors_per_mapping() { let writer = MockWriter::new(vec![Ok(UpsertResult::Written)]); let mappings = vec![ mapping("x", "u1", 1), @@ -575,7 +581,7 @@ mod tests { mapping(&format!("{}.ABC123", "a".repeat(64)), "u3", 1), ]; - let (accepted, errors) = process_mappings(&writer, "partner", &mappings); + let (accepted, errors) = process_mappings(&writer, "partner", &mappings).await; assert_eq!(accepted, 1, "should count successful writes as accepted"); assert_eq!(errors.len(), 2, "should reject invalid mappings only"); @@ -585,8 +591,8 @@ mod tests { assert_eq!(errors[1].reason, REASON_INVALID_PARTNER_UID); } - #[test] - fn process_mappings_aborts_on_kv_unavailable() { + #[tokio::test(flavor = "current_thread")] + async fn process_mappings_aborts_on_kv_unavailable() { let writer = MockWriter::new(vec![ Ok(UpsertResult::Written), Err(Report::new(TrustedServerError::KvStore { @@ -601,7 +607,7 @@ mod tests { mapping(&format!("{}.ABC123", "c".repeat(64)), "u3", 1), ]; - let (accepted, errors) = process_mappings(&writer, "partner", &mappings); + let (accepted, errors) = process_mappings(&writer, "partner", &mappings).await; assert_eq!(accepted, 1, "should keep accepted count before failure"); assert_eq!( @@ -620,8 +626,8 @@ mod tests { ); } - #[test] - fn handle_batch_sync_rejects_missing_auth() { + #[tokio::test(flavor = "current_thread")] + async fn handle_batch_sync_rejects_missing_auth() { let kv = KvIdentityGraph::failing("test_store"); let registry = PartnerRegistry::empty(); let limiter = MockRateLimiter { @@ -686,8 +692,8 @@ mod tests { assert_eq!(json["errors"][0]["reason"], REASON_INELIGIBLE); } - #[test] - fn process_mappings_collapses_missing_and_withdrawn_to_ineligible() { + #[tokio::test(flavor = "current_thread")] + async fn process_mappings_collapses_missing_and_withdrawn_to_ineligible() { let writer = MockWriter::new(vec![ Ok(UpsertResult::NotFound), Ok(UpsertResult::ConsentWithdrawn), @@ -699,7 +705,7 @@ mod tests { mapping(&withdrawn_ec_id, "uid-2", 101), ]; - let (accepted, errors) = process_mappings(&writer, "partner", &mappings); + let (accepted, errors) = process_mappings(&writer, "partner", &mappings).await; assert_eq!(accepted, 0, "should not accept ineligible mappings"); assert_eq!(errors.len(), 2, "should report both errors"); @@ -710,13 +716,13 @@ mod tests { assert_eq!(writer.calls().len(), 2, "should exercise both outcomes"); } - #[test] - fn process_mappings_fans_out_unchanged_to_group_members() { + #[tokio::test(flavor = "current_thread")] + async fn process_mappings_fans_out_unchanged_to_group_members() { let writer = MockWriter::new(vec![Ok(UpsertResult::Unchanged)]); let ec_id = format!("{}.ABC123", "a".repeat(64)); let mappings = vec![mapping(&ec_id, "uid-1", 100), mapping(&ec_id, "uid-1", 101)]; - let (accepted, errors) = process_mappings(&writer, "partner", &mappings); + let (accepted, errors) = process_mappings(&writer, "partner", &mappings).await; assert_eq!(accepted, 2, "should accept every unchanged group member"); assert!( @@ -726,8 +732,8 @@ mod tests { assert_eq!(writer.calls().len(), 1, "should call once for the group"); } - #[test] - fn process_mappings_does_not_order_by_timestamp() { + #[tokio::test(flavor = "current_thread")] + async fn process_mappings_does_not_order_by_timestamp() { let writer = MockWriter::new(vec![Ok(UpsertResult::Written)]); let ec_id = format!("{}.ABC123", "a".repeat(64)); let mappings = vec![ @@ -735,7 +741,7 @@ mod tests { mapping(&ec_id, "uid-old", 100), ]; - let (accepted, errors) = process_mappings(&writer, "partner", &mappings); + let (accepted, errors) = process_mappings(&writer, "partner", &mappings).await; assert_eq!( accepted, 2, @@ -753,8 +759,8 @@ mod tests { ); } - #[test] - fn process_mappings_groups_normalized_ids_in_first_occurrence_order() { + #[tokio::test(flavor = "current_thread")] + async fn process_mappings_groups_normalized_ids_in_first_occurrence_order() { let writer = MockWriter::new(vec![Ok(UpsertResult::Written), Ok(UpsertResult::Unchanged)]); let ec_id_a = format!("{}.ABC123", "a".repeat(64)); let ec_id_a_upper = format!("{}.ABC123", "A".repeat(64)); @@ -765,7 +771,7 @@ mod tests { mapping(&ec_id_a_upper, "a-last", 3), ]; - let (accepted, errors) = process_mappings(&writer, "partner", &mappings); + let (accepted, errors) = process_mappings(&writer, "partner", &mappings).await; assert_eq!(accepted, 3, "should accept every valid group member"); assert!(errors.is_empty(), "should report no errors"); @@ -787,15 +793,15 @@ mod tests { ); } - #[test] - fn process_mappings_bounds_writer_calls_at_max_batch_size() { + #[tokio::test(flavor = "current_thread")] + async fn process_mappings_bounds_writer_calls_at_max_batch_size() { let writer = MockWriter::new(vec![Ok(UpsertResult::Written)]); let ec_id = format!("{}.ABC123", "a".repeat(64)); let mappings = (0..MAX_BATCH_SIZE) .map(|index| mapping(&ec_id, &format!("uid-{index}"), index as u64)) .collect::>(); - let (accepted, errors) = process_mappings(&writer, "partner", &mappings); + let (accepted, errors) = process_mappings(&writer, "partner", &mappings).await; assert_eq!( accepted, MAX_BATCH_SIZE, @@ -813,8 +819,8 @@ mod tests { ); } - #[test] - fn process_mappings_keeps_suffix_case_distinct() { + #[tokio::test(flavor = "current_thread")] + async fn process_mappings_keeps_suffix_case_distinct() { let writer = MockWriter::new(vec![Ok(UpsertResult::Written), Ok(UpsertResult::Written)]); let upper_suffix = format!("{}.ABC123", "a".repeat(64)); let mixed_suffix = format!("{}.AbC123", "a".repeat(64)); @@ -823,7 +829,7 @@ mod tests { mapping(&mixed_suffix, "mixed", 2), ]; - let (accepted, errors) = process_mappings(&writer, "partner", &mappings); + let (accepted, errors) = process_mappings(&writer, "partner", &mappings).await; assert_eq!(accepted, 2, "should accept both distinct EC IDs"); assert!(errors.is_empty(), "should report no errors"); @@ -845,8 +851,8 @@ mod tests { ); } - #[test] - fn process_mappings_invalid_duplicate_does_not_replace_last_valid_uid() { + #[tokio::test(flavor = "current_thread")] + async fn process_mappings_invalid_duplicate_does_not_replace_last_valid_uid() { let writer = MockWriter::new(vec![Ok(UpsertResult::Written)]); let ec_id = format!("{}.ABC123", "a".repeat(64)); let mappings = vec![ @@ -855,7 +861,7 @@ mod tests { mapping(&ec_id, " ", 3), ]; - let (accepted, errors) = process_mappings(&writer, "partner", &mappings); + let (accepted, errors) = process_mappings(&writer, "partner", &mappings).await; assert_eq!(accepted, 2, "should accept valid group members"); assert_eq!(errors.len(), 1, "should retain the invalid UID error"); @@ -868,8 +874,8 @@ mod tests { ); } - #[test] - fn process_mappings_fans_out_ineligible_outcomes_to_group_members() { + #[tokio::test(flavor = "current_thread")] + async fn process_mappings_fans_out_ineligible_outcomes_to_group_members() { let writer = MockWriter::new(vec![ Ok(UpsertResult::NotFound), Ok(UpsertResult::ConsentWithdrawn), @@ -883,7 +889,7 @@ mod tests { mapping(&ec_id_b, "b-2", 4), ]; - let (accepted, errors) = process_mappings(&writer, "partner", &mappings); + let (accepted, errors) = process_mappings(&writer, "partner", &mappings).await; assert_eq!(accepted, 0, "should reject all ineligible group members"); assert_eq!( @@ -903,8 +909,8 @@ mod tests { assert_eq!(writer.calls().len(), 2, "should call once per group"); } - #[test] - fn process_mappings_aborts_by_group_and_preserves_sorted_accounting() { + #[tokio::test(flavor = "current_thread")] + async fn process_mappings_aborts_by_group_and_preserves_sorted_accounting() { let writer = MockWriter::new(vec![ Ok(UpsertResult::Written), Err(Report::new(TrustedServerError::KvStore { @@ -924,7 +930,7 @@ mod tests { mapping(&ec_id_c, "", 6), ]; - let (accepted, errors) = process_mappings(&writer, "partner", &mappings); + let (accepted, errors) = process_mappings(&writer, "partner", &mappings).await; assert_eq!( accepted, 2, @@ -955,8 +961,8 @@ mod tests { ); } - #[test] - fn handle_batch_sync_accepts_every_member_of_a_written_group() { + #[tokio::test(flavor = "current_thread")] + async fn handle_batch_sync_accepts_every_member_of_a_written_group() { let registry = test_registry(); let limiter = MockRateLimiter { should_exceed: false, @@ -985,8 +991,8 @@ mod tests { assert_eq!(json["rejected"], 0, "should reject no group members"); } - #[test] - fn handle_batch_sync_rejects_every_member_of_an_ineligible_group() { + #[tokio::test(flavor = "current_thread")] + async fn handle_batch_sync_rejects_every_member_of_an_ineligible_group() { let registry = test_registry(); let limiter = MockRateLimiter { should_exceed: false, @@ -1023,8 +1029,8 @@ mod tests { ); } - #[test] - fn handle_batch_sync_reports_validation_errors_without_writer_calls() { + #[tokio::test(flavor = "current_thread")] + async fn handle_batch_sync_reports_validation_errors_without_writer_calls() { let writer = MockWriter::new(vec![]); let registry = test_registry(); let limiter = MockRateLimiter { @@ -1060,8 +1066,8 @@ mod tests { ); } - #[test] - fn handle_batch_sync_reports_grouped_infrastructure_failure() { + #[tokio::test(flavor = "current_thread")] + async fn handle_batch_sync_reports_grouped_infrastructure_failure() { let writer = MockWriter::new(vec![ Ok(UpsertResult::Written), Err(Report::new(TrustedServerError::KvStore { diff --git a/crates/trusted-server-core/src/ec/finalize.rs b/crates/trusted-server-core/src/ec/finalize.rs index 019ef2aa4..ccb0ef349 100644 --- a/crates/trusted-server-core/src/ec/finalize.rs +++ b/crates/trusted-server-core/src/ec/finalize.rs @@ -44,7 +44,7 @@ const EC_RESPONSE_HEADERS: &[&str] = &[ /// /// `eids_cookie` should be the raw value of the `ts-eids` cookie extracted /// from the request *before* routing consumes it. -pub fn ec_finalize_response( +pub async fn ec_finalize_response( settings: &Settings, ec_context: &mut EcContext, kv: Option<&KvIdentityGraph>, @@ -71,7 +71,8 @@ pub fn ec_finalize_response( registry, consent_withdrawn, response, - ); + ) + .await; return; } @@ -107,17 +108,18 @@ pub fn ec_finalize_response( && ec_context.recovery_eligible() && matches!(ec_context.kv_snapshot(), EcKvSnapshot::NotRead) { - let snapshot = graph.load_snapshot(&ec_id); + let snapshot = graph.load_snapshot(&ec_id).await; ec_context.set_kv_snapshot(snapshot); } else if let Some(source) = source { - sync_eid_cookie_updates(graph, ec_context, &ec_id, &updates, source); + sync_eid_cookie_updates(graph, ec_context, &ec_id, &updates, source).await; } if matches!(ec_context.kv_snapshot(), EcKvSnapshot::Missing { .. }) && ec_context.recovery_eligible() { confirm_then_recover_orphaned_ec( settings, ec_context, graph, &ec_id, &updates, response, - ); + ) + .await; } } @@ -147,7 +149,7 @@ pub fn ec_finalize_response( ), ec_context.consent(), ); - sync_eid_cookie_updates(graph, ec_context, &ec_id, &updates, EidSyncSource::NewEc); + sync_eid_cookie_updates(graph, ec_context, &ec_id, &updates, EidSyncSource::NewEc).await; if ec_context.kv_snapshot().entry_for(&ec_id).is_some() { set_ec_cookie_on_response(settings, ec_context, response); } else { @@ -158,7 +160,7 @@ pub fn ec_finalize_response( reconcile_pull_sync_marker(settings, registry, ec_context, response); } -fn sync_eid_cookie_updates( +async fn sync_eid_cookie_updates( graph: &KvIdentityGraph, ec_context: &mut EcContext, ec_id: &str, @@ -169,11 +171,9 @@ fn sync_eid_cookie_updates( return; } - let (snapshot, outcome) = graph.sync_eid_cookie_updates_from_snapshot( - ec_id, - updates, - ec_context.kv_snapshot().clone(), - ); + let (snapshot, outcome) = graph + .sync_eid_cookie_updates_from_snapshot(ec_id, updates, ec_context.kv_snapshot().clone()) + .await; ec_context.set_kv_snapshot(snapshot); record_eid_sync_terminal(source, outcome); } @@ -268,7 +268,7 @@ fn gate_eid_updates_by_consent( Vec::new() } -fn recover_orphaned_ec( +async fn recover_orphaned_ec( settings: &Settings, ec_context: &mut EcContext, graph: &KvIdentityGraph, @@ -311,7 +311,7 @@ fn recover_orphaned_ec( .map(super::device::DeviceSignals::to_kv_device); apply_partner_id_updates(&mut entry, updates); - match graph.create_if_absent(&ec_id, &entry) { + match graph.create_if_absent(&ec_id, &entry).await { Ok(CreateIfAbsentOutcome::Written) => { let snapshot = EcKvSnapshot::Present { ec_id: ec_id.clone(), @@ -357,7 +357,7 @@ fn recover_orphaned_ec( /// the identity stays intact and recovery is retried on a later navigation; /// - neither a read failure nor a failed existence check is a miss, and neither /// rotates. -fn confirm_then_recover_orphaned_ec( +async fn confirm_then_recover_orphaned_ec( settings: &Settings, ec_context: &mut EcContext, graph: &KvIdentityGraph, @@ -365,16 +365,17 @@ fn confirm_then_recover_orphaned_ec( updates: &[super::kv::PartnerIdUpdate], response: &mut Response, ) { - let confirmed = graph.load_snapshot(ec_id); + let confirmed = graph.load_snapshot(ec_id).await; match confirmed { EcKvSnapshot::Present { .. } => { // The row became visible after the origin round trip: adopt it and // merge any pending updates rather than rotating a valid identity. ec_context.set_kv_snapshot(confirmed); - sync_eid_cookie_updates(graph, ec_context, ec_id, updates, EidSyncSource::Navigation); + sync_eid_cookie_updates(graph, ec_context, ec_id, updates, EidSyncSource::Navigation) + .await; } - EcKvSnapshot::Missing { .. } => match graph.key_exists_confirmed(ec_id) { - Ok(false) => recover_orphaned_ec(settings, ec_context, graph, updates, response), + EcKvSnapshot::Missing { .. } => match graph.key_exists_confirmed(ec_id).await { + Ok(false) => recover_orphaned_ec(settings, ec_context, graph, updates, response).await, Ok(true) => { log::warn!( "Orphan EC recovery skipped: both point reads missed a row the store still \ @@ -443,7 +444,7 @@ pub fn clear_ec_on_response(settings: &Settings, response: &mut Response, @@ -465,7 +466,7 @@ fn finalize_unusable_consent( // The identity-graph tombstone is the authoritative withdrawal marker // for subsequent EC behavior. if let Some(graph) = kv { - apply_withdrawal_tombstones(&ids_to_withdraw, |ec_id| { + for ec_id in &ids_to_withdraw { // The graph hands back the post-withdrawal snapshot rather than // leaving the caller to rebuild it, so post-send work that reads // the context — pull sync discloses the raw EC ID to partners — @@ -476,7 +477,7 @@ fn finalize_unusable_consent( } else { EcKvSnapshot::NotRead }; - let outcome = graph.tombstone_existing_from_snapshot(ec_id, initial); + let outcome = graph.tombstone_existing_from_snapshot(ec_id, initial).await; // The browser cookie is already cleared, so a failed tombstone // leaves a live row that server-side consumers still read as // consented. Report every failure, including the non-active cookie @@ -491,7 +492,7 @@ fn finalize_unusable_consent( if ec_context.ec_value() == Some(ec_id) { ec_context.set_kv_snapshot(outcome); } - }); + } } } @@ -513,15 +514,6 @@ fn withdrawal_ec_ids(ec_context: &EcContext) -> HashSet { hashes } -fn apply_withdrawal_tombstones(ec_ids: &HashSet, mut write_tombstone: F) -where - F: FnMut(&str), -{ - for ec_id in ec_ids { - write_tombstone(ec_id); - } -} - #[cfg(test)] mod tests { use http::HeaderValue; @@ -687,23 +679,6 @@ mod tests { assert!(ids.contains(&valid_ec), "should keep the valid EC ID"); } - #[test] - fn apply_withdrawal_tombstones_invokes_writer_for_each_ec_id() { - let first = sample_ec_id("first1"); - let second = sample_ec_id("second"); - let mut ids = HashSet::new(); - ids.insert(first.clone()); - ids.insert(second.clone()); - - let mut written = Vec::new(); - apply_withdrawal_tombstones(&ids, |ec_id| written.push(ec_id.to_owned())); - written.sort(); - - let mut expected = vec![first, second]; - expected.sort(); - assert_eq!(written, expected, "should write a tombstone for each EC ID"); - } - #[test] fn clear_ec_on_response_removes_headers_and_expires_cookie() { let settings = create_test_settings(); @@ -739,8 +714,8 @@ mod tests { ); } - #[test] - fn finalize_withdrawal_does_not_create_a_row_for_an_unheld_identity() { + #[tokio::test(flavor = "current_thread")] + async fn finalize_withdrawal_does_not_create_a_row_for_an_unheld_identity() { let settings = create_test_settings(); // The cookie value is chosen by the client, so a withdrawal naming an // identity this deployment never issued must not put a row in the @@ -766,10 +741,11 @@ mod tests { None, None, &mut response, - ); + ) + .await; assert!( - kv.get(&ec_id).expect("should read back").is_none(), + kv.get(&ec_id).await.expect("should read back").is_none(), "should not write a tombstone for an identity that was never issued" ); assert!( @@ -783,8 +759,8 @@ mod tests { ); } - #[test] - fn finalize_withdrawal_tombstones_a_held_identity() { + #[tokio::test(flavor = "current_thread")] + async fn finalize_withdrawal_tombstones_a_held_identity() { let settings = create_test_settings(); let ec_id = sample_ec_id("held01"); let consent = ConsentContext { @@ -800,8 +776,9 @@ mod tests { &ec_id, &crate::ec::kv_types::KvEntry::minimal("p.example", "uid", 1), ) + .await .expect("should seed the identity"); - ec_context.set_kv_snapshot(kv.load_snapshot(&ec_id)); + ec_context.set_kv_snapshot(kv.load_snapshot(&ec_id).await); ec_context.set_eid_sync_source(EidSyncSource::Auction); assert!(matches!( ec_context.kv_snapshot(), @@ -818,10 +795,12 @@ mod tests { None, None, &mut response, - ); + ) + .await; let (entry, _) = kv .get(&ec_id) + .await .expect("should read back") .expect("should still hold the identity"); assert!( @@ -836,8 +815,8 @@ mod tests { assert_eq!(ec_context.kv_snapshot().generation_for(&ec_id), None); } - #[test] - fn withdrawal_still_expires_the_cookie_when_the_store_is_unavailable() { + #[tokio::test(flavor = "current_thread")] + async fn withdrawal_still_expires_the_cookie_when_the_store_is_unavailable() { // Cookie expiry is the primary enforcement, so it has to survive a // store that cannot answer at all — the case where the identity-graph // marker is exactly what goes missing. @@ -869,7 +848,8 @@ mod tests { None, None, &mut response, - ); + ) + .await; let set_cookie = get_header_str(&response, "set-cookie").unwrap_or_default(); assert!( @@ -886,8 +866,8 @@ mod tests { ); } - #[test] - fn finalize_withdrawal_clears_cookie_and_headers() { + #[tokio::test(flavor = "current_thread")] + async fn finalize_withdrawal_clears_cookie_and_headers() { let settings = create_test_settings(); let ec_id = sample_ec_id("aBc123"); let consent = ConsentContext { @@ -914,7 +894,8 @@ mod tests { None, None, &mut response, - ); + ) + .await; assert!( get_header(&response, "x-ts-ec").is_none(), @@ -943,8 +924,8 @@ mod tests { ); } - #[test] - fn finalize_returning_user_with_cookie_mismatch_sets_no_header_or_cookie() { + #[tokio::test(flavor = "current_thread")] + async fn finalize_returning_user_with_cookie_mismatch_sets_no_header_or_cookie() { let settings = create_test_settings(); let active_ec = sample_ec_id("activ1"); let cookie_ec = sample_ec_id("cook1e"); @@ -966,7 +947,8 @@ mod tests { None, None, &mut response, - ); + ) + .await; assert!( get_header(&response, "x-ts-ec").is_none(), @@ -978,8 +960,8 @@ mod tests { ); } - #[test] - fn finalize_returning_user_sets_no_header_or_cookie() { + #[tokio::test(flavor = "current_thread")] + async fn finalize_returning_user_sets_no_header_or_cookie() { let settings = create_test_settings(); let ec_id = sample_ec_id("mtch01"); let mut ec_context = make_context( @@ -1000,7 +982,8 @@ mod tests { None, None, &mut response, - ); + ) + .await; assert!( get_header(&response, "x-ts-ec").is_none(), @@ -1012,8 +995,8 @@ mod tests { ); } - #[test] - fn finalize_generated_ec_without_kv_skips_cookie_and_header() { + #[tokio::test(flavor = "current_thread")] + async fn finalize_generated_ec_without_kv_skips_cookie_and_header() { let settings = create_test_settings(); let generated_ec = sample_ec_id("gen123"); let mut ec_context = make_context( @@ -1034,7 +1017,8 @@ mod tests { None, None, &mut response, - ); + ) + .await; assert!( get_header(&response, "x-ts-ec").is_none(), @@ -1046,8 +1030,8 @@ mod tests { ); } - #[test] - fn finalize_rotates_orphaned_cookie_to_new_backed_ec() { + #[tokio::test(flavor = "current_thread")] + async fn finalize_rotates_orphaned_cookie_to_new_backed_ec() { let settings = create_test_settings(); let orphaned_ec = sample_ec_id("orphn1"); let consent = ConsentContext { @@ -1075,13 +1059,15 @@ mod tests { None, None, &mut response, - ); + ) + .await; let replacement = ec_context.ec_value().expect("should rotate orphan"); assert_ne!(replacement, orphaned_ec); assert!( graph .get(replacement) + .await .expect("should read replacement") .is_some(), "replacement cookie should have a backing row" @@ -1092,8 +1078,8 @@ mod tests { ); } - #[test] - fn finalize_recovers_orphaned_ec_when_purpose_four_denial_empties_updates() { + #[tokio::test(flavor = "current_thread")] + async fn finalize_recovers_orphaned_ec_when_purpose_four_denial_empties_updates() { // Regression test: this request has real EID data (a configured // partner and a captured client EID), but TCF Purpose 4 denial gates // it away, emptying the update list. That must not also suppress the @@ -1136,7 +1122,8 @@ mod tests { None, None, &mut response, - ); + ) + .await; let replacement = ec_context.ec_value().expect( "should rotate the orphan even though Purpose 4 denial emptied the gated update list", @@ -1144,6 +1131,7 @@ mod tests { assert_ne!(replacement, &orphaned_ec); let (stored, _) = graph .get(replacement) + .await .expect("should read replacement") .expect("replacement cookie should have a backing row"); assert!( @@ -1156,8 +1144,8 @@ mod tests { ); } - #[test] - fn valid_marker_with_unread_snapshot_defers_orphan_recovery() { + #[tokio::test(flavor = "current_thread")] + async fn valid_marker_with_unread_snapshot_defers_orphan_recovery() { let settings = create_test_settings(); let orphaned_ec = sample_ec_id("orphn2"); let consent = ConsentContext { @@ -1191,7 +1179,8 @@ mod tests { None, None, &mut response, - ); + ) + .await; assert_eq!( ec_context.ec_value(), @@ -1208,8 +1197,8 @@ mod tests { ); } - #[test] - fn finalize_returning_user_subresource_does_not_persist_eid_updates() { + #[tokio::test(flavor = "current_thread")] + async fn finalize_returning_user_subresource_does_not_persist_eid_updates() { let settings = create_test_settings(); let ec_id = sample_ec_id("subeid"); let graph = KvIdentityGraph::in_memory("test_store"); @@ -1221,8 +1210,10 @@ mod tests { ); graph .create(&ec_id, &live) + .await .expect("should seed the live row"); - let mut ec_context = returning_user_context(&ec_id, graph.load_snapshot(&ec_id), false); + let mut ec_context = + returning_user_context(&ec_id, graph.load_snapshot(&ec_id).await, false); let partners = vec![make_partner("sharedid.org")]; let registry = PartnerRegistry::from_config(&partners).expect("should build registry"); let mut response = empty_response(); @@ -1235,10 +1226,12 @@ mod tests { None, Some("shared-cookie-id"), &mut response, - ); + ) + .await; let (stored, _) = graph .get(&ec_id) + .await .expect("should read store") .expect("row should remain"); assert!( @@ -1247,8 +1240,8 @@ mod tests { ); } - #[test] - fn finalize_navigation_routes_persist_returning_user_eid_updates() { + #[tokio::test(flavor = "current_thread")] + async fn finalize_navigation_routes_persist_returning_user_eid_updates() { for (source, suffix, cookie_id) in [ (EidSyncSource::Navigation, "naveid", "navigation-cookie-id"), (EidSyncSource::PageBids, "spaeid", "page-bids-cookie-id"), @@ -1264,8 +1257,10 @@ mod tests { ); graph .create(&ec_id, &live) + .await .expect("should seed the live row"); - let mut ec_context = returning_user_context(&ec_id, graph.load_snapshot(&ec_id), true); + let mut ec_context = + returning_user_context(&ec_id, graph.load_snapshot(&ec_id).await, true); ec_context.set_eid_sync_source(source); let partners = vec![make_partner("sharedid.org")]; let registry = PartnerRegistry::from_config(&partners).expect("should build registry"); @@ -1279,10 +1274,12 @@ mod tests { None, Some(cookie_id), &mut response, - ); + ) + .await; let (stored, _) = graph .get(&ec_id) + .await .expect("should read store") .expect("row should remain"); assert_eq!( @@ -1293,8 +1290,8 @@ mod tests { } } - #[test] - fn finalize_generated_ec_persists_eid_updates() { + #[tokio::test(flavor = "current_thread")] + async fn finalize_generated_ec_persists_eid_updates() { let settings = create_test_settings(); let ec_id = sample_ec_id("geneid"); let graph = KvIdentityGraph::in_memory("test_store"); @@ -1306,10 +1303,11 @@ mod tests { ); graph .create(&ec_id, &live) + .await .expect("should seed generated row"); let mut ec_context = make_context(Some(&ec_id), None, false, true, Jurisdiction::NonRegulated); - ec_context.set_kv_snapshot(graph.load_snapshot(&ec_id)); + ec_context.set_kv_snapshot(graph.load_snapshot(&ec_id).await); let partners = vec![make_partner("sharedid.org")]; let registry = PartnerRegistry::from_config(&partners).expect("should build registry"); let mut response = empty_response(); @@ -1322,10 +1320,12 @@ mod tests { None, Some("generated-cookie-id"), &mut response, - ); + ) + .await; let (stored, _) = graph .get(&ec_id) + .await .expect("should read store") .expect("row should remain"); assert_eq!( @@ -1430,8 +1430,8 @@ mod tests { ); } - #[test] - fn finalize_auction_transient_miss_still_persists_eid_updates() { + #[tokio::test(flavor = "current_thread")] + async fn finalize_auction_transient_miss_still_persists_eid_updates() { // `/auction` saves its first lookup into the context and is never // recovery eligible, so a stale miss there has no later chance to // retry. Finalization must revalidate before dropping collected IDs. @@ -1446,6 +1446,7 @@ mod tests { ); graph .create(&ec_id, &live) + .await .expect("should seed the live row the endpoint lookup missed"); let mut ec_context = returning_user_context( &ec_id, @@ -1467,11 +1468,13 @@ mod tests { None, Some("shared-cookie-id"), &mut response, - ); + ) + .await; assert_did_not_rotate(&ec_context, &ec_id, &response); let (stored, _) = graph .get(&ec_id) + .await .expect("should read store") .expect("row should remain"); assert_eq!( @@ -1481,8 +1484,8 @@ mod tests { ); } - #[test] - fn finalize_persists_every_configured_partner_from_client_eids_over_a_trimmed_cookie() { + #[tokio::test(flavor = "current_thread")] + async fn finalize_persists_every_configured_partner_from_client_eids_over_a_trimmed_cookie() { // Regression test for #1184: the `ts-eids` cookie only carries what // the browser could fit under its size cap, but `/auction` also // hands finalization the full EID set from the request body via @@ -1499,6 +1502,7 @@ mod tests { ); graph .create(&ec_id, &live) + .await .expect("should seed the live row this request updates"); let mut ec_context = returning_user_context( &ec_id, @@ -1546,10 +1550,12 @@ mod tests { Some(&eids_cookie), None, &mut response, - ); + ) + .await; let (stored, _) = graph .get(&ec_id) + .await .expect("should read store") .expect("row should exist after ingestion"); assert_eq!( @@ -1564,8 +1570,8 @@ mod tests { ); } - #[test] - fn finalize_withholds_eid_kv_writes_when_purpose_four_is_denied() { + #[tokio::test(flavor = "current_thread")] + async fn finalize_withholds_eid_kv_writes_when_purpose_four_is_denied() { // A GDPR user can grant TCF Purpose 1 (storage/EC) while denying // Purpose 4 (personalized ads). `gate_eids_by_consent` already strips // EIDs from the outbound /auction bid request in that case; the KV @@ -1583,6 +1589,7 @@ mod tests { ); graph .create(&ec_id, &live) + .await .expect("should seed the live row this request updates"); let mut ec_context = EcContext::new_for_test_with_cookie( Some(ec_id.clone()), @@ -1615,10 +1622,12 @@ mod tests { None, None, &mut response, - ); + ) + .await; let (stored, _) = graph .get(&ec_id) + .await .expect("should read store") .expect("row should remain"); assert!( @@ -1628,8 +1637,8 @@ mod tests { ); } - #[test] - fn finalize_auction_confirmed_miss_does_not_create_a_row() { + #[tokio::test(flavor = "current_thread")] + async fn finalize_auction_confirmed_miss_does_not_create_a_row() { // The same path with a genuinely absent row must stay a no-op: a route // without orphan recovery must never mint an identity-graph entry. let settings = create_test_settings(); @@ -1655,17 +1664,22 @@ mod tests { None, Some("shared-cookie-id"), &mut response, - ); + ) + .await; assert_did_not_rotate(&ec_context, &ec_id, &response); assert!( - graph.get(&ec_id).expect("should read store").is_none(), + graph + .get(&ec_id) + .await + .expect("should read store") + .is_none(), "a confirmed miss must not create a root entry" ); } - #[test] - fn finalize_transient_missing_row_confirms_present_and_does_not_rotate() { + #[tokio::test(flavor = "current_thread")] + async fn finalize_transient_missing_row_confirms_present_and_does_not_rotate() { // The origin-overlapped preload transiently read `Missing` on an // eventually-consistent store, but the row actually exists. The // confirming re-read at finalize must adopt the live row instead of @@ -1681,6 +1695,7 @@ mod tests { ); graph .create(&orphan, &live) + .await .expect("should seed the live row the preload missed"); let mut ec_context = returning_user_context( &orphan, @@ -1699,7 +1714,8 @@ mod tests { None, None, &mut response, - ); + ) + .await; assert_did_not_rotate(&ec_context, &orphan, &response); assert!( @@ -1708,8 +1724,8 @@ mod tests { ); } - #[test] - fn finalize_two_missing_reads_do_not_rotate_a_row_the_store_still_lists() { + #[tokio::test(flavor = "current_thread")] + async fn finalize_two_missing_reads_do_not_rotate_a_row_the_store_still_lists() { // Both the origin-overlapped preload and the confirming re-read missed, // but the row is live — the point reads were stale. Two stale reads are // not an absence proof, so the identity graph must be left intact and @@ -1728,6 +1744,7 @@ mod tests { ); graph .create(&orphan, &live) + .await .expect("should seed the live row both point reads miss"); let mut ec_context = returning_user_context( &orphan, @@ -1746,12 +1763,14 @@ mod tests { None, None, &mut response, - ); + ) + .await; assert_did_not_rotate(&ec_context, &orphan, &response); assert_eq!( graph .get(&orphan) + .await .expect("should read store") .map(|(entry, _)| entry.consent.ok), Some(true), @@ -1759,8 +1778,8 @@ mod tests { ); } - #[test] - fn finalize_rotates_when_only_a_longer_key_shares_the_orphan_prefix() { + #[tokio::test(flavor = "current_thread")] + async fn finalize_rotates_when_only_a_longer_key_shares_the_orphan_prefix() { // `key_exists_confirmed` gates orphan recovery as well as withdrawal. // It matches whole keys, so a neighbouring key that merely starts with // the orphaned ID cannot answer for it. Under the prefix count this @@ -1779,10 +1798,12 @@ mod tests { ); graph .create(&neighbour, &live) + .await .expect("should seed the neighbouring row"); assert!( !graph .key_exists_confirmed(&orphan) + .await .expect("should confirm against the store"), "a longer key sharing the prefix must not prove the orphan exists" ); @@ -1803,7 +1824,8 @@ mod tests { None, None, &mut response, - ); + ) + .await; let replacement = ec_context.ec_value().expect("should rotate the orphan"); assert_ne!( @@ -1813,6 +1835,7 @@ mod tests { assert!( graph .get(replacement) + .await .expect("should read the replacement") .is_some(), "the replacement cookie should have a backing row" @@ -1820,14 +1843,15 @@ mod tests { assert!( graph .get(&neighbour) + .await .expect("should read the neighbour") .is_some(), "the neighbouring identity must be left untouched" ); } - #[test] - fn finalize_does_not_rotate_when_the_existence_check_fails() { + #[tokio::test(flavor = "current_thread")] + async fn finalize_does_not_rotate_when_the_existence_check_fails() { // Absence is unprovable when the list itself errors. Rotation abandons a // year-lived identity, so it must not run on an unproven miss. let settings = create_test_settings(); @@ -1850,13 +1874,14 @@ mod tests { None, None, &mut response, - ); + ) + .await; assert_did_not_rotate(&ec_context, &orphan, &response); } - #[test] - fn finalize_generated_ec_does_not_emit_cookie_for_authoritative_missing_row() { + #[tokio::test(flavor = "current_thread")] + async fn finalize_generated_ec_does_not_emit_cookie_for_authoritative_missing_row() { let settings = create_test_settings(); let generated_ec = sample_ec_id("genmis"); let mut ec_context = make_context( @@ -1880,7 +1905,8 @@ mod tests { None, None, &mut response, - ); + ) + .await; assert!( get_header(&response, "set-cookie").is_none(), @@ -1888,8 +1914,8 @@ mod tests { ); } - #[test] - fn finalize_denied_without_cookie_is_noop() { + #[tokio::test(flavor = "current_thread")] + async fn finalize_denied_without_cookie_is_noop() { let settings = create_test_settings(); let mut ec_context = make_context(None, None, false, false, Jurisdiction::Unknown); let mut response = empty_response(); @@ -1903,7 +1929,8 @@ mod tests { None, None, &mut response, - ); + ) + .await; assert!( get_header(&response, "x-ts-ec").is_none(), @@ -1915,8 +1942,8 @@ mod tests { ); } - #[test] - fn finalize_unknown_jurisdiction_strips_headers_without_expiring_cookie() { + #[tokio::test(flavor = "current_thread")] + async fn finalize_unknown_jurisdiction_strips_headers_without_expiring_cookie() { let settings = create_test_settings(); let ec_id = sample_ec_id("unk001"); let mut ec_context = make_context( @@ -1939,7 +1966,8 @@ mod tests { None, None, &mut response, - ); + ) + .await; assert!( get_header(&response, "x-ts-ec").is_none(), @@ -1995,8 +2023,8 @@ mod tests { ); } - #[test] - fn finalize_not_read_snapshot_does_not_rotate() { + #[tokio::test(flavor = "current_thread")] + async fn finalize_not_read_snapshot_does_not_rotate() { let settings = create_test_settings(); let orphan = sample_ec_id("notrd1"); let mut ec_context = returning_user_context(&orphan, EcKvSnapshot::NotRead, true); @@ -2011,13 +2039,14 @@ mod tests { None, None, &mut response, - ); + ) + .await; assert_did_not_rotate(&ec_context, &orphan, &response); } - #[test] - fn finalize_failed_snapshot_does_not_rotate() { + #[tokio::test(flavor = "current_thread")] + async fn finalize_failed_snapshot_does_not_rotate() { let settings = create_test_settings(); let orphan = sample_ec_id("faild1"); let mut ec_context = returning_user_context( @@ -2038,13 +2067,14 @@ mod tests { None, None, &mut response, - ); + ) + .await; assert_did_not_rotate(&ec_context, &orphan, &response); } - #[test] - fn finalize_tombstone_snapshot_does_not_rotate() { + #[tokio::test(flavor = "current_thread")] + async fn finalize_tombstone_snapshot_does_not_rotate() { let settings = create_test_settings(); let orphan = sample_ec_id("tomb01"); let tombstone = EcKvSnapshot::Present { @@ -2064,13 +2094,14 @@ mod tests { None, None, &mut response, - ); + ) + .await; assert_did_not_rotate(&ec_context, &orphan, &response); } - #[test] - fn finalize_subresource_missing_row_does_not_rotate() { + #[tokio::test(flavor = "current_thread")] + async fn finalize_subresource_missing_row_does_not_rotate() { let settings = create_test_settings(); let orphan = sample_ec_id("subrs1"); // Missing row, but the request is not a recovery-eligible browser navigation. @@ -2092,17 +2123,22 @@ mod tests { None, None, &mut response, - ); + ) + .await; assert_did_not_rotate(&ec_context, &orphan, &response); assert!( - graph.get(&orphan).expect("should read store").is_none(), + graph + .get(&orphan) + .await + .expect("should read store") + .is_none(), "a non-eligible request must not create the missing root" ); } - #[test] - fn finalize_withdrawal_tombstones_present_id_and_skips_missing_other() { + #[tokio::test(flavor = "current_thread")] + async fn finalize_withdrawal_tombstones_present_id_and_skips_missing_other() { let settings = create_test_settings(); let active_ec = sample_ec_id("activ2"); let cookie_ec = sample_ec_id("cook2e"); @@ -2119,8 +2155,9 @@ mod tests { let graph = KvIdentityGraph::in_memory("test_store"); graph .create(&active_ec, &live_entry()) + .await .expect("should seed active row"); - ec_context.set_kv_snapshot(graph.load_snapshot(&active_ec)); + ec_context.set_kv_snapshot(graph.load_snapshot(&active_ec).await); let mut response = empty_response(); ec_finalize_response( @@ -2131,10 +2168,12 @@ mod tests { None, None, &mut response, - ); + ) + .await; let (active_stored, _) = graph .get(&active_ec) + .await .expect("should read active row") .expect("active row should remain as a tombstone"); assert!( @@ -2142,13 +2181,17 @@ mod tests { "the present active ID should be tombstoned after strong confirmation" ); assert!( - graph.get(&cookie_ec).expect("should read store").is_none(), + graph + .get(&cookie_ec) + .await + .expect("should read store") + .is_none(), "a missing second ID must never be created by withdrawal" ); } - #[test] - fn finalize_withdrawal_tombstones_both_present_ids_once() { + #[tokio::test(flavor = "current_thread")] + async fn finalize_withdrawal_tombstones_both_present_ids_once() { let settings = create_test_settings(); let active_ec = sample_ec_id("activ3"); let cookie_ec = sample_ec_id("cook3e"); @@ -2166,14 +2209,16 @@ mod tests { &active_ec, &KvEntry::minimal("active.example.com", "active-uid", 1_000), ) + .await .expect("should seed active row"); graph .create( &cookie_ec, &KvEntry::minimal("cookie.example.com", "cookie-uid", 1_000), ) + .await .expect("should seed cookie row"); - ec_context.set_kv_snapshot(graph.load_snapshot(&active_ec)); + ec_context.set_kv_snapshot(graph.load_snapshot(&active_ec).await); let mut response = empty_response(); ec_finalize_response( @@ -2184,14 +2229,17 @@ mod tests { None, None, &mut response, - ); + ) + .await; let (active_tombstone, active_generation) = graph .get(&active_ec) + .await .expect("should read active row") .expect("should retain active tombstone"); let (cookie_tombstone, cookie_generation) = graph .get(&cookie_ec) + .await .expect("should read cookie row") .expect("should retain cookie tombstone"); assert!( @@ -2220,11 +2268,13 @@ mod tests { None, None, &mut repeated_response, - ); + ) + .await; assert_eq!( graph .get(&active_ec) + .await .expect("should read active row") .expect("should retain active tombstone") .1, @@ -2234,6 +2284,7 @@ mod tests { assert_eq!( graph .get(&cookie_ec) + .await .expect("should read cookie row") .expect("should retain cookie tombstone") .1, @@ -2242,8 +2293,8 @@ mod tests { ); } - #[test] - fn finalize_withdrawal_keeps_cookie_deletion_on_kv_failure() { + #[tokio::test(flavor = "current_thread")] + async fn finalize_withdrawal_keeps_cookie_deletion_on_kv_failure() { let settings = create_test_settings(); let ec_id = sample_ec_id("failw1"); let consent = ConsentContext { @@ -2268,7 +2319,8 @@ mod tests { None, None, &mut response, - ); + ) + .await; let cookies = response .headers() @@ -2295,8 +2347,8 @@ mod tests { ); } - #[test] - fn finalize_sets_marker_for_complete_pull_partner_snapshot() { + #[tokio::test(flavor = "current_thread")] + async fn finalize_sets_marker_for_complete_pull_partner_snapshot() { let settings = create_test_settings(); let ec_id = sample_ec_id("compl1"); let mut partner = make_partner("ssp.example.com"); @@ -2334,7 +2386,8 @@ mod tests { None, None, &mut response, - ); + ) + .await; let cookies = response .headers() @@ -2350,8 +2403,8 @@ mod tests { ); } - #[test] - fn explicit_withdrawal_without_marker_or_ec_cookie_does_not_set_cookie() { + #[tokio::test(flavor = "current_thread")] + async fn explicit_withdrawal_without_marker_or_ec_cookie_does_not_set_cookie() { let settings = create_test_settings(); let consent = ConsentContext { jurisdiction: Jurisdiction::UsState("CA".to_owned()), @@ -2370,7 +2423,8 @@ mod tests { None, None, &mut response, - ); + ) + .await; assert!( response.headers().get(http::header::SET_COOKIE).is_none(), @@ -2378,8 +2432,8 @@ mod tests { ); } - #[test] - fn explicit_withdrawal_expires_marker_without_ec_cookie() { + #[tokio::test(flavor = "current_thread")] + async fn explicit_withdrawal_expires_marker_without_ec_cookie() { let settings = create_test_settings(); let consent = ConsentContext { jurisdiction: Jurisdiction::UsState("CA".to_owned()), @@ -2401,7 +2455,8 @@ mod tests { None, None, &mut response, - ); + ) + .await; let cookies = response .headers() diff --git a/crates/trusted-server-core/src/ec/identify.rs b/crates/trusted-server-core/src/ec/identify.rs index 94464d1a5..5c755b7a3 100644 --- a/crates/trusted-server-core/src/ec/identify.rs +++ b/crates/trusted-server-core/src/ec/identify.rs @@ -33,7 +33,7 @@ use super::registry::PartnerRegistry; /// /// Panics if response builder produces an invalid status or body, which cannot /// happen with the hardcoded values used here. -pub fn handle_identify( +pub async fn handle_identify( settings: &Settings, kv: &KvIdentityGraph, registry: &PartnerRegistry, @@ -87,7 +87,7 @@ pub fn handle_identify( let mut uid: Option = None; let mut cluster_size: Option = None; - match kv.get(ec_id) { + match kv.get(ec_id).await { Ok(Some((entry, generation))) => { if !entry.consent.ok { // Tombstone entries preserve the withdrawal signal for 24 hours. @@ -104,7 +104,7 @@ pub fn handle_identify( // Evaluate cluster size lazily for identify responses. Existing // stored cluster_size values are reused without a prefix-list call. - match kv.evaluate_cluster(ec_id, &entry, generation) { + match kv.evaluate_cluster(ec_id, &entry, generation).await { Ok(size) => { cluster_size = size; } @@ -462,8 +462,8 @@ mod tests { ); } - #[test] - fn handle_identify_rejects_missing_bearer_token() { + #[tokio::test(flavor = "current_thread")] + async fn handle_identify_rejects_missing_bearer_token() { let settings = create_test_settings(); let kv = KvIdentityGraph::failing("missing_store"); let registry = PartnerRegistry::empty(); @@ -475,6 +475,7 @@ mod tests { let ec_context = make_ec_context(Jurisdiction::NonRegulated, None); let response = handle_identify(&settings, &kv, ®istry, &req, &ec_context) + .await .expect("should construct unauthorized response"); assert_eq!( @@ -502,8 +503,8 @@ mod tests { ); } - #[test] - fn handle_identify_rejects_invalid_bearer_token() { + #[tokio::test(flavor = "current_thread")] + async fn handle_identify_rejects_invalid_bearer_token() { let settings = create_test_settings(); let kv = KvIdentityGraph::failing("missing_store"); let partners = vec![make_test_partner("ssp.example.com", VALID_API_TOKEN)]; @@ -517,6 +518,7 @@ mod tests { let ec_context = make_ec_context(Jurisdiction::NonRegulated, None); let response = handle_identify(&settings, &kv, ®istry, &req, &ec_context) + .await .expect("should construct unauthorized response"); assert_eq!( @@ -527,8 +529,8 @@ mod tests { assert_no_store(&response); } - #[test] - fn handle_identify_denied_consent_returns_403() { + #[tokio::test(flavor = "current_thread")] + async fn handle_identify_denied_consent_returns_403() { let settings = create_test_settings(); let kv = KvIdentityGraph::failing("missing_store"); let partners = vec![make_test_partner("ssp.example.com", VALID_API_TOKEN)]; @@ -542,6 +544,7 @@ mod tests { let ec_context = make_ec_context(Jurisdiction::Unknown, None); let response = handle_identify(&settings, &kv, ®istry, &req, &ec_context) + .await .expect("should construct denied response"); assert_eq!( @@ -561,8 +564,8 @@ mod tests { ); } - #[test] - fn handle_identify_without_ec_returns_204() { + #[tokio::test(flavor = "current_thread")] + async fn handle_identify_without_ec_returns_204() { let settings = create_test_settings(); let kv = KvIdentityGraph::failing("missing_store"); let partners = vec![make_test_partner("ssp.example.com", VALID_API_TOKEN)]; @@ -576,6 +579,7 @@ mod tests { let ec_context = make_ec_context(Jurisdiction::NonRegulated, None); let response = handle_identify(&settings, &kv, ®istry, &req, &ec_context) + .await .expect("should construct no-content response"); assert_eq!( @@ -586,8 +590,8 @@ mod tests { assert_no_store(&response); } - #[test] - fn handle_identify_kv_failure_sets_degraded_true() { + #[tokio::test(flavor = "current_thread")] + async fn handle_identify_kv_failure_sets_degraded_true() { let settings = create_test_settings(); let kv = KvIdentityGraph::failing("missing_store"); let partners = vec![make_test_partner("ssp.example.com", VALID_API_TOKEN)]; @@ -602,6 +606,7 @@ mod tests { let ec_context = make_ec_context(Jurisdiction::NonRegulated, Some(&ec_id)); let response = handle_identify(&settings, &kv, ®istry, &req, &ec_context) + .await .expect("should construct degraded identify response"); assert_eq!( @@ -639,8 +644,8 @@ mod tests { ); } - #[test] - fn handle_identify_denies_mismatched_browser_origin() { + #[tokio::test(flavor = "current_thread")] + async fn handle_identify_denies_mismatched_browser_origin() { let settings = create_test_settings(); let kv = KvIdentityGraph::failing("missing_store"); let partners = vec![make_test_partner("ssp.example.com", VALID_API_TOKEN)]; @@ -655,6 +660,7 @@ mod tests { let ec_context = make_ec_context(Jurisdiction::NonRegulated, None); let response = handle_identify(&settings, &kv, ®istry, &req, &ec_context) + .await .expect("should construct forbidden response"); assert_eq!( @@ -665,8 +671,8 @@ mod tests { assert_no_store(&response); } - #[test] - fn handle_identify_allows_browser_origin_and_reflects_cors_headers() { + #[tokio::test(flavor = "current_thread")] + async fn handle_identify_allows_browser_origin_and_reflects_cors_headers() { let settings = create_test_settings(); let kv = KvIdentityGraph::failing("missing_store"); let partners = vec![make_test_partner("ssp.example.com", VALID_API_TOKEN)]; @@ -681,6 +687,7 @@ mod tests { let ec_context = make_ec_context(Jurisdiction::NonRegulated, None); let response = handle_identify(&settings, &kv, ®istry, &req, &ec_context) + .await .expect("should construct no-content response with CORS headers"); assert_eq!( diff --git a/crates/trusted-server-core/src/ec/kv.rs b/crates/trusted-server-core/src/ec/kv.rs index b0eb83d3f..286fb3c99 100644 --- a/crates/trusted-server-core/src/ec/kv.rs +++ b/crates/trusted-server-core/src/ec/kv.rs @@ -305,11 +305,11 @@ impl KvIdentityGraph { /// # Errors /// /// Returns [`TrustedServerError::KvStore`] on store open or read failure. - pub fn lookup_raw( + pub async fn lookup_raw( &self, ec_id: &str, ) -> Result, Report> { - self.store.lookup(ec_id) + self.store.lookup(ec_id).await } /// Reads the full entry and its generation marker for CAS writes. @@ -319,8 +319,11 @@ impl KvIdentityGraph { /// # Errors /// /// Returns [`TrustedServerError::KvStore`] on store open or read failure. - pub fn get(&self, ec_id: &str) -> Result, Report> { - let Some(lookup) = self.store.lookup(ec_id)? else { + pub async fn get( + &self, + ec_id: &str, + ) -> Result, Report> { + let Some(lookup) = self.store.lookup(ec_id).await? else { return Ok(None); }; @@ -330,8 +333,8 @@ impl KvIdentityGraph { /// Loads one request-scoped snapshot, preserving miss versus failure at the caller boundary. #[must_use] - pub fn load_snapshot(&self, ec_id: &str) -> EcKvSnapshot { - match self.get(ec_id) { + pub async fn load_snapshot(&self, ec_id: &str) -> EcKvSnapshot { + match self.get(ec_id).await { Ok(Some((entry, generation))) => EcKvSnapshot::Present { ec_id: ec_id.to_owned(), entry: Box::new(entry), @@ -383,11 +386,11 @@ impl KvIdentityGraph { /// # Errors /// /// Returns [`TrustedServerError::KvStore`] on store open or read failure. - pub fn get_metadata( + pub async fn get_metadata( &self, ec_id: &str, ) -> Result, Report> { - let Some(lookup) = self.store.lookup(ec_id)? else { + let Some(lookup) = self.store.lookup(ec_id).await? else { return Ok(None); }; @@ -413,9 +416,16 @@ impl KvIdentityGraph { /// /// Returns [`TrustedServerError::KvStore`] on store error or if the /// key already exists. - pub fn create(&self, ec_id: &str, entry: &KvEntry) -> Result<(), Report> { + pub async fn create( + &self, + ec_id: &str, + entry: &KvEntry, + ) -> Result<(), Report> { let (body, meta_str) = Self::serialize_entry(entry, self.store_name())?; - match self.write_entry(ec_id, &body, &meta_str, ENTRY_TTL, EcKvWriteMode::Add)? { + match self + .write_entry(ec_id, &body, &meta_str, ENTRY_TTL, EcKvWriteMode::Add) + .await? + { EcKvWriteOutcome::Written => Ok(()), EcKvWriteOutcome::PreconditionFailed => { Err(self.kv_error(format!("Key '{}' already exists", log_id(ec_id)))) @@ -428,20 +438,23 @@ impl KvIdentityGraph { /// # Errors /// /// Returns [`TrustedServerError::KvStore`] when serialization or store I/O fails. - pub fn create_if_absent( + pub async fn create_if_absent( &self, ec_id: &str, entry: &KvEntry, ) -> Result> { let (body, meta_str) = Self::serialize_entry(entry, self.store_name())?; - match self.write_entry(ec_id, &body, &meta_str, ENTRY_TTL, EcKvWriteMode::Add)? { + match self + .write_entry(ec_id, &body, &meta_str, ENTRY_TTL, EcKvWriteMode::Add) + .await? + { EcKvWriteOutcome::Written => Ok(CreateIfAbsentOutcome::Written), EcKvWriteOutcome::PreconditionFailed => Ok(CreateIfAbsentOutcome::AlreadyExists), } } /// Low-level write with shared error context. - fn write_entry( + async fn write_entry( &self, ec_id: &str, body: &str, @@ -449,15 +462,17 @@ impl KvIdentityGraph { ttl: Duration, mode: EcKvWriteMode, ) -> Result> { - self.store.insert( - ec_id, - EcKvWrite { - body, - metadata: meta_str, - ttl, - mode, - }, - ) + self.store + .insert( + ec_id, + EcKvWrite { + body, + metadata: meta_str, + ttl, + mode, + }, + ) + .await } /// Creates a new entry, or overwrites an existing tombstone on re-consent. @@ -475,7 +490,7 @@ impl KvIdentityGraph { /// /// Returns [`TrustedServerError::KvStore`] on store error or CAS /// exhaustion. - pub fn create_or_revive( + pub async fn create_or_revive( &self, ec_id: &str, entry: &KvEntry, @@ -484,17 +499,19 @@ impl KvIdentityGraph { let (body, meta_str) = Self::serialize_entry(entry, self.store_name())?; // Try create first — fast path for new entries. - if self.write_entry(ec_id, &body, &meta_str, ENTRY_TTL, EcKvWriteMode::Add)? + if self + .write_entry(ec_id, &body, &meta_str, ENTRY_TTL, EcKvWriteMode::Add) + .await? == EcKvWriteOutcome::Written { return Ok(()); } // Key exists — read it to determine if it's live or a tombstone. - let (existing, generation) = match self.get(ec_id)? { + let (existing, generation) = match self.get(ec_id).await? { Some(pair) => pair, // Raced with a delete — try create again. - None => return self.create(ec_id, entry), + None => return self.create(ec_id, entry).await, }; // Live entry — nothing to do. @@ -517,15 +534,18 @@ impl KvIdentityGraph { // A completion marker belongs to the tombstone generation. Remove // it before making this key live so a later withdrawal cannot be // suppressed by stale fallback state. - self.clear_withdrawal_marker(ec_id)?; + self.clear_withdrawal_marker(ec_id).await?; - match self.write_entry( - ec_id, - &body, - &meta_str, - ENTRY_TTL, - EcKvWriteMode::IfGenerationMatch(current_gen), - )? { + match self + .write_entry( + ec_id, + &body, + &meta_str, + ENTRY_TTL, + EcKvWriteMode::IfGenerationMatch(current_gen), + ) + .await? + { EcKvWriteOutcome::Written => return Ok(()), EcKvWriteOutcome::PreconditionFailed => { log::debug!( @@ -535,7 +555,7 @@ impl KvIdentityGraph { ); // Re-read immediately to get a fresh generation. Sleeping in // the CAS loop would block the edge compute request worker. - match self.get(ec_id)? { + match self.get(ec_id).await? { Some((refreshed, generation)) => { if refreshed.consent.ok { // Someone else revived it — done. @@ -543,7 +563,7 @@ impl KvIdentityGraph { } current_gen = generation; } - None => return self.create(ec_id, entry), + None => return self.create(ec_id, entry).await, } } } @@ -562,7 +582,7 @@ impl KvIdentityGraph { /// conflict this method never writes again. A live follow-up becomes the /// authoritative snapshot; a missing or failed follow-up retains the live /// pre-write snapshot as proof and defers the update. - pub(crate) fn sync_eid_cookie_updates_from_snapshot( + pub(crate) async fn sync_eid_cookie_updates_from_snapshot( &self, ec_id: &str, updates: &[PartnerIdUpdate], @@ -589,7 +609,7 @@ impl KvIdentityGraph { EcKvSnapshot::Failed { ec_id: ref snapshot_id, } if snapshot_id == ec_id => return (snapshot, EidCookieSyncOutcome::Failed), - _ => self.load_snapshot(ec_id), + _ => self.load_snapshot(ec_id).await, }; let (mut entry, generation) = match current { @@ -642,13 +662,16 @@ impl KvIdentityGraph { EidCookieSyncOutcome::Failed, ); }; - match self.write_entry( - ec_id, - &body, - &meta_str, - ENTRY_TTL, - EcKvWriteMode::IfGenerationMatch(generation), - ) { + match self + .write_entry( + ec_id, + &body, + &meta_str, + ENTRY_TTL, + EcKvWriteMode::IfGenerationMatch(generation), + ) + .await + { Ok(EcKvWriteOutcome::Written) => ( EcKvSnapshot::Present { ec_id: ec_id.to_owned(), @@ -662,7 +685,7 @@ impl KvIdentityGraph { }, ), Ok(EcKvWriteOutcome::PreconditionFailed) => { - let refreshed = self.load_snapshot(ec_id); + let refreshed = self.load_snapshot(ec_id).await; let Some(refreshed_entry) = refreshed.entry_for(ec_id) else { // The failed CAS proved `current`'s generation stale. Keep // only its existence proof so later writes must reread. @@ -720,7 +743,7 @@ impl KvIdentityGraph { /// later chance to retry. /// /// [`generate_if_needed`]: super::generate_if_needed - pub(crate) fn upsert_partner_ids_from_snapshot( + pub(crate) async fn upsert_partner_ids_from_snapshot( &self, ec_id: &str, updates: &[PartnerIdUpdate], @@ -760,7 +783,7 @@ impl KvIdentityGraph { EcKvSnapshot::Failed { ec_id: ref snapshot_id, } if snapshot_id == ec_id => return snapshot, - _ => self.load_snapshot(ec_id), + _ => self.load_snapshot(ec_id).await, }; for _attempt in 0..MAX_CAS_RETRIES { @@ -797,13 +820,16 @@ impl KvIdentityGraph { ec_id: ec_id.to_owned(), }; }; - match self.write_entry( - ec_id, - &body, - &meta_str, - ENTRY_TTL, - EcKvWriteMode::IfGenerationMatch(generation), - ) { + match self + .write_entry( + ec_id, + &body, + &meta_str, + ENTRY_TTL, + EcKvWriteMode::IfGenerationMatch(generation), + ) + .await + { Ok(EcKvWriteOutcome::Written) => { return EcKvSnapshot::Present { ec_id: ec_id.to_owned(), @@ -812,7 +838,7 @@ impl KvIdentityGraph { }; } Ok(EcKvWriteOutcome::PreconditionFailed) => { - current = self.load_snapshot(ec_id); + current = self.load_snapshot(ec_id).await; } Err(err) => { log::warn!( @@ -876,14 +902,14 @@ impl KvIdentityGraph { /// /// Returns [`TrustedServerError::KvStore`] on store error or CAS /// exhaustion after [`MAX_CAS_RETRIES`] attempts. - pub fn upsert_partner_id( + pub async fn upsert_partner_id( &self, ec_id: &str, partner_id: &str, uid: &str, ) -> Result<(), Report> { for attempt in 0..MAX_CAS_RETRIES { - let (mut entry, generation) = match self.get(ec_id)? { + let (mut entry, generation) = match self.get(ec_id).await? { Some(pair) => pair, None => { log::info!( @@ -927,13 +953,16 @@ impl KvIdentityGraph { let (body, meta_str) = Self::serialize_entry(&entry, self.store_name())?; - match self.write_entry( - ec_id, - &body, - &meta_str, - ENTRY_TTL, - EcKvWriteMode::IfGenerationMatch(generation), - )? { + match self + .write_entry( + ec_id, + &body, + &meta_str, + ENTRY_TTL, + EcKvWriteMode::IfGenerationMatch(generation), + ) + .await? + { EcKvWriteOutcome::Written => return Ok(()), EcKvWriteOutcome::PreconditionFailed => { log::debug!( @@ -966,14 +995,14 @@ impl KvIdentityGraph { /// /// Returns [`TrustedServerError::KvStore`] on store I/O or CAS /// exhaustion errors. - pub fn upsert_partner_id_if_exists( + pub async fn upsert_partner_id_if_exists( &self, ec_id: &str, partner_id: &str, uid: &str, ) -> Result> { for attempt in 0..MAX_CAS_RETRIES { - let (mut entry, generation) = match self.get(ec_id)? { + let (mut entry, generation) = match self.get(ec_id).await? { Some(pair) => pair, None => return Ok(UpsertResult::NotFound), }; @@ -999,13 +1028,16 @@ impl KvIdentityGraph { let (body, meta_str) = Self::serialize_entry(&entry, self.store_name())?; - match self.write_entry( - ec_id, - &body, - &meta_str, - ENTRY_TTL, - EcKvWriteMode::IfGenerationMatch(generation), - )? { + match self + .write_entry( + ec_id, + &body, + &meta_str, + ENTRY_TTL, + EcKvWriteMode::IfGenerationMatch(generation), + ) + .await? + { EcKvWriteOutcome::Written => return Ok(UpsertResult::Written), EcKvWriteOutcome::PreconditionFailed => { log::debug!( @@ -1029,8 +1061,11 @@ impl KvIdentityGraph { /// # Errors /// /// Returns [`TrustedServerError::KvStore`] if existence cannot be confirmed. - pub fn key_exists_confirmed(&self, ec_id: &str) -> Result> { - self.store.key_exists(ec_id) + pub async fn key_exists_confirmed( + &self, + ec_id: &str, + ) -> Result> { + self.store.key_exists(ec_id).await } fn withdrawal_marker_prefix(ec_id: &str) -> String { @@ -1041,17 +1076,19 @@ impl KvIdentityGraph { format!("{}{valid_until}", Self::withdrawal_marker_prefix(ec_id)) } - fn withdrawal_marker_keys( + async fn withdrawal_marker_keys( &self, ec_id: &str, ) -> Result, Report> { - self.store.list_keys_with_prefix( - &Self::withdrawal_marker_prefix(ec_id), - WITHDRAWAL_MARKER_LIST_LIMIT, - ) + self.store + .list_keys_with_prefix( + &Self::withdrawal_marker_prefix(ec_id), + WITHDRAWAL_MARKER_LIST_LIMIT, + ) + .await } - fn withdrawal_marker_exists( + async fn withdrawal_marker_exists( &self, ec_id: &str, now: Option, @@ -1064,35 +1101,40 @@ impl KvIdentityGraph { let marker_prefix = Self::withdrawal_marker_prefix(ec_id); Ok(self .store - .list_keys_with_prefix(&marker_prefix, WITHDRAWAL_MARKER_LIST_LIMIT)? + .list_keys_with_prefix(&marker_prefix, WITHDRAWAL_MARKER_LIST_LIMIT) + .await? .iter() .filter_map(|key| key.strip_prefix(&marker_prefix)) .filter_map(|valid_until| valid_until.parse::().ok()) .any(|valid_until| valid_until > now)) } - fn write_withdrawal_marker( + async fn write_withdrawal_marker( &self, ec_id: &str, tombstone_updated: u64, ) -> Result<(), Report> { let valid_until = tombstone_updated.saturating_add(TOMBSTONE_TTL.as_secs()); let marker_key = Self::withdrawal_marker_key(ec_id, valid_until); - match self.store.insert( - &marker_key, - EcKvWrite { - body: "1", - metadata: "{}", - ttl: TOMBSTONE_TTL, - mode: EcKvWriteMode::Add, - }, - )? { + match self + .store + .insert( + &marker_key, + EcKvWrite { + body: "1", + metadata: "{}", + ttl: TOMBSTONE_TTL, + mode: EcKvWriteMode::Add, + }, + ) + .await? + { EcKvWriteOutcome::Written | EcKvWriteOutcome::PreconditionFailed => Ok(()), } } - fn clear_withdrawal_marker(&self, ec_id: &str) -> Result<(), Report> { - let marker_keys = self.withdrawal_marker_keys(ec_id)?; + async fn clear_withdrawal_marker(&self, ec_id: &str) -> Result<(), Report> { + let marker_keys = self.withdrawal_marker_keys(ec_id).await?; if marker_keys.len() >= WITHDRAWAL_MARKER_LIST_LIMIT as usize { return Err(self.kv_error(format!( "Withdrawal marker cleanup exceeds list budget for '{}'", @@ -1101,8 +1143,8 @@ impl KvIdentityGraph { } for marker_key in marker_keys { - if let Err(delete_err) = self.store.delete(&marker_key) { - match self.store.key_exists(&marker_key) { + if let Err(delete_err) = self.store.delete(&marker_key).await { + match self.store.key_exists(&marker_key).await { // Another request removed the marker first. Ok(false) => {} Ok(true) | Err(_) => return Err(delete_err), @@ -1112,8 +1154,8 @@ impl KvIdentityGraph { Ok(()) } - fn record_withdrawal_completion(&self, ec_id: &str, tombstone_updated: u64) { - if let Err(err) = self.write_withdrawal_marker(ec_id, tombstone_updated) { + async fn record_withdrawal_completion(&self, ec_id: &str, tombstone_updated: u64) { + if let Err(err) = self.write_withdrawal_marker(ec_id, tombstone_updated).await { // The root is already tombstoned. Preserve that successful privacy // write even if the cost-control marker cannot be recorded. log::warn!( @@ -1177,12 +1219,12 @@ impl KvIdentityGraph { /// `error` level and continue: cookie deletion is the primary enforcement /// mechanism. #[cfg(test)] - pub(crate) fn write_withdrawal_tombstone( + pub(crate) async fn write_withdrawal_tombstone( &self, ec_id: &str, record_snapshot: impl FnOnce(EcKvSnapshot), ) -> Result> { - let written = self.tombstone_held_identity(ec_id); + let written = self.tombstone_held_identity(ec_id).await; record_snapshot(match &written { Ok(Some(entry)) => EcKvSnapshot::Present { @@ -1199,7 +1241,8 @@ impl KvIdentityGraph { }); if let Ok(Some(entry)) = &written { - self.record_withdrawal_completion(ec_id, entry.consent.updated); + self.record_withdrawal_completion(ec_id, entry.consent.updated) + .await; } written.map(|entry| { if entry.is_some() { @@ -1215,7 +1258,7 @@ impl KvIdentityGraph { /// `Ok(None)` means the store does not hold the identity, so nothing was /// written. #[cfg(test)] - fn tombstone_held_identity( + async fn tombstone_held_identity( &self, ec_id: &str, ) -> Result, Report> { @@ -1223,14 +1266,14 @@ impl KvIdentityGraph { // would restore the unconditional write whenever the store can be made // to fail, and an extra `Ok` variant would be discarded in silence by a // caller that only inspects the error case. - if !self.key_exists_confirmed(ec_id)? { + if !self.key_exists_confirmed(ec_id).await? { return Ok(None); } - self.overwrite_withdrawal_tombstone(ec_id).map(Some) + self.overwrite_withdrawal_tombstone(ec_id).await.map(Some) } - fn overwrite_withdrawal_tombstone( + async fn overwrite_withdrawal_tombstone( &self, ec_id: &str, ) -> Result> { @@ -1244,6 +1287,7 @@ impl KvIdentityGraph { TOMBSTONE_TTL, EcKvWriteMode::Overwrite, ) + .await .map(|_| entry) .map_err(|report| { report.change_context(TrustedServerError::KvStore { @@ -1262,13 +1306,13 @@ impl KvIdentityGraph { /// generation is available after a missed read. Marker-check failure falls /// back to that privacy write, as does an unusable clock; root-existence /// failure leaves withdrawal unresolved rather than silently dropped. - fn tombstone_unproven_missing( + async fn tombstone_unproven_missing( &self, ec_id: &str, missing: EcKvSnapshot, now: Option, ) -> EcKvSnapshot { - match self.withdrawal_marker_exists(ec_id, now) { + match self.withdrawal_marker_exists(ec_id, now).await { Ok(true) => { log::debug!( "withdrawal tombstone for '{}': completion marker already exists", @@ -1287,7 +1331,7 @@ impl KvIdentityGraph { } } - match self.key_exists_confirmed(ec_id) { + match self.key_exists_confirmed(ec_id).await { Ok(false) => missing, Ok(true) => { log::warn!( @@ -1295,9 +1339,10 @@ impl KvIdentityGraph { lists; writing an unconditional tombstone", log_id(ec_id) ); - match self.overwrite_withdrawal_tombstone(ec_id) { + match self.overwrite_withdrawal_tombstone(ec_id).await { Ok(tombstone) => { - self.record_withdrawal_completion(ec_id, tombstone.consent.updated); + self.record_withdrawal_completion(ec_id, tombstone.consent.updated) + .await; EcKvSnapshot::Present { ec_id: ec_id.to_owned(), entry: Box::new(tombstone), @@ -1346,7 +1391,7 @@ impl KvIdentityGraph { /// /// Resolving the initial snapshot happens outside the retry counter, so all /// [`MAX_CAS_RETRIES`] iterations stay available for the tombstone write. - pub(crate) fn tombstone_existing_from_snapshot( + pub(crate) async fn tombstone_existing_from_snapshot( &self, ec_id: &str, snapshot: EcKvSnapshot, @@ -1362,7 +1407,7 @@ impl KvIdentityGraph { generation: Some(_), .. } if snapshot_id == ec_id => snapshot, - _ => self.load_snapshot(ec_id), + _ => self.load_snapshot(ec_id).await, }; for _attempt in 0..MAX_CAS_RETRIES { @@ -1383,11 +1428,9 @@ impl KvIdentityGraph { EcKvSnapshot::Missing { ec_id: ref snapshot_id, } if snapshot_id == ec_id => { - return self.tombstone_unproven_missing( - ec_id, - current, - checked_current_timestamp(), - ); + return self + .tombstone_unproven_missing(ec_id, current, checked_current_timestamp()) + .await; } // A refreshed read that failed (or any other unusable state) // fails closed rather than silently dropping the withdrawal. @@ -1404,15 +1447,19 @@ impl KvIdentityGraph { ec_id: ec_id.to_owned(), }; }; - match self.write_entry( - ec_id, - &body, - &meta_str, - TOMBSTONE_TTL, - EcKvWriteMode::IfGenerationMatch(generation), - ) { + match self + .write_entry( + ec_id, + &body, + &meta_str, + TOMBSTONE_TTL, + EcKvWriteMode::IfGenerationMatch(generation), + ) + .await + { Ok(EcKvWriteOutcome::Written) => { - self.record_withdrawal_completion(ec_id, tombstone.consent.updated); + self.record_withdrawal_completion(ec_id, tombstone.consent.updated) + .await; return EcKvSnapshot::Present { ec_id: ec_id.to_owned(), entry: Box::new(tombstone), @@ -1420,7 +1467,7 @@ impl KvIdentityGraph { }; } Ok(EcKvWriteOutcome::PreconditionFailed) => { - current = self.load_snapshot(ec_id); + current = self.load_snapshot(ec_id).await; } Err(err) => { log::warn!( @@ -1457,7 +1504,7 @@ impl KvIdentityGraph { /// # Errors /// /// Returns [`TrustedServerError::KvStore`] on store error. - pub fn count_hash_prefix_keys( + pub async fn count_hash_prefix_keys( &self, hash_prefix: &str, ) -> Result> { @@ -1466,6 +1513,7 @@ impl KvIdentityGraph { // store context to list failures, so propagate without re-wrapping. self.store .count_keys_with_prefix(hash_prefix, CLUSTER_LIST_LIMIT) + .await } /// Evaluates the cluster size for an EC entry. @@ -1481,7 +1529,7 @@ impl KvIdentityGraph { /// # Errors /// /// Returns [`TrustedServerError::KvStore`] on store or list failure. - pub fn evaluate_cluster( + pub async fn evaluate_cluster( &self, ec_id: &str, entry: &KvEntry, @@ -1503,7 +1551,7 @@ impl KvIdentityGraph { // Compute cluster size via prefix list. let hash_prefix = ec_hash(ec_id); - let cluster_size = self.count_hash_prefix_keys(hash_prefix)?; + let cluster_size = self.count_hash_prefix_keys(hash_prefix).await?; log::debug!( "evaluate_cluster: computed cluster_size={cluster_size} for '{}'", @@ -1521,13 +1569,16 @@ impl KvIdentityGraph { let (body, meta_str) = Self::serialize_entry(&updated_entry, self.store_name())?; - match self.write_entry( - ec_id, - &body, - &meta_str, - ENTRY_TTL, - EcKvWriteMode::IfGenerationMatch(generation), - ) { + match self + .write_entry( + ec_id, + &body, + &meta_str, + ENTRY_TTL, + EcKvWriteMode::IfGenerationMatch(generation), + ) + .await + { Ok(EcKvWriteOutcome::Written) => {} Ok(EcKvWriteOutcome::PreconditionFailed) => { log::debug!( @@ -1556,11 +1607,11 @@ impl KvIdentityGraph { /// # Errors /// /// Returns [`TrustedServerError::KvStore`] on store error. - pub fn delete(&self, ec_id: &str) -> Result<(), Report> { + pub async fn delete(&self, ec_id: &str) -> Result<(), Report> { // The backend's delete already attaches store context, so propagate // without re-wrapping the same message. - self.clear_withdrawal_marker(ec_id)?; - self.store.delete(ec_id) + self.clear_withdrawal_marker(ec_id).await?; + self.store.delete(ec_id).await } } @@ -1636,7 +1687,7 @@ mod tests { } } - fn seed_live(&self, ec_id: &str) { + async fn seed_live(&self, ec_id: &str) { let (body, meta) = KvIdentityGraph::serialize_entry(&live_entry(), self.inner.store_name()) .expect("should serialize seeded entry"); @@ -1650,55 +1701,146 @@ mod tests { mode: EcKvWriteMode::Add, }, ) + .await .expect("should seed live entry"); } } + #[async_trait::async_trait(?Send)] impl EcKvStore for DisappearOnConflictEcKv { fn store_name(&self) -> &str { self.inner.store_name() } - fn lookup(&self, key: &str) -> Result, Report> { - self.inner.lookup(key) + async fn lookup( + &self, + key: &str, + ) -> Result, Report> { + self.inner.lookup(key).await } - fn key_exists(&self, key: &str) -> Result> { - self.inner.key_exists(key) + async fn key_exists(&self, key: &str) -> Result> { + self.inner.key_exists(key).await } - fn insert( + async fn insert( &self, key: &str, write: EcKvWrite<'_>, ) -> Result> { if matches!(write.mode, EcKvWriteMode::IfGenerationMatch(_)) { - let mut remaining = self - .conflicts_remaining - .lock() - .expect("should lock conflict counter"); - if *remaining > 0 { - *remaining -= 1; - self.inner.delete(key).expect("should delete on conflict"); + let inject_conflict = { + let mut remaining = self + .conflicts_remaining + .lock() + .expect("should lock conflict counter"); + let inject_conflict = *remaining > 0; + if inject_conflict { + *remaining -= 1; + } + inject_conflict + }; + if inject_conflict { + self.inner + .delete(key) + .await + .expect("should delete on conflict"); return Ok(EcKvWriteOutcome::PreconditionFailed); } } - self.inner.insert(key, write) + self.inner.insert(key, write).await + } + + async fn list_keys_with_prefix( + &self, + prefix: &str, + limit: u32, + ) -> Result, Report> { + self.inner.list_keys_with_prefix(prefix, limit).await + } + + async fn delete(&self, key: &str) -> Result<(), Report> { + self.inner.delete(key).await + } + } + + struct DelayedLocalEcKv { + inner: InMemoryEcKv, + lookup_started: std::rc::Rc>, + release: std::cell::RefCell>>, + } + + #[async_trait::async_trait(?Send)] + impl EcKvStore for DelayedLocalEcKv { + fn store_name(&self) -> &str { + self.inner.store_name() + } + + async fn lookup( + &self, + key: &str, + ) -> Result, Report> { + self.lookup_started.set(true); + let release = self.release.borrow_mut().take().expect("should read once"); + release + .await + .expect("should release lookup from other work"); + self.inner.lookup(key).await } - fn list_keys_with_prefix( + async fn key_exists(&self, key: &str) -> Result> { + self.inner.key_exists(key).await + } + + async fn insert( + &self, + key: &str, + write: EcKvWrite<'_>, + ) -> Result> { + self.inner.insert(key, write).await + } + + async fn list_keys_with_prefix( &self, prefix: &str, limit: u32, ) -> Result, Report> { - self.inner.list_keys_with_prefix(prefix, limit) + self.inner.list_keys_with_prefix(prefix, limit).await } - fn delete(&self, key: &str) -> Result<(), Report> { - self.inner.delete(key) + async fn delete(&self, key: &str) -> Result<(), Report> { + self.inner.delete(key).await } } + #[tokio::test(flavor = "current_thread")] + async fn graph_lookup_yields_to_other_work_with_a_non_send_store() { + let lookup_started = std::rc::Rc::new(std::cell::Cell::new(false)); + let (release, waiting) = tokio::sync::oneshot::channel(); + let graph = KvIdentityGraph::new(DelayedLocalEcKv { + inner: InMemoryEcKv::new("local-store"), + lookup_started: lookup_started.clone(), + release: std::cell::RefCell::new(Some(waiting)), + }); + let ec_id = snapshot_ec_id(); + graph + .create(&ec_id, &live_entry()) + .await + .expect("should seed row"); + + let (entry, ()) = tokio::join!(biased; graph.get(&ec_id), async { + assert!( + lookup_started.get(), + "lookup should suspend before other work" + ); + release.send(()).expect("should keep lookup pending"); + }); + assert!( + entry.expect("should read row").is_some(), + "should read the seeded row after yielding" + ); + } + fn snapshot_ec_id() -> String { format!("{}.ABC123", "a".repeat(64)) } @@ -1837,7 +1979,7 @@ mod tests { } } - fn seed_tombstone(&self, ec_id: &str) { + async fn seed_tombstone(&self, ec_id: &str) { let (body, meta) = KvIdentityGraph::serialize_entry( &KvEntry::tombstone(1000), self.inner.store_name(), @@ -1853,10 +1995,11 @@ mod tests { mode: EcKvWriteMode::Add, }, ) + .await .expect("should seed tombstone"); } - fn seed_live(&self, ec_id: &str) { + async fn seed_live(&self, ec_id: &str) { let (body, meta) = KvIdentityGraph::serialize_entry(&live_entry(), self.inner.store_name()) .expect("should serialize live entry"); @@ -1870,35 +2013,46 @@ mod tests { mode: EcKvWriteMode::Add, }, ) + .await .expect("should seed live entry"); } } + #[async_trait::async_trait(?Send)] impl EcKvStore for ConflictInjectingEcKv { fn store_name(&self) -> &str { self.inner.store_name() } - fn lookup(&self, key: &str) -> Result, Report> { - self.inner.lookup(key) + async fn lookup( + &self, + key: &str, + ) -> Result, Report> { + self.inner.lookup(key).await } - fn key_exists(&self, key: &str) -> Result> { - self.inner.key_exists(key) + async fn key_exists(&self, key: &str) -> Result> { + self.inner.key_exists(key).await } - fn insert( + async fn insert( &self, key: &str, write: EcKvWrite<'_>, ) -> Result> { if matches!(write.mode, EcKvWriteMode::IfGenerationMatch(_)) { - let mut remaining = self - .conflicts_remaining - .lock() - .expect("should lock conflict counter"); - if *remaining > 0 { - *remaining -= 1; + let inject_conflict = { + let mut remaining = self + .conflicts_remaining + .lock() + .expect("should lock conflict counter"); + let inject_conflict = *remaining > 0; + if inject_conflict { + *remaining -= 1; + } + inject_conflict + }; + if inject_conflict { if self.revive_on_conflict { // Simulate a concurrent writer reviving the entry // between this writer's read and its CAS write. @@ -1922,24 +2076,25 @@ mod tests { mode: EcKvWriteMode::Overwrite, }, ) + .await .expect("should apply concurrent revive"); } return Ok(EcKvWriteOutcome::PreconditionFailed); } } - self.inner.insert(key, write) + self.inner.insert(key, write).await } - fn list_keys_with_prefix( + async fn list_keys_with_prefix( &self, prefix: &str, limit: u32, ) -> Result, Report> { - self.inner.list_keys_with_prefix(prefix, limit) + self.inner.list_keys_with_prefix(prefix, limit).await } - fn delete(&self, key: &str) -> Result<(), Report> { - self.inner.delete(key) + async fn delete(&self, key: &str) -> Result<(), Report> { + self.inner.delete(key).await } } @@ -1975,7 +2130,7 @@ mod tests { self } - fn seed_live(&self, ec_id: &str) { + async fn seed_live(&self, ec_id: &str) { let (body, meta) = KvIdentityGraph::serialize_entry(&live_entry(), self.inner.store_name()) .expect("should serialize initial entry"); @@ -1989,16 +2144,21 @@ mod tests { mode: EcKvWriteMode::Add, }, ) + .await .expect("should seed initial entry"); } } + #[async_trait::async_trait(?Send)] impl EcKvStore for EidConflictEcKv { fn store_name(&self) -> &str { self.inner.store_name() } - fn lookup(&self, key: &str) -> Result, Report> { + async fn lookup( + &self, + key: &str, + ) -> Result, Report> { self.lookups .fetch_add(1, std::sync::atomic::Ordering::Relaxed); if self @@ -2007,14 +2167,14 @@ mod tests { { return Ok(None); } - self.inner.lookup(key) + self.inner.lookup(key).await } - fn key_exists(&self, key: &str) -> Result> { - self.inner.key_exists(key) + async fn key_exists(&self, key: &str) -> Result> { + self.inner.key_exists(key).await } - fn insert( + async fn insert( &self, key: &str, write: EcKvWrite<'_>, @@ -2037,6 +2197,7 @@ mod tests { mode: EcKvWriteMode::Overwrite, }, ) + .await .expect("should write concurrent entry"); if self.follow_up_miss { self.miss_next_lookup @@ -2044,42 +2205,44 @@ mod tests { } return Ok(EcKvWriteOutcome::PreconditionFailed); } - self.inner.insert(key, write) + self.inner.insert(key, write).await } - fn list_keys_with_prefix( + async fn list_keys_with_prefix( &self, prefix: &str, limit: u32, ) -> Result, Report> { - self.inner.list_keys_with_prefix(prefix, limit) + self.inner.list_keys_with_prefix(prefix, limit).await } - fn count_keys_with_prefix( + async fn count_keys_with_prefix( &self, prefix: &str, limit: u32, ) -> Result> { - self.inner.count_keys_with_prefix(prefix, limit) + self.inner.count_keys_with_prefix(prefix, limit).await } - fn delete(&self, key: &str) -> Result<(), Report> { - self.inner.delete(key) + async fn delete(&self, key: &str) -> Result<(), Report> { + self.inner.delete(key).await } } - #[test] - fn create_or_revive_retries_cas_conflict_and_succeeds() { + #[tokio::test(flavor = "current_thread")] + async fn create_or_revive_retries_cas_conflict_and_succeeds() { let store = ConflictInjectingEcKv::new(2, false); - store.seed_tombstone("ec-1"); + store.seed_tombstone("ec-1").await; let graph = KvIdentityGraph::new(store); graph .create_or_revive("ec-1", &live_entry()) + .await .expect("should revive after re-reading a fresh generation"); let (entry, _) = graph .get("ec-1") + .await .expect("should read entry") .expect("entry should exist"); assert!( @@ -2088,27 +2251,29 @@ mod tests { ); } - #[test] - fn create_or_revive_short_circuits_on_concurrent_revive() { + #[tokio::test(flavor = "current_thread")] + async fn create_or_revive_short_circuits_on_concurrent_revive() { // Inject more conflicts than MAX_CAS_RETRIES so the only way the call // can succeed is the concurrent-revive short-circuit on re-read. let store = ConflictInjectingEcKv::new(MAX_CAS_RETRIES + 1, true); - store.seed_tombstone("ec-2"); + store.seed_tombstone("ec-2").await; let graph = KvIdentityGraph::new(store); graph .create_or_revive("ec-2", &live_entry()) + .await .expect("should return Ok when a concurrent writer already revived the entry"); } - #[test] - fn create_or_revive_errors_after_cas_exhaustion() { + #[tokio::test(flavor = "current_thread")] + async fn create_or_revive_errors_after_cas_exhaustion() { let store = ConflictInjectingEcKv::new(MAX_CAS_RETRIES + 1, false); - store.seed_tombstone("ec-3"); + store.seed_tombstone("ec-3").await; let graph = KvIdentityGraph::new(store); let err = graph .create_or_revive("ec-3", &live_entry()) + .await .expect_err("should fail after exhausting CAS retries"); assert!( format!("{err}").contains("CAS conflict after"), @@ -2209,8 +2374,8 @@ mod tests { assert_eq!(entry.ids["ssp_x"].uid, "original"); } - #[test] - fn eid_cookie_sync_conflict_matching_value_stops_after_one_write() { + #[tokio::test(flavor = "current_thread")] + async fn eid_cookie_sync_conflict_matching_value_stops_after_one_write() { let ec_id = snapshot_ec_id(); let mut concurrent = live_entry(); concurrent.ids.insert( @@ -2226,15 +2391,17 @@ mod tests { std::sync::Arc::clone(&lookups), std::sync::Arc::clone(&writes), ); - store.seed_live(&ec_id); + store.seed_live(&ec_id).await; let graph = KvIdentityGraph::new(store); - let snapshot = graph.load_snapshot(&ec_id); + let snapshot = graph.load_snapshot(&ec_id).await; - let (snapshot, outcome) = graph.sync_eid_cookie_updates_from_snapshot( - &ec_id, - &[PartnerIdUpdate::new("ssp_x", "desired-uid")], - snapshot, - ); + let (snapshot, outcome) = graph + .sync_eid_cookie_updates_from_snapshot( + &ec_id, + &[PartnerIdUpdate::new("ssp_x", "desired-uid")], + snapshot, + ) + .await; assert_eq!(outcome, EidCookieSyncOutcome::ConflictMatched); assert_eq!( @@ -2257,8 +2424,8 @@ mod tests { ); } - #[test] - fn eid_cookie_sync_keeps_live_proof_when_conflict_follow_up_misses() { + #[tokio::test(flavor = "current_thread")] + async fn eid_cookie_sync_keeps_live_proof_when_conflict_follow_up_misses() { let ec_id = snapshot_ec_id(); let lookups = std::sync::Arc::new(std::sync::atomic::AtomicUsize::new(0)); let writes = std::sync::Arc::new(std::sync::atomic::AtomicUsize::new(0)); @@ -2268,16 +2435,18 @@ mod tests { std::sync::Arc::clone(&writes), ) .with_follow_up_miss(); - store.seed_live(&ec_id); + store.seed_live(&ec_id).await; let graph = KvIdentityGraph::new(store); - let (snapshot, outcome) = graph.sync_eid_cookie_updates_from_snapshot( - &ec_id, - &[PartnerIdUpdate::new("ssp_x", "desired-uid")], - EcKvSnapshot::Missing { - ec_id: ec_id.clone(), - }, - ); + let (snapshot, outcome) = graph + .sync_eid_cookie_updates_from_snapshot( + &ec_id, + &[PartnerIdUpdate::new("ssp_x", "desired-uid")], + EcKvSnapshot::Missing { + ec_id: ec_id.clone(), + }, + ) + .await; assert_eq!(outcome, EidCookieSyncOutcome::DeferredConflict); assert!( @@ -2301,8 +2470,8 @@ mod tests { ); } - #[test] - fn eid_cookie_sync_reports_proven_refresh_miss_as_deferred() { + #[tokio::test(flavor = "current_thread")] + async fn eid_cookie_sync_reports_proven_refresh_miss_as_deferred() { let ec_id = snapshot_ec_id(); let graph = KvIdentityGraph::in_memory("empty-store"); let proven = EcKvSnapshot::Present { @@ -2311,11 +2480,13 @@ mod tests { generation: None, }; - let (snapshot, outcome) = graph.sync_eid_cookie_updates_from_snapshot( - &ec_id, - &[PartnerIdUpdate::new("ssp_x", "desired-uid")], - proven, - ); + let (snapshot, outcome) = graph + .sync_eid_cookie_updates_from_snapshot( + &ec_id, + &[PartnerIdUpdate::new("ssp_x", "desired-uid")], + proven, + ) + .await; assert!( snapshot.entry_for(&ec_id).is_some(), @@ -2348,8 +2519,8 @@ mod tests { ); } - #[test] - fn eid_cookie_sync_conflicting_value_defers_after_one_write() { + #[tokio::test(flavor = "current_thread")] + async fn eid_cookie_sync_conflicting_value_defers_after_one_write() { let ec_id = snapshot_ec_id(); let mut concurrent = live_entry(); concurrent.ids.insert( @@ -2365,15 +2536,17 @@ mod tests { std::sync::Arc::clone(&lookups), std::sync::Arc::clone(&writes), ); - store.seed_live(&ec_id); + store.seed_live(&ec_id).await; let graph = KvIdentityGraph::new(store); - let snapshot = graph.load_snapshot(&ec_id); + let snapshot = graph.load_snapshot(&ec_id).await; - let (_, outcome) = graph.sync_eid_cookie_updates_from_snapshot( - &ec_id, - &[PartnerIdUpdate::new("ssp_x", "stale-uid")], - snapshot, - ); + let (_, outcome) = graph + .sync_eid_cookie_updates_from_snapshot( + &ec_id, + &[PartnerIdUpdate::new("ssp_x", "stale-uid")], + snapshot, + ) + .await; assert_eq!(outcome, EidCookieSyncOutcome::DeferredConflict); assert_eq!( @@ -2388,13 +2561,14 @@ mod tests { ); let (stored, _) = graph .get(&ec_id) + .await .expect("should read concurrent row") .expect("row should remain"); assert_eq!(stored.ids["ssp_x"].uid, "concurrent-uid"); } - #[test] - fn eid_cookie_sync_preserves_concurrent_withdrawal_after_conflict() { + #[tokio::test(flavor = "current_thread")] + async fn eid_cookie_sync_preserves_concurrent_withdrawal_after_conflict() { let ec_id = snapshot_ec_id(); let lookups = std::sync::Arc::new(std::sync::atomic::AtomicUsize::new(0)); let writes = std::sync::Arc::new(std::sync::atomic::AtomicUsize::new(0)); @@ -2403,15 +2577,17 @@ mod tests { std::sync::Arc::clone(&lookups), std::sync::Arc::clone(&writes), ); - store.seed_live(&ec_id); + store.seed_live(&ec_id).await; let graph = KvIdentityGraph::new(store); - let snapshot = graph.load_snapshot(&ec_id); + let snapshot = graph.load_snapshot(&ec_id).await; - let (snapshot, outcome) = graph.sync_eid_cookie_updates_from_snapshot( - &ec_id, - &[PartnerIdUpdate::new("ssp_x", "desired-uid")], - snapshot, - ); + let (snapshot, outcome) = graph + .sync_eid_cookie_updates_from_snapshot( + &ec_id, + &[PartnerIdUpdate::new("ssp_x", "desired-uid")], + snapshot, + ) + .await; assert_eq!(outcome, EidCookieSyncOutcome::ConsentWithdrawn); assert!( @@ -2432,8 +2608,8 @@ mod tests { ); } - #[test] - fn eid_cookie_sync_defers_different_values_without_value_owned_freshness() { + #[tokio::test(flavor = "current_thread")] + async fn eid_cookie_sync_defers_different_values_without_value_owned_freshness() { let ec_id = snapshot_ec_id(); let graph = KvIdentityGraph::in_memory("freshness-store"); let mut entry = live_entry(); @@ -2443,24 +2619,30 @@ mod tests { uid: "protected-uid".to_owned(), }, ); - graph.create(&ec_id, &entry).expect("should seed entry"); + graph + .create(&ec_id, &entry) + .await + .expect("should seed entry"); - let (_, outcome) = graph.sync_eid_cookie_updates_from_snapshot( - &ec_id, - &[PartnerIdUpdate::new("ssp_x", "incoming-uid")], - graph.load_snapshot(&ec_id), - ); + let (_, outcome) = graph + .sync_eid_cookie_updates_from_snapshot( + &ec_id, + &[PartnerIdUpdate::new("ssp_x", "incoming-uid")], + graph.load_snapshot(&ec_id).await, + ) + .await; assert_eq!(outcome, EidCookieSyncOutcome::DeferredFreshness); let (stored, _) = graph .get(&ec_id) + .await .expect("should read protected row") .expect("row should remain"); assert_eq!(stored.ids["ssp_x"].uid, "protected-uid"); } - #[test] - fn eid_cookie_sync_adds_missing_ids_while_deferring_different_values() { + #[tokio::test(flavor = "current_thread")] + async fn eid_cookie_sync_adds_missing_ids_while_deferring_different_values() { let ec_id = snapshot_ec_id(); let graph = KvIdentityGraph::in_memory("mixed-freshness-store"); let mut entry = live_entry(); @@ -2470,28 +2652,34 @@ mod tests { uid: "protected-uid".to_owned(), }, ); - graph.create(&ec_id, &entry).expect("should seed entry"); + graph + .create(&ec_id, &entry) + .await + .expect("should seed entry"); - let (_, outcome) = graph.sync_eid_cookie_updates_from_snapshot( - &ec_id, - &[ - PartnerIdUpdate::new("ssp_x", "different-uid"), - PartnerIdUpdate::new("ssp_y", "new-uid"), - ], - graph.load_snapshot(&ec_id), - ); + let (_, outcome) = graph + .sync_eid_cookie_updates_from_snapshot( + &ec_id, + &[ + PartnerIdUpdate::new("ssp_x", "different-uid"), + PartnerIdUpdate::new("ssp_y", "new-uid"), + ], + graph.load_snapshot(&ec_id).await, + ) + .await; assert_eq!(outcome, EidCookieSyncOutcome::WrittenWithDeferredFreshness); let (stored, _) = graph .get(&ec_id) + .await .expect("should read updated row") .expect("row should remain"); assert_eq!(stored.ids["ssp_x"].uid, "protected-uid"); assert_eq!(stored.ids["ssp_y"].uid, "new-uid"); } - #[test] - fn evaluate_cluster_returns_stored_value_without_store_io() { + #[tokio::test(flavor = "current_thread")] + async fn evaluate_cluster_returns_stored_value_without_store_io() { let kv = KvIdentityGraph::failing("nonexistent_store_for_cluster_cache_test"); let ec_id = format!("{}.ABC123", "a".repeat(64)); let mut entry = live_entry(); @@ -2501,6 +2689,7 @@ mod tests { let cluster_size = kv .evaluate_cluster(&ec_id, &entry, 0) + .await .expect("should not touch store when cluster_size is already known"); assert_eq!( @@ -2510,14 +2699,15 @@ mod tests { ); } - #[test] - fn evaluate_cluster_skips_tombstone_without_store_io() { + #[tokio::test(flavor = "current_thread")] + async fn evaluate_cluster_skips_tombstone_without_store_io() { let kv = KvIdentityGraph::failing("nonexistent_store_for_tombstone_cluster_test"); let ec_id = format!("{}.ABC123", "a".repeat(64)); let entry = KvEntry::tombstone(1000); let cluster_size = kv .evaluate_cluster(&ec_id, &entry, 0) + .await .expect("should not touch store for tombstone entries"); assert_eq!( @@ -2526,15 +2716,18 @@ mod tests { ); } - #[test] - fn create_then_get_roundtrips_entry() { + #[tokio::test(flavor = "current_thread")] + async fn create_then_get_roundtrips_entry() { let kv = KvIdentityGraph::in_memory("test_store"); let ec_id = format!("{}.ABC123", "a".repeat(64)); let entry = live_entry(); - kv.create(&ec_id, &entry).expect("should create new entry"); + kv.create(&ec_id, &entry) + .await + .expect("should create new entry"); let (loaded, generation) = kv .get(&ec_id) + .await .expect("should read entry back") .expect("should find created entry"); @@ -2542,15 +2735,18 @@ mod tests { assert!(generation > 0, "should expose a generation marker"); } - #[test] - fn create_rejects_existing_key() { + #[tokio::test(flavor = "current_thread")] + async fn create_rejects_existing_key() { let kv = KvIdentityGraph::in_memory("test_store"); let ec_id = format!("{}.ABC123", "a".repeat(64)); let entry = live_entry(); - kv.create(&ec_id, &entry).expect("should create new entry"); + kv.create(&ec_id, &entry) + .await + .expect("should create new entry"); let err = kv .create(&ec_id, &entry) + .await .expect_err("should reject duplicate create"); assert!( format!("{err}").contains("already exists"), @@ -2558,53 +2754,58 @@ mod tests { ); } - #[test] - fn create_if_absent_reports_written_and_collision() { + #[tokio::test(flavor = "current_thread")] + async fn create_if_absent_reports_written_and_collision() { let kv = KvIdentityGraph::in_memory("test_store"); let ec_id = format!("{}.ABC123", "a".repeat(64)); assert_eq!( kv.create_if_absent(&ec_id, &live_entry()) + .await .expect("should create absent entry"), CreateIfAbsentOutcome::Written ); assert_eq!( kv.create_if_absent(&ec_id, &live_entry()) + .await .expect("should report collision"), CreateIfAbsentOutcome::AlreadyExists ); } - #[test] - fn create_if_absent_propagates_store_error() { + #[tokio::test(flavor = "current_thread")] + async fn create_if_absent_propagates_store_error() { let kv = KvIdentityGraph::failing("test_store"); let ec_id = format!("{}.ABC123", "a".repeat(64)); assert!( - kv.create_if_absent(&ec_id, &live_entry()).is_err(), + kv.create_if_absent(&ec_id, &live_entry()).await.is_err(), "should preserve store failures instead of reporting a collision" ); } - #[test] - fn create_or_revive_revives_tombstone() { + #[tokio::test(flavor = "current_thread")] + async fn create_or_revive_revives_tombstone() { let kv = KvIdentityGraph::in_memory("test_store"); let ec_id = format!("{}.ABC123", "a".repeat(64)); kv.create(&ec_id, &KvEntry::tombstone(1000)) + .await .expect("should create tombstone"); kv.create_or_revive(&ec_id, &live_entry()) + .await .expect("should revive tombstone"); let (loaded, _) = kv .get(&ec_id) + .await .expect("should read entry back") .expect("should find revived entry"); assert!(loaded.consent.ok, "should be live after revive"); } - #[test] - fn create_or_revive_fresh_entry_does_not_access_marker_store() { + #[tokio::test(flavor = "current_thread")] + async fn create_or_revive_fresh_entry_does_not_access_marker_store() { let operations = Arc::new(RecordedEcKvOperations::default()); let kv = KvIdentityGraph::new(RecordingEcKv::with_marker_failures( Arc::clone(&operations), @@ -2613,6 +2814,7 @@ mod tests { let ec_id = format!("{}.ABC123", "a".repeat(64)); kv.create_or_revive(&ec_id, &live_entry()) + .await .expect("should create without reading withdrawal markers"); assert_eq!( @@ -2627,43 +2829,49 @@ mod tests { ); assert!( kv.get(&ec_id) + .await .expect("should read entry") .is_some_and(|(entry, _)| entry.consent.ok), "fresh create should persist a live entry" ); } - #[test] - fn create_or_revive_clears_withdrawal_marker() { + #[tokio::test(flavor = "current_thread")] + async fn create_or_revive_clears_withdrawal_marker() { let kv = KvIdentityGraph::in_memory("test_store"); let ec_id = format!("{}.ABC123", "a".repeat(64)); kv.create(&ec_id, &live_entry()) + .await .expect("should create live entry"); - let snapshot = kv.load_snapshot(&ec_id); - kv.tombstone_existing_from_snapshot(&ec_id, snapshot); + let snapshot = kv.load_snapshot(&ec_id).await; + kv.tombstone_existing_from_snapshot(&ec_id, snapshot).await; assert!( kv.withdrawal_marker_exists(&ec_id, checked_current_timestamp()) + .await .expect("should read withdrawal marker"), "withdrawal should record completion" ); kv.create_or_revive(&ec_id, &live_entry()) + .await .expect("should revive tombstone"); let (loaded, _) = kv .get(&ec_id) + .await .expect("should read revived entry") .expect("should find revived entry"); assert!(loaded.consent.ok, "should be live after revive"); assert!( !kv.withdrawal_marker_exists(&ec_id, checked_current_timestamp()) + .await .expect("should read withdrawal marker"), "revival should clear stale withdrawal completion" ); } - #[test] - fn stale_completion_marker_does_not_suppress_withdrawal_after_root_recreation() { + #[tokio::test(flavor = "current_thread")] + async fn stale_completion_marker_does_not_suppress_withdrawal_after_root_recreation() { let operations = Arc::new(RecordedEcKvOperations::default()); let graph = KvIdentityGraph::new(RecordingEcKv::with_stale_lookups( Arc::clone(&operations), @@ -2674,28 +2882,37 @@ mod tests { let expired_tombstone = KvEntry::tombstone(expired_updated); graph .create(&ec_id, &expired_tombstone) + .await .expect("should seed old tombstone"); graph .write_withdrawal_marker(&ec_id, expired_updated) + .await .expect("should seed old completion marker"); // Model the root expiring just before its later-written completion // marker, followed by recreation of the same key. - graph.store.delete(&ec_id).expect("should expire root row"); + graph + .store + .delete(&ec_id) + .await + .expect("should expire root row"); assert_eq!( graph .create_if_absent(&ec_id, &live_entry()) + .await .expect("should recreate expired root"), CreateIfAbsentOutcome::Written, "should make the same key live again after root expiry" ); - let outcome = graph.tombstone_existing_from_snapshot( - &ec_id, - EcKvSnapshot::Missing { - ec_id: ec_id.clone(), - }, - ); + let outcome = graph + .tombstone_existing_from_snapshot( + &ec_id, + EcKvSnapshot::Missing { + ec_id: ec_id.clone(), + }, + ) + .await; assert!( outcome @@ -2705,13 +2922,14 @@ mod tests { ); let (stored, _) = graph .get(&ec_id) + .await .expect("should read recreated row") .expect("should retain recreated row"); assert!(!stored.consent.ok, "recreated row should be tombstoned"); } - #[test] - fn withdrawal_failed_clock_stale_miss_tombstones_live_root() { + #[tokio::test(flavor = "current_thread")] + async fn withdrawal_failed_clock_stale_miss_tombstones_live_root() { let operations = Arc::new(RecordedEcKvOperations::default()); let graph = KvIdentityGraph::new(RecordingEcKv::with_stale_lookups( Arc::clone(&operations), @@ -2720,22 +2938,27 @@ mod tests { let ec_id = snapshot_ec_id(); graph .create(&ec_id, &concurrent_live_entry()) + .await .expect("should seed a live row with partner IDs"); graph .write_withdrawal_marker(&ec_id, 1000) + .await .expect("should seed a completion marker left after root recreation"); operations.reset(); - let missing = graph.load_snapshot(&ec_id); + let missing = graph.load_snapshot(&ec_id).await; assert!( matches!(missing, EcKvSnapshot::Missing { .. }), "should reproduce a stale point-read miss for the live row" ); // None is the checked clock's failure result, not Unix epoch zero. - let outcome = graph.tombstone_unproven_missing(&ec_id, missing, None); + let outcome = graph + .tombstone_unproven_missing(&ec_id, missing, None) + .await; let (stored, generation) = graph .get(&ec_id) + .await .expect("should read back persisted withdrawal") .expect("should retain the root"); assert!( @@ -2777,18 +3000,21 @@ mod tests { ); } - #[test] - fn withdrawal_failed_clock_does_not_create_absent_root() { + #[tokio::test(flavor = "current_thread")] + async fn withdrawal_failed_clock_does_not_create_absent_root() { let operations = Arc::new(RecordedEcKvOperations::default()); let graph = KvIdentityGraph::new(RecordingEcKv::new(Arc::clone(&operations))); let ec_id = snapshot_ec_id(); graph .write_withdrawal_marker(&ec_id, 1000) + .await .expect("should seed a marker without a root"); operations.reset(); - let missing = graph.load_snapshot(&ec_id); + let missing = graph.load_snapshot(&ec_id).await; - let outcome = graph.tombstone_unproven_missing(&ec_id, missing, None); + let outcome = graph + .tombstone_unproven_missing(&ec_id, missing, None) + .await; assert!( matches!(outcome, EcKvSnapshot::Missing { .. }), @@ -2797,6 +3023,7 @@ mod tests { assert!( graph .get(&ec_id) + .await .expect("should read absent root") .is_none(), "clock failure must not mint an unknown root" @@ -2812,12 +3039,13 @@ mod tests { ); } - #[test] - fn withdrawal_marker_validity_requires_a_usable_clock_before_expiry() { + #[tokio::test(flavor = "current_thread")] + async fn withdrawal_marker_validity_requires_a_usable_clock_before_expiry() { let graph = KvIdentityGraph::in_memory("test_store"); let ec_id = snapshot_ec_id(); graph .write_withdrawal_marker(&ec_id, 1000) + .await .expect("should seed marker"); let valid_until = 1000 + TOMBSTONE_TTL.as_secs(); @@ -2830,6 +3058,7 @@ mod tests { assert_eq!( graph .withdrawal_marker_exists(&ec_id, now) + .await .expect("should check marker validity"), expected, "marker validity should honor clock availability and exclusive expiry at {now:?}" @@ -2837,8 +3066,8 @@ mod tests { } } - #[test] - fn withdrawal_marker_existence_rejects_malformed_expiry_suffix() { + #[tokio::test(flavor = "current_thread")] + async fn withdrawal_marker_existence_rejects_malformed_expiry_suffix() { let ec_id = format!("{}.ABC123", "a".repeat(64)); let marker_key = KvIdentityGraph::withdrawal_marker_key( &ec_id, @@ -2856,23 +3085,26 @@ mod tests { mode: EcKvWriteMode::Add, }, ) + .await .expect("should seed longer marker key"); let kv = KvIdentityGraph::new(store); assert!( !kv.withdrawal_marker_exists(&ec_id, checked_current_timestamp()) + .await .expect("should validate withdrawal marker expiry"), "a malformed expiry suffix must not prove withdrawal completion" ); } - #[test] - fn clear_absent_withdrawal_marker_lists_once_without_delete() { + #[tokio::test(flavor = "current_thread")] + async fn clear_absent_withdrawal_marker_lists_once_without_delete() { let operations = Arc::new(RecordedEcKvOperations::default()); let kv = KvIdentityGraph::new(RecordingEcKv::new(Arc::clone(&operations))); let ec_id = snapshot_ec_id(); kv.clear_withdrawal_marker(&ec_id) + .await .expect("should accept an absent marker"); assert_eq!( @@ -2887,16 +3119,18 @@ mod tests { ); } - #[test] - fn writing_an_existing_withdrawal_marker_is_idempotent() { + #[tokio::test(flavor = "current_thread")] + async fn writing_an_existing_withdrawal_marker_is_idempotent() { let operations = Arc::new(RecordedEcKvOperations::default()); let kv = KvIdentityGraph::new(RecordingEcKv::new(Arc::clone(&operations))); let ec_id = snapshot_ec_id(); let updated = current_timestamp(); kv.write_withdrawal_marker(&ec_id, updated) + .await .expect("should write first completion marker"); kv.write_withdrawal_marker(&ec_id, updated) + .await .expect("should accept completion marker add collision"); assert_eq!( @@ -2915,18 +3149,20 @@ mod tests { ); assert!( kv.withdrawal_marker_exists(&ec_id, checked_current_timestamp()) + .await .expect("should read completion marker"), "completion marker should remain valid" ); } - #[test] - fn marker_list_saturation_blocks_revival_and_hard_delete_before_mutation() { + #[tokio::test(flavor = "current_thread")] + async fn marker_list_saturation_blocks_revival_and_hard_delete_before_mutation() { let operations = Arc::new(RecordedEcKvOperations::default()); let graph = KvIdentityGraph::new(RecordingEcKv::new(Arc::clone(&operations))); let ec_id = snapshot_ec_id(); graph .create(&ec_id, &KvEntry::tombstone(current_timestamp())) + .await .expect("should seed tombstone"); let first_valid_until = current_timestamp() + TOMBSTONE_TTL.as_secs(); for offset in 0..WITHDRAWAL_MARKER_LIST_LIMIT { @@ -2945,11 +3181,12 @@ mod tests { mode: EcKvWriteMode::Add, }, ) + .await .expect("should seed completion marker"); } operations.reset(); - let result = graph.create_or_revive(&ec_id, &live_entry()); + let result = graph.create_or_revive(&ec_id, &live_entry()).await; assert!( result.is_err(), @@ -2957,6 +3194,7 @@ mod tests { ); let (stored, _) = graph .get(&ec_id) + .await .expect("should read root") .expect("should retain root"); assert!(!stored.consent.ok, "root should remain tombstoned"); @@ -2967,7 +3205,7 @@ mod tests { ); operations.reset(); - let delete_result = graph.delete(&ec_id); + let delete_result = graph.delete(&ec_id).await; assert!( delete_result.is_err(), @@ -2975,6 +3213,7 @@ mod tests { ); let (stored, _) = graph .get(&ec_id) + .await .expect("should read root after rejected hard delete") .expect("rejected hard delete should retain root"); assert!( @@ -2988,106 +3227,120 @@ mod tests { ); } - #[test] - fn clear_withdrawal_marker_accepts_concurrent_removal_after_delete_error() { + #[tokio::test(flavor = "current_thread")] + async fn clear_withdrawal_marker_accepts_concurrent_removal_after_delete_error() { let ec_id = format!("{}.ABC123", "a".repeat(64)); let operations = Arc::new(RecordedEcKvOperations::default()); let kv = KvIdentityGraph::new(RecordingEcKv::with_marker_delete_failure(operations, true)); kv.create(&ec_id, &live_entry()) + .await .expect("should create live entry"); - let snapshot = kv.load_snapshot(&ec_id); - kv.tombstone_existing_from_snapshot(&ec_id, snapshot); + let snapshot = kv.load_snapshot(&ec_id).await; + kv.tombstone_existing_from_snapshot(&ec_id, snapshot).await; kv.clear_withdrawal_marker(&ec_id) + .await .expect("should accept a marker removed by another request"); assert!( !kv.withdrawal_marker_exists(&ec_id, checked_current_timestamp()) + .await .expect("should confirm marker removal"), "completion marker should remain absent" ); } - #[test] - fn clear_withdrawal_marker_preserves_delete_error_when_marker_remains() { + #[tokio::test(flavor = "current_thread")] + async fn clear_withdrawal_marker_preserves_delete_error_when_marker_remains() { let ec_id = format!("{}.ABC123", "a".repeat(64)); let operations = Arc::new(RecordedEcKvOperations::default()); let kv = KvIdentityGraph::new(RecordingEcKv::with_marker_delete_failure(operations, false)); kv.create(&ec_id, &live_entry()) + .await .expect("should create live entry"); - let snapshot = kv.load_snapshot(&ec_id); - kv.tombstone_existing_from_snapshot(&ec_id, snapshot); + let snapshot = kv.load_snapshot(&ec_id).await; + kv.tombstone_existing_from_snapshot(&ec_id, snapshot).await; assert!( - kv.clear_withdrawal_marker(&ec_id).is_err(), + kv.clear_withdrawal_marker(&ec_id).await.is_err(), "a marker that remains after delete failure should block revival" ); assert!( kv.withdrawal_marker_exists(&ec_id, checked_current_timestamp()) + .await .expect("should confirm marker remains"), "completion marker should remain present" ); } - #[test] - fn withdrawal_marker_is_excluded_from_hash_prefix_count() { + #[tokio::test(flavor = "current_thread")] + async fn withdrawal_marker_is_excluded_from_hash_prefix_count() { let hash = "a".repeat(64); let ec_id = format!("{hash}.ABC123"); let kv = KvIdentityGraph::in_memory("test_store"); kv.create(&ec_id, &live_entry()) + .await .expect("should create live entry"); assert_eq!( kv.count_hash_prefix_keys(&hash) + .await .expect("should count live root"), 1, "the root should be the only hash-prefix key" ); - let snapshot = kv.load_snapshot(&ec_id); - kv.tombstone_existing_from_snapshot(&ec_id, snapshot); + let snapshot = kv.load_snapshot(&ec_id).await; + kv.tombstone_existing_from_snapshot(&ec_id, snapshot).await; assert!( kv.withdrawal_marker_exists(&ec_id, checked_current_timestamp()) + .await .expect("should read withdrawal marker"), "withdrawal should record completion" ); assert_eq!( kv.count_hash_prefix_keys(&hash) + .await .expect("should count tombstoned root"), 1, "the marker namespace must not affect cluster counts" ); } - #[test] - fn delete_removes_withdrawal_marker() { + #[tokio::test(flavor = "current_thread")] + async fn delete_removes_withdrawal_marker() { let kv = KvIdentityGraph::in_memory("test_store"); let ec_id = format!("{}.ABC123", "a".repeat(64)); kv.create(&ec_id, &live_entry()) + .await .expect("should create live entry"); - let snapshot = kv.load_snapshot(&ec_id); - kv.tombstone_existing_from_snapshot(&ec_id, snapshot); + let snapshot = kv.load_snapshot(&ec_id).await; + kv.tombstone_existing_from_snapshot(&ec_id, snapshot).await; - kv.delete(&ec_id).expect("should delete entry and marker"); + kv.delete(&ec_id) + .await + .expect("should delete entry and marker"); assert!( - kv.get(&ec_id).expect("should read store").is_none(), + kv.get(&ec_id).await.expect("should read store").is_none(), "hard delete should remove the root" ); assert!( !kv.withdrawal_marker_exists(&ec_id, checked_current_timestamp()) + .await .expect("should read withdrawal marker"), "hard delete should remove withdrawal completion" ); } - #[test] - fn upsert_partner_id_if_exists_reports_missing_key() { + #[tokio::test(flavor = "current_thread")] + async fn upsert_partner_id_if_exists_reports_missing_key() { let kv = KvIdentityGraph::in_memory("test_store"); let ec_id = format!("{}.ABC123", "a".repeat(64)); let result = kv .upsert_partner_id_if_exists(&ec_id, "ssp_x", "uid-1") + .await .expect("should not error on missing key"); assert_eq!( result, @@ -3096,14 +3349,17 @@ mod tests { ); } - #[test] - fn upsert_partner_id_if_exists_writes_and_detects_unchanged() { + #[tokio::test(flavor = "current_thread")] + async fn upsert_partner_id_if_exists_writes_and_detects_unchanged() { let kv = KvIdentityGraph::in_memory("test_store"); let ec_id = format!("{}.ABC123", "a".repeat(64)); - kv.create(&ec_id, &live_entry()).expect("should create"); + kv.create(&ec_id, &live_entry()) + .await + .expect("should create"); let first = kv .upsert_partner_id_if_exists(&ec_id, "ssp_x", "uid-1") + .await .expect("should write partner id"); assert_eq!( first, @@ -3113,6 +3369,7 @@ mod tests { let second = kv .upsert_partner_id_if_exists(&ec_id, "ssp_x", "uid-1") + .await .expect("should detect unchanged uid"); assert_eq!( second, @@ -3121,21 +3378,24 @@ mod tests { ); } - #[test] - fn upsert_partner_id_if_exists_retries_cas_conflict() { + #[tokio::test(flavor = "current_thread")] + async fn upsert_partner_id_if_exists_retries_cas_conflict() { let graph = KvIdentityGraph::new(ConflictInjectingEcKv::new(1, false)); let ec_id = format!("{}.ABC123", "a".repeat(64)); graph .create(&ec_id, &live_entry()) + .await .expect("should seed live entry"); let result = graph .upsert_partner_id_if_exists(&ec_id, "ssp_x", "uid-1") + .await .expect("should retry and write after one generation conflict"); assert_eq!(result, UpsertResult::Written); let (entry, _) = graph .get(&ec_id) + .await .expect("should read persisted entry") .expect("should retain entry after conflict"); assert_eq!( @@ -3145,15 +3405,17 @@ mod tests { ); } - #[test] - fn upsert_partner_id_if_exists_rejects_tombstone() { + #[tokio::test(flavor = "current_thread")] + async fn upsert_partner_id_if_exists_rejects_tombstone() { let kv = KvIdentityGraph::in_memory("test_store"); let ec_id = format!("{}.ABC123", "a".repeat(64)); kv.create(&ec_id, &KvEntry::tombstone(1000)) + .await .expect("should create tombstone"); let result = kv .upsert_partner_id_if_exists(&ec_id, "ssp_x", "uid-1") + .await .expect("should not error on tombstone"); assert_eq!( result, @@ -3162,15 +3424,19 @@ mod tests { ); } - #[test] - fn snapshot_bulk_upsert_returns_persisted_entry_without_stale_generation() { + #[tokio::test(flavor = "current_thread")] + async fn snapshot_bulk_upsert_returns_persisted_entry_without_stale_generation() { let kv = KvIdentityGraph::in_memory("test_store"); let ec_id = format!("{}.ABC123", "a".repeat(64)); - kv.create(&ec_id, &live_entry()).expect("should create"); - let snapshot = kv.load_snapshot(&ec_id); + kv.create(&ec_id, &live_entry()) + .await + .expect("should create"); + let snapshot = kv.load_snapshot(&ec_id).await; let updates = [PartnerIdUpdate::new("ssp_x", "uid-1")]; - let outcome = kv.upsert_partner_ids_from_snapshot(&ec_id, &updates, snapshot); + let outcome = kv + .upsert_partner_ids_from_snapshot(&ec_id, &updates, snapshot) + .await; let entry = outcome .entry_for(&ec_id) @@ -3186,55 +3452,59 @@ mod tests { ); } - #[test] - fn snapshot_bulk_upsert_does_not_create_missing_root() { + #[tokio::test(flavor = "current_thread")] + async fn snapshot_bulk_upsert_does_not_create_missing_root() { let kv = KvIdentityGraph::in_memory("test_store"); let ec_id = format!("{}.ABC123", "a".repeat(64)); let updates = [PartnerIdUpdate::new("ssp_x", "uid-1")]; - let outcome = kv.upsert_partner_ids_from_snapshot( - &ec_id, - &updates, - EcKvSnapshot::Missing { - ec_id: ec_id.clone(), - }, - ); + let outcome = kv + .upsert_partner_ids_from_snapshot( + &ec_id, + &updates, + EcKvSnapshot::Missing { + ec_id: ec_id.clone(), + }, + ) + .await; assert!( matches!(outcome, EcKvSnapshot::Missing { .. }), "missing snapshot should remain missing" ); - assert!(kv.get(&ec_id).expect("should read store").is_none()); + assert!(kv.get(&ec_id).await.expect("should read store").is_none()); } - #[test] - fn tombstone_existing_from_snapshot_never_creates_missing_key() { + #[tokio::test(flavor = "current_thread")] + async fn tombstone_existing_from_snapshot_never_creates_missing_key() { let kv = KvIdentityGraph::in_memory("test_store"); let ec_id = format!("{}.ABC123", "a".repeat(64)); let snapshot = EcKvSnapshot::Missing { ec_id: ec_id.clone(), }; - let outcome = kv.tombstone_existing_from_snapshot(&ec_id, snapshot); + let outcome = kv.tombstone_existing_from_snapshot(&ec_id, snapshot).await; assert!( matches!(outcome, EcKvSnapshot::Missing { .. }), "withdrawal should preserve the missing snapshot for an absent identity" ); assert!( - kv.get(&ec_id).expect("should read store").is_none(), + kv.get(&ec_id).await.expect("should read store").is_none(), "withdrawal must not create a tombstone for an absent key" ); } - #[test] - fn tombstone_existing_from_snapshot_uses_existing_generation() { + #[tokio::test(flavor = "current_thread")] + async fn tombstone_existing_from_snapshot_uses_existing_generation() { let kv = KvIdentityGraph::in_memory("test_store"); let ec_id = format!("{}.ABC123", "a".repeat(64)); - kv.create(&ec_id, &live_entry()).expect("should create"); - let snapshot = kv.load_snapshot(&ec_id); + kv.create(&ec_id, &live_entry()) + .await + .expect("should create"); + let snapshot = kv.load_snapshot(&ec_id).await; - let outcome = kv.tombstone_existing_from_snapshot(&ec_id, snapshot); + let outcome = kv.tombstone_existing_from_snapshot(&ec_id, snapshot).await; assert!( outcome @@ -3244,19 +3514,23 @@ mod tests { ); let (stored, _) = kv .get(&ec_id) + .await .expect("should read store") .expect("should preserve existing key"); assert!(!stored.consent.ok, "should persist withdrawal state"); } - #[test] - fn write_withdrawal_tombstone_overwrites_live_entry() { + #[tokio::test(flavor = "current_thread")] + async fn write_withdrawal_tombstone_overwrites_live_entry() { let kv = KvIdentityGraph::in_memory("test_store"); let ec_id = format!("{}.ABC123", "a".repeat(64)); - kv.create(&ec_id, &live_entry()).expect("should create"); + kv.create(&ec_id, &live_entry()) + .await + .expect("should create"); assert_eq!( kv.write_withdrawal_tombstone(&ec_id, drop) + .await .expect("should write tombstone"), TombstoneOutcome::Written, "should tombstone an identity the store holds" @@ -3264,6 +3538,7 @@ mod tests { let (loaded, _) = kv .get(&ec_id) + .await .expect("should read entry back") .expect("should find tombstone entry"); assert!(!loaded.consent.ok, "should be withdrawn after tombstone"); @@ -3378,7 +3653,7 @@ mod tests { } } - fn completed_with_root_check_failure( + async fn completed_with_root_check_failure( operations: Arc, ec_id: &str, ) -> Self { @@ -3402,6 +3677,7 @@ mod tests { mode: EcKvWriteMode::Add, }, ) + .await .expect("should seed tombstone"); store .inner @@ -3417,6 +3693,7 @@ mod tests { mode: EcKvWriteMode::Add, }, ) + .await .expect("should seed completion marker"); store } @@ -3446,24 +3723,30 @@ mod tests { } } + #[async_trait::async_trait(?Send)] impl EcKvStore for RecordingEcKv { fn store_name(&self) -> &str { self.inner.store_name() } - fn lookup(&self, key: &str) -> Result, Report> { + async fn lookup( + &self, + key: &str, + ) -> Result, Report> { self.operations .lookups .fetch_add(1, std::sync::atomic::Ordering::Relaxed); - let mut stale_lookups = self - .stale_lookups_remaining - .lock() - .expect("should lock stale lookup counter"); - if *stale_lookups > 0 { - *stale_lookups -= 1; - return Ok(None); + { + let mut stale_lookups = self + .stale_lookups_remaining + .lock() + .expect("should lock stale lookup counter"); + if *stale_lookups > 0 { + *stale_lookups -= 1; + return Ok(None); + } } - let found = self.inner.lookup(key)?; + let found = self.inner.lookup(key).await?; if self.lag_live_entries && found.as_ref().is_some_and(|entry| { serde_json::from_slice::(&entry.body) @@ -3477,17 +3760,17 @@ mod tests { Ok(found) } - fn key_exists(&self, key: &str) -> Result> { + async fn key_exists(&self, key: &str) -> Result> { self.operations .exact_checks .fetch_add(1, std::sync::atomic::Ordering::Relaxed); if !key.starts_with(WITHDRAWAL_MARKER_PREFIX) && self.root_check_fails { return Err(self.root_check_error()); } - self.inner.key_exists(key) + self.inner.key_exists(key).await } - fn insert( + async fn insert( &self, key: &str, write: EcKvWrite<'_>, @@ -3503,10 +3786,10 @@ mod tests { if key.starts_with(WITHDRAWAL_MARKER_PREFIX) && self.marker_operations_fail { return Err(self.marker_error("write")); } - self.inner.insert(key, write) + self.inner.insert(key, write).await } - fn list_keys_with_prefix( + async fn list_keys_with_prefix( &self, prefix: &str, limit: u32, @@ -3517,10 +3800,10 @@ mod tests { if prefix.starts_with(WITHDRAWAL_MARKER_PREFIX) && self.marker_operations_fail { return Err(self.marker_error("lookup")); } - self.inner.list_keys_with_prefix(prefix, limit) + self.inner.list_keys_with_prefix(prefix, limit).await } - fn delete(&self, key: &str) -> Result<(), Report> { + async fn delete(&self, key: &str) -> Result<(), Report> { self.operations .deletes .fetch_add(1, std::sync::atomic::Ordering::Relaxed); @@ -3528,11 +3811,11 @@ mod tests { && let Some(remove_before_error) = self.marker_delete_failure_removes_key { if remove_before_error { - self.inner.delete(key)?; + self.inner.delete(key).await?; } return Err(self.marker_error("delete")); } - self.inner.delete(key) + self.inner.delete(key).await } } @@ -3550,18 +3833,22 @@ mod tests { } } + #[async_trait::async_trait(?Send)] impl EcKvStore for WriteFailingEcKv { fn store_name(&self) -> &str { self.inner.store_name() } - fn lookup(&self, key: &str) -> Result, Report> { - self.inner.lookup(key) + async fn lookup( + &self, + key: &str, + ) -> Result, Report> { + self.inner.lookup(key).await } - fn key_exists(&self, key: &str) -> Result> { - self.inner.key_exists(key) + async fn key_exists(&self, key: &str) -> Result> { + self.inner.key_exists(key).await } - fn insert( + async fn insert( &self, _key: &str, _write: EcKvWrite<'_>, @@ -3571,24 +3858,27 @@ mod tests { message: "write failing test store".to_owned(), })) } - fn list_keys_with_prefix( + async fn list_keys_with_prefix( &self, prefix: &str, limit: u32, ) -> Result, Report> { - self.inner.list_keys_with_prefix(prefix, limit) + self.inner.list_keys_with_prefix(prefix, limit).await } - fn delete(&self, key: &str) -> Result<(), Report> { - self.inner.delete(key) + async fn delete(&self, key: &str) -> Result<(), Report> { + self.inner.delete(key).await } } - #[test] - fn snapshot_upsert_with_generation_writes_without_reading() { + #[tokio::test(flavor = "current_thread")] + async fn snapshot_upsert_with_generation_writes_without_reading() { let lookups = std::sync::Arc::new(std::sync::atomic::AtomicUsize::new(0)); let graph = KvIdentityGraph::counting("counting-store", lookups.clone()); let ec_id = snapshot_ec_id(); - graph.create(&ec_id, &live_entry()).expect("should seed"); + graph + .create(&ec_id, &live_entry()) + .await + .expect("should seed"); let snapshot = EcKvSnapshot::Present { ec_id: ec_id.clone(), entry: Box::new(live_entry()), @@ -3596,7 +3886,9 @@ mod tests { }; let updates = [PartnerIdUpdate::new("ssp_x", "uid-1")]; - let outcome = graph.upsert_partner_ids_from_snapshot(&ec_id, &updates, snapshot); + let outcome = graph + .upsert_partner_ids_from_snapshot(&ec_id, &updates, snapshot) + .await; assert_eq!( lookups.load(std::sync::atomic::Ordering::Relaxed), @@ -3617,25 +3909,27 @@ mod tests { ); } - #[test] - fn snapshot_upsert_unchanged_updates_preserve_generation() { + #[tokio::test(flavor = "current_thread")] + async fn snapshot_upsert_unchanged_updates_preserve_generation() { let kv = KvIdentityGraph::in_memory("test_store"); let ec_id = snapshot_ec_id(); let mut seeded = live_entry(); apply_partner_id_updates(&mut seeded, &[PartnerIdUpdate::new("ssp_x", "uid-1")]); - kv.create(&ec_id, &seeded).expect("should seed"); - let snapshot = kv.load_snapshot(&ec_id); + kv.create(&ec_id, &seeded).await.expect("should seed"); + let snapshot = kv.load_snapshot(&ec_id).await; assert_eq!( snapshot.generation_for(&ec_id), Some(1), "seeded snapshot should carry its stored generation" ); - let outcome = kv.upsert_partner_ids_from_snapshot( - &ec_id, - &[PartnerIdUpdate::new("ssp_x", "uid-1")], - snapshot, - ); + let outcome = kv + .upsert_partner_ids_from_snapshot( + &ec_id, + &[PartnerIdUpdate::new("ssp_x", "uid-1")], + snapshot, + ) + .await; assert_eq!( outcome.generation_for(&ec_id), @@ -3644,12 +3938,15 @@ mod tests { ); } - #[test] - fn snapshot_upsert_refreshes_unavailable_generation_exactly_once() { + #[tokio::test(flavor = "current_thread")] + async fn snapshot_upsert_refreshes_unavailable_generation_exactly_once() { let lookups = std::sync::Arc::new(std::sync::atomic::AtomicUsize::new(0)); let graph = KvIdentityGraph::counting("counting-store", lookups.clone()); let ec_id = snapshot_ec_id(); - graph.create(&ec_id, &live_entry()).expect("should seed"); + graph + .create(&ec_id, &live_entry()) + .await + .expect("should seed"); // Finalize-written style snapshot: entry known, generation unavailable. let snapshot = EcKvSnapshot::Present { ec_id: ec_id.clone(), @@ -3658,7 +3955,9 @@ mod tests { }; let updates = [PartnerIdUpdate::new("ssp_x", "uid-1")]; - let outcome = graph.upsert_partner_ids_from_snapshot(&ec_id, &updates, snapshot); + let outcome = graph + .upsert_partner_ids_from_snapshot(&ec_id, &updates, snapshot) + .await; assert_eq!( lookups.load(std::sync::atomic::Ordering::Relaxed), @@ -3673,15 +3972,18 @@ mod tests { ); } - #[test] - fn snapshot_upsert_gen_unavailable_survives_four_conflicts_then_writes() { + #[tokio::test(flavor = "current_thread")] + async fn snapshot_upsert_gen_unavailable_survives_four_conflicts_then_writes() { // A generation-unavailable snapshot (finalize-written style) refreshes // once to obtain a usable generation. That refresh must not consume a // CAS attempt, so all five write attempts remain: four conflicts // followed by a successful fifth write still persist the update. let graph = KvIdentityGraph::new(ConflictInjectingEcKv::new(4, false)); let ec_id = snapshot_ec_id(); - graph.create(&ec_id, &live_entry()).expect("should seed"); + graph + .create(&ec_id, &live_entry()) + .await + .expect("should seed"); let snapshot = EcKvSnapshot::Present { ec_id: ec_id.clone(), entry: Box::new(live_entry()), @@ -3689,7 +3991,9 @@ mod tests { }; let updates = [PartnerIdUpdate::new("ssp_x", "uid-1")]; - let outcome = graph.upsert_partner_ids_from_snapshot(&ec_id, &updates, snapshot); + let outcome = graph + .upsert_partner_ids_from_snapshot(&ec_id, &updates, snapshot) + .await; assert_eq!( outcome @@ -3701,15 +4005,20 @@ mod tests { ); } - #[test] - fn snapshot_upsert_cas_conflict_remerges_concurrent_data() { + #[tokio::test(flavor = "current_thread")] + async fn snapshot_upsert_cas_conflict_remerges_concurrent_data() { let graph = KvIdentityGraph::new(ConflictInjectingEcKv::new(1, true)); let ec_id = snapshot_ec_id(); - graph.create(&ec_id, &live_entry()).expect("should seed"); - let snapshot = graph.load_snapshot(&ec_id); + graph + .create(&ec_id, &live_entry()) + .await + .expect("should seed"); + let snapshot = graph.load_snapshot(&ec_id).await; let updates = [PartnerIdUpdate::new("ssp_x", "uid-1")]; - let outcome = graph.upsert_partner_ids_from_snapshot(&ec_id, &updates, snapshot); + let outcome = graph + .upsert_partner_ids_from_snapshot(&ec_id, &updates, snapshot) + .await; let entry = outcome .entry_for(&ec_id) @@ -3722,24 +4031,28 @@ mod tests { assert!(entry.consent.ok, "concurrent revive keeps the row live"); } - #[test] - fn snapshot_upsert_revalidates_transient_missing_and_persists() { + #[tokio::test(flavor = "current_thread")] + async fn snapshot_upsert_revalidates_transient_missing_and_persists() { // An eventually-consistent point read earlier in the request missed a // row that exists. Named routes such as `/auction` never run orphan // recovery, so this refresh is the request's only chance to persist the // collected partner IDs. let kv = KvIdentityGraph::in_memory("test_store"); let ec_id = snapshot_ec_id(); - kv.create(&ec_id, &live_entry()).expect("should seed live"); + kv.create(&ec_id, &live_entry()) + .await + .expect("should seed live"); let updates = [PartnerIdUpdate::new("ssp_x", "uid-1")]; - let outcome = kv.upsert_partner_ids_from_snapshot( - &ec_id, - &updates, - EcKvSnapshot::Missing { - ec_id: ec_id.clone(), - }, - ); + let outcome = kv + .upsert_partner_ids_from_snapshot( + &ec_id, + &updates, + EcKvSnapshot::Missing { + ec_id: ec_id.clone(), + }, + ) + .await; assert_eq!( outcome @@ -3750,6 +4063,7 @@ mod tests { ); let (stored, _) = kv .get(&ec_id) + .await .expect("should read store") .expect("row should remain"); assert_eq!( @@ -3759,48 +4073,54 @@ mod tests { ); } - #[test] - fn snapshot_upsert_confirmed_missing_still_never_creates() { + #[tokio::test(flavor = "current_thread")] + async fn snapshot_upsert_confirmed_missing_still_never_creates() { // Revalidation only changes what a *stale* miss does. A row that is // genuinely absent on the refresh must stay absent. let kv = KvIdentityGraph::in_memory("test_store"); let ec_id = snapshot_ec_id(); let updates = [PartnerIdUpdate::new("ssp_x", "uid-1")]; - let outcome = kv.upsert_partner_ids_from_snapshot( - &ec_id, - &updates, - EcKvSnapshot::Missing { - ec_id: ec_id.clone(), - }, - ); + let outcome = kv + .upsert_partner_ids_from_snapshot( + &ec_id, + &updates, + EcKvSnapshot::Missing { + ec_id: ec_id.clone(), + }, + ) + .await; assert!( matches!(outcome, EcKvSnapshot::Missing { .. }), "a confirmed miss must stay missing" ); assert!( - kv.get(&ec_id).expect("should read store").is_none(), + kv.get(&ec_id).await.expect("should read store").is_none(), "must not create a root entry for a missing key" ); } - #[test] - fn snapshot_upsert_failed_snapshot_is_not_revalidated() { + #[tokio::test(flavor = "current_thread")] + async fn snapshot_upsert_failed_snapshot_is_not_revalidated() { // A lookup that already errored is not retried on the hot path, even // though the row exists. let kv = KvIdentityGraph::in_memory("test_store"); let ec_id = snapshot_ec_id(); - kv.create(&ec_id, &live_entry()).expect("should seed live"); + kv.create(&ec_id, &live_entry()) + .await + .expect("should seed live"); let updates = [PartnerIdUpdate::new("ssp_x", "uid-1")]; - let outcome = kv.upsert_partner_ids_from_snapshot( - &ec_id, - &updates, - EcKvSnapshot::Failed { - ec_id: ec_id.clone(), - }, - ); + let outcome = kv + .upsert_partner_ids_from_snapshot( + &ec_id, + &updates, + EcKvSnapshot::Failed { + ec_id: ec_id.clone(), + }, + ) + .await; assert!( matches!(outcome, EcKvSnapshot::Failed { .. }), @@ -3808,16 +4128,19 @@ mod tests { ); } - #[test] - fn snapshot_upsert_rejects_tombstone() { + #[tokio::test(flavor = "current_thread")] + async fn snapshot_upsert_rejects_tombstone() { let kv = KvIdentityGraph::in_memory("test_store"); let ec_id = snapshot_ec_id(); kv.create(&ec_id, &KvEntry::tombstone(1000)) + .await .expect("should seed tombstone"); - let snapshot = kv.load_snapshot(&ec_id); + let snapshot = kv.load_snapshot(&ec_id).await; let updates = [PartnerIdUpdate::new("ssp_x", "uid-1")]; - let outcome = kv.upsert_partner_ids_from_snapshot(&ec_id, &updates, snapshot); + let outcome = kv + .upsert_partner_ids_from_snapshot(&ec_id, &updates, snapshot) + .await; assert!( outcome @@ -3827,13 +4150,14 @@ mod tests { ); let (stored, _) = kv .get(&ec_id) + .await .expect("should read store") .expect("tombstone should remain"); assert!(stored.ids.is_empty(), "no update should reach the store"); } - #[test] - fn snapshot_upsert_store_failure_returns_failed_not_request_local() { + #[tokio::test(flavor = "current_thread")] + async fn snapshot_upsert_store_failure_returns_failed_not_request_local() { let graph = KvIdentityGraph::new(WriteFailingEcKv::new()); let ec_id = snapshot_ec_id(); let snapshot = EcKvSnapshot::Present { @@ -3843,7 +4167,9 @@ mod tests { }; let updates = [PartnerIdUpdate::new("ssp_x", "uid-1")]; - let outcome = graph.upsert_partner_ids_from_snapshot(&ec_id, &updates, snapshot); + let outcome = graph + .upsert_partner_ids_from_snapshot(&ec_id, &updates, snapshot) + .await; assert!( matches!(outcome, EcKvSnapshot::Failed { .. }), @@ -3851,8 +4177,8 @@ mod tests { ); } - #[test] - fn tombstone_existing_from_snapshot_skips_backend_for_authoritative_tombstone() { + #[tokio::test(flavor = "current_thread")] + async fn tombstone_existing_from_snapshot_skips_backend_for_authoritative_tombstone() { for generation in [Some(7), None] { let operations = Arc::new(RecordedEcKvOperations::default()); let graph = KvIdentityGraph::new(RecordingEcKv::new(Arc::clone(&operations))); @@ -3863,7 +4189,9 @@ mod tests { generation, }; - let outcome = graph.tombstone_existing_from_snapshot(&ec_id, snapshot.clone()); + let outcome = graph + .tombstone_existing_from_snapshot(&ec_id, snapshot.clone()) + .await; assert_eq!( outcome, snapshot, @@ -3877,18 +4205,21 @@ mod tests { } } - #[test] - fn tombstone_existing_from_snapshot_repeated_request_preserves_first_write() { + #[tokio::test(flavor = "current_thread")] + async fn tombstone_existing_from_snapshot_repeated_request_preserves_first_write() { let operations = Arc::new(RecordedEcKvOperations::default()); let graph = KvIdentityGraph::new(RecordingEcKv::new(Arc::clone(&operations))); let ec_id = snapshot_ec_id(); graph .create(&ec_id, &live_entry()) + .await .expect("should seed live row"); - let live_snapshot = graph.load_snapshot(&ec_id); + let live_snapshot = graph.load_snapshot(&ec_id).await; operations.reset(); - graph.tombstone_existing_from_snapshot(&ec_id, live_snapshot); + graph + .tombstone_existing_from_snapshot(&ec_id, live_snapshot) + .await; assert_eq!( operations.lookup_count(), @@ -3909,7 +4240,7 @@ mod tests { ], "first withdrawal should write the root and its completion marker" ); - let first_snapshot = graph.load_snapshot(&ec_id); + let first_snapshot = graph.load_snapshot(&ec_id).await; let (first_entry, first_generation) = match &first_snapshot { EcKvSnapshot::Present { entry, generation, .. @@ -3918,7 +4249,9 @@ mod tests { }; operations.reset(); - let second_outcome = graph.tombstone_existing_from_snapshot(&ec_id, first_snapshot); + let second_outcome = graph + .tombstone_existing_from_snapshot(&ec_id, first_snapshot) + .await; assert_eq!( operations.operation_count(), @@ -3939,8 +4272,8 @@ mod tests { ); } - #[test] - fn tombstone_existing_from_repeated_stale_miss_preserves_first_write() { + #[tokio::test(flavor = "current_thread")] + async fn tombstone_existing_from_repeated_stale_miss_preserves_first_write() { let operations = Arc::new(RecordedEcKvOperations::default()); let graph = KvIdentityGraph::new(RecordingEcKv::with_stale_lookups( Arc::clone(&operations), @@ -3949,15 +4282,18 @@ mod tests { let ec_id = snapshot_ec_id(); graph .create(&ec_id, &live_entry()) + .await .expect("should seed live row"); operations.reset(); - let first_outcome = graph.tombstone_existing_from_snapshot( - &ec_id, - EcKvSnapshot::Missing { - ec_id: ec_id.clone(), - }, - ); + let first_outcome = graph + .tombstone_existing_from_snapshot( + &ec_id, + EcKvSnapshot::Missing { + ec_id: ec_id.clone(), + }, + ) + .await; let first_updated = first_outcome .entry_for(&ec_id) .expect("should return first tombstone") @@ -3985,12 +4321,14 @@ mod tests { ); operations.reset(); - graph.tombstone_existing_from_snapshot( - &ec_id, - EcKvSnapshot::Missing { - ec_id: ec_id.clone(), - }, - ); + graph + .tombstone_existing_from_snapshot( + &ec_id, + EcKvSnapshot::Missing { + ec_id: ec_id.clone(), + }, + ) + .await; assert_eq!( operations.lookup_count(), @@ -4018,6 +4356,7 @@ mod tests { ); let (stored, generation) = graph .get(&ec_id) + .await .expect("should read stored tombstone") .expect("should preserve tombstone"); assert_eq!( @@ -4030,19 +4369,24 @@ mod tests { ); } - #[test] - fn tombstone_existing_from_stale_parallel_snapshot_stops_after_conflict() { + #[tokio::test(flavor = "current_thread")] + async fn tombstone_existing_from_stale_parallel_snapshot_stops_after_conflict() { let operations = Arc::new(RecordedEcKvOperations::default()); let graph = KvIdentityGraph::new(RecordingEcKv::new(Arc::clone(&operations))); let ec_id = snapshot_ec_id(); graph .create(&ec_id, &live_entry()) + .await .expect("should seed live row"); - let stale_snapshot = graph.load_snapshot(&ec_id); + let stale_snapshot = graph.load_snapshot(&ec_id).await; operations.reset(); - let first_outcome = graph.tombstone_existing_from_snapshot(&ec_id, stale_snapshot.clone()); - let second_outcome = graph.tombstone_existing_from_snapshot(&ec_id, stale_snapshot); + let first_outcome = graph + .tombstone_existing_from_snapshot(&ec_id, stale_snapshot.clone()) + .await; + let second_outcome = graph + .tombstone_existing_from_snapshot(&ec_id, stale_snapshot) + .await; assert_eq!( operations.inserts(), @@ -4083,8 +4427,8 @@ mod tests { ); } - #[test] - fn tombstone_existing_from_snapshot_succeeds_without_backend_for_tombstone() { + #[tokio::test(flavor = "current_thread")] + async fn tombstone_existing_from_snapshot_succeeds_without_backend_for_tombstone() { let graph = KvIdentityGraph::failing("unavailable-store"); let ec_id = snapshot_ec_id(); let snapshot = EcKvSnapshot::Present { @@ -4093,7 +4437,9 @@ mod tests { generation: Some(3), }; - let outcome = graph.tombstone_existing_from_snapshot(&ec_id, snapshot.clone()); + let outcome = graph + .tombstone_existing_from_snapshot(&ec_id, snapshot.clone()) + .await; assert_eq!( outcome, snapshot, @@ -4101,8 +4447,8 @@ mod tests { ); } - #[test] - fn tombstone_existing_from_snapshot_non_authoritative_states_reread_live_row() { + #[tokio::test(flavor = "current_thread")] + async fn tombstone_existing_from_snapshot_non_authoritative_states_reread_live_row() { let ec_id = snapshot_ec_id(); let states = [ EcKvSnapshot::NotRead, @@ -4126,10 +4472,11 @@ mod tests { let graph = KvIdentityGraph::new(RecordingEcKv::new(Arc::clone(&operations))); graph .create(&ec_id, &live_entry()) + .await .expect("should seed live row"); operations.reset(); - let outcome = graph.tombstone_existing_from_snapshot(&ec_id, state); + let outcome = graph.tombstone_existing_from_snapshot(&ec_id, state).await; assert_eq!( operations.lookup_count(), @@ -4150,22 +4497,23 @@ mod tests { } } - #[test] - fn tombstone_stale_miss_valid_marker_skips_root_check_and_writes() { + #[tokio::test(flavor = "current_thread")] + async fn tombstone_stale_miss_valid_marker_skips_root_check_and_writes() { let ec_id = snapshot_ec_id(); let operations = Arc::new(RecordedEcKvOperations::default()); - let graph = KvIdentityGraph::new(RecordingEcKv::completed_with_root_check_failure( - Arc::clone(&operations), - &ec_id, - )); - - let outcome = graph.tombstone_existing_from_snapshot( - &ec_id, - EcKvSnapshot::Missing { - ec_id: ec_id.clone(), - }, + let graph = KvIdentityGraph::new( + RecordingEcKv::completed_with_root_check_failure(Arc::clone(&operations), &ec_id).await, ); + let outcome = graph + .tombstone_existing_from_snapshot( + &ec_id, + EcKvSnapshot::Missing { + ec_id: ec_id.clone(), + }, + ) + .await; + assert!( matches!(outcome, EcKvSnapshot::Missing { .. }), "a completion marker should resolve a repeated stale miss" @@ -4191,8 +4539,8 @@ mod tests { ); } - #[test] - fn tombstone_stale_miss_still_writes_when_marker_operations_fail() { + #[tokio::test(flavor = "current_thread")] + async fn tombstone_stale_miss_still_writes_when_marker_operations_fail() { let operations = Arc::new(RecordedEcKvOperations::default()); let graph = KvIdentityGraph::new(RecordingEcKv::with_marker_failures( Arc::clone(&operations), @@ -4201,15 +4549,18 @@ mod tests { let ec_id = snapshot_ec_id(); graph .create(&ec_id, &live_entry()) + .await .expect("should seed live row"); operations.reset(); - let outcome = graph.tombstone_existing_from_snapshot( - &ec_id, - EcKvSnapshot::Missing { - ec_id: ec_id.clone(), - }, - ); + let outcome = graph + .tombstone_existing_from_snapshot( + &ec_id, + EcKvSnapshot::Missing { + ec_id: ec_id.clone(), + }, + ) + .await; assert!( outcome @@ -4233,19 +4584,25 @@ mod tests { ); let (stored, _) = graph .get(&ec_id) + .await .expect("should read stored row") .expect("should preserve the root"); assert!(!stored.consent.ok, "root should remain tombstoned"); } - #[test] - fn tombstone_existing_from_snapshot_retries_cas_conflict() { + #[tokio::test(flavor = "current_thread")] + async fn tombstone_existing_from_snapshot_retries_cas_conflict() { let graph = KvIdentityGraph::new(ConflictInjectingEcKv::new(1, false)); let ec_id = snapshot_ec_id(); - graph.create(&ec_id, &live_entry()).expect("should seed"); - let snapshot = graph.load_snapshot(&ec_id); + graph + .create(&ec_id, &live_entry()) + .await + .expect("should seed"); + let snapshot = graph.load_snapshot(&ec_id).await; - let outcome = graph.tombstone_existing_from_snapshot(&ec_id, snapshot); + let outcome = graph + .tombstone_existing_from_snapshot(&ec_id, snapshot) + .await; assert!( outcome @@ -4255,21 +4612,26 @@ mod tests { ); } - #[test] - fn tombstone_gen_unavailable_survives_four_conflicts_then_writes() { + #[tokio::test(flavor = "current_thread")] + async fn tombstone_gen_unavailable_survives_four_conflicts_then_writes() { // A generation-unavailable snapshot refreshes once before its CAS. That // refresh must not spend a CAS attempt, so a withdrawal tombstone still // persists after four conflicts and a successful fifth write. let graph = KvIdentityGraph::new(ConflictInjectingEcKv::new(4, false)); let ec_id = snapshot_ec_id(); - graph.create(&ec_id, &live_entry()).expect("should seed"); + graph + .create(&ec_id, &live_entry()) + .await + .expect("should seed"); let snapshot = EcKvSnapshot::Present { ec_id: ec_id.clone(), entry: Box::new(live_entry()), generation: None, }; - let outcome = graph.tombstone_existing_from_snapshot(&ec_id, snapshot); + let outcome = graph + .tombstone_existing_from_snapshot(&ec_id, snapshot) + .await; assert!( outcome @@ -4279,17 +4641,22 @@ mod tests { ); } - #[test] - fn tombstone_existing_from_snapshot_returns_failed_after_cas_exhaustion() { + #[tokio::test(flavor = "current_thread")] + async fn tombstone_existing_from_snapshot_returns_failed_after_cas_exhaustion() { // Every CAS attempt loses its race, so the row stays live with consent // granted while the browser cookie is already cleared. The caller must // see a failure it can report rather than a silent no-op. let graph = KvIdentityGraph::new(ConflictInjectingEcKv::new(MAX_CAS_RETRIES, false)); let ec_id = snapshot_ec_id(); - graph.create(&ec_id, &live_entry()).expect("should seed"); - let snapshot = graph.load_snapshot(&ec_id); + graph + .create(&ec_id, &live_entry()) + .await + .expect("should seed"); + let snapshot = graph.load_snapshot(&ec_id).await; - let outcome = graph.tombstone_existing_from_snapshot(&ec_id, snapshot); + let outcome = graph + .tombstone_existing_from_snapshot(&ec_id, snapshot) + .await; assert!( matches!(outcome, EcKvSnapshot::Failed { .. }), @@ -4297,6 +4664,7 @@ mod tests { ); let (stored, _) = graph .get(&ec_id) + .await .expect("should read store") .expect("row should remain"); assert!( @@ -4305,16 +4673,19 @@ mod tests { ); } - #[test] - fn tombstone_existing_from_snapshot_overrides_concurrent_live_update() { + #[tokio::test(flavor = "current_thread")] + async fn tombstone_existing_from_snapshot_overrides_concurrent_live_update() { let graph = KvIdentityGraph::new(ConflictInjectingEcKv::with_partner_update_on_conflict(1)); let ec_id = snapshot_ec_id(); graph .create(&ec_id, &live_entry()) + .await .expect("should seed live row"); - let snapshot = graph.load_snapshot(&ec_id); + let snapshot = graph.load_snapshot(&ec_id).await; - let outcome = graph.tombstone_existing_from_snapshot(&ec_id, snapshot); + let outcome = graph + .tombstone_existing_from_snapshot(&ec_id, snapshot) + .await; let entry = outcome .entry_for(&ec_id) @@ -4326,19 +4697,23 @@ mod tests { ); } - #[test] - fn upsert_partner_id_rejects_tombstone() { + #[tokio::test(flavor = "current_thread")] + async fn upsert_partner_id_rejects_tombstone() { let graph = KvIdentityGraph::in_memory("test-store"); let ec_id = snapshot_ec_id(); graph .create(&ec_id, &KvEntry::tombstone(1_000)) + .await .expect("should seed tombstone"); - let result = graph.upsert_partner_id(&ec_id, "ssp.example.com", "uid-1"); + let result = graph + .upsert_partner_id(&ec_id, "ssp.example.com", "uid-1") + .await; assert!(result.is_err(), "public upsert should reject a tombstone"); let (stored, _) = graph .get(&ec_id) + .await .expect("should read store") .expect("should preserve tombstone"); assert!(!stored.consent.ok, "entry should remain withdrawn"); @@ -4348,8 +4723,8 @@ mod tests { ); } - #[test] - fn tombstone_existing_from_snapshot_store_failure_returns_failed() { + #[tokio::test(flavor = "current_thread")] + async fn tombstone_existing_from_snapshot_store_failure_returns_failed() { let graph = KvIdentityGraph::new(WriteFailingEcKv::new()); let ec_id = snapshot_ec_id(); let snapshot = EcKvSnapshot::Present { @@ -4358,7 +4733,9 @@ mod tests { generation: Some(1), }; - let outcome = graph.tombstone_existing_from_snapshot(&ec_id, snapshot); + let outcome = graph + .tombstone_existing_from_snapshot(&ec_id, snapshot) + .await; assert!( matches!(outcome, EcKvSnapshot::Failed { .. }), @@ -4366,41 +4743,51 @@ mod tests { ); } - #[test] - fn tombstone_existing_from_snapshot_noop_when_row_disappears_on_retry() { + #[tokio::test(flavor = "current_thread")] + async fn tombstone_existing_from_snapshot_noop_when_row_disappears_on_retry() { let store = DisappearOnConflictEcKv::new(1); - store.seed_live(&snapshot_ec_id()); + store.seed_live(&snapshot_ec_id()).await; let graph = KvIdentityGraph::new(store); let ec_id = snapshot_ec_id(); - let snapshot = graph.load_snapshot(&ec_id); + let snapshot = graph.load_snapshot(&ec_id).await; - let outcome = graph.tombstone_existing_from_snapshot(&ec_id, snapshot); + let outcome = graph + .tombstone_existing_from_snapshot(&ec_id, snapshot) + .await; assert!( matches!(outcome, EcKvSnapshot::Missing { .. }), "a row that disappears during retry becomes a no-op" ); assert!( - graph.get(&ec_id).expect("should read store").is_none(), + graph + .get(&ec_id) + .await + .expect("should read store") + .is_none(), "must not recreate the disappeared key" ); } - #[test] - fn tombstone_existing_from_snapshot_reretries_failed_snapshot_read() { + #[tokio::test(flavor = "current_thread")] + async fn tombstone_existing_from_snapshot_reretries_failed_snapshot_read() { // A prior request-scoped read failed, so the snapshot is `Failed`. A // withdrawal must not silently drop consent removal: re-read the store // and tombstone the row if it is authoritatively present. let kv = KvIdentityGraph::in_memory("test_store"); let ec_id = snapshot_ec_id(); - kv.create(&ec_id, &live_entry()).expect("should seed live"); + kv.create(&ec_id, &live_entry()) + .await + .expect("should seed live"); - let outcome = kv.tombstone_existing_from_snapshot( - &ec_id, - EcKvSnapshot::Failed { - ec_id: ec_id.clone(), - }, - ); + let outcome = kv + .tombstone_existing_from_snapshot( + &ec_id, + EcKvSnapshot::Failed { + ec_id: ec_id.clone(), + }, + ) + .await; assert!( outcome @@ -4410,26 +4797,33 @@ mod tests { ); let (stored, _) = kv .get(&ec_id) + .await .expect("should read store") .expect("should preserve existing key"); assert!(!stored.consent.ok, "withdrawal must reach the store"); } - #[test] - fn key_exists_confirmed_distinguishes_absence_from_a_stale_point_read() { + #[tokio::test(flavor = "current_thread")] + async fn key_exists_confirmed_distinguishes_absence_from_a_stale_point_read() { let graph = KvIdentityGraph::stale_lookup("stale-store", 1); let ec_id = snapshot_ec_id(); graph .create(&ec_id, &live_entry()) + .await .expect("should seed live"); assert!( - graph.get(&ec_id).expect("should read store").is_none(), + graph + .get(&ec_id) + .await + .expect("should read store") + .is_none(), "the first point read is stale by construction" ); assert!( graph .key_exists_confirmed(&ec_id) + .await .expect("should list the store"), "the list API must still see a row the point read missed" ); @@ -4438,13 +4832,14 @@ mod tests { assert!( !absent .key_exists_confirmed(&ec_id) + .await .expect("should list the store"), "an empty store must prove absence" ); } - #[test] - fn snapshot_upsert_keeps_add_confirmed_present_when_refresh_misses() { + #[tokio::test(flavor = "current_thread")] + async fn snapshot_upsert_keeps_add_confirmed_present_when_refresh_misses() { // `generate_if_needed` records a successful `Add` as `Present` without a // generation. EID ingestion refreshes that snapshot to obtain one; on an // eventually-consistent store the refresh can miss. Enrichment is best @@ -4459,20 +4854,26 @@ mod tests { }; let updates = [PartnerIdUpdate::new("ssp_x", "uid-1")]; - let outcome = graph.upsert_partner_ids_from_snapshot(&ec_id, &updates, add_confirmed); + let outcome = graph + .upsert_partner_ids_from_snapshot(&ec_id, &updates, add_confirmed) + .await; assert!( outcome.entry_for(&ec_id).is_some(), "an Add-confirmed row must survive a non-authoritative refresh miss" ); assert!( - graph.get(&ec_id).expect("should read store").is_none(), + graph + .get(&ec_id) + .await + .expect("should read store") + .is_none(), "a missed refresh must not create or overwrite a root" ); } - #[test] - fn snapshot_upsert_keeps_add_confirmed_present_when_refresh_fails() { + #[tokio::test(flavor = "current_thread")] + async fn snapshot_upsert_keeps_add_confirmed_present_when_refresh_fails() { let graph = KvIdentityGraph::failing("failing-store"); let ec_id = snapshot_ec_id(); let add_confirmed = EcKvSnapshot::Present { @@ -4482,7 +4883,9 @@ mod tests { }; let updates = [PartnerIdUpdate::new("ssp_x", "uid-1")]; - let outcome = graph.upsert_partner_ids_from_snapshot(&ec_id, &updates, add_confirmed); + let outcome = graph + .upsert_partner_ids_from_snapshot(&ec_id, &updates, add_confirmed) + .await; assert!( outcome.entry_for(&ec_id).is_some(), @@ -4490,16 +4893,17 @@ mod tests { ); } - #[test] - fn snapshot_upsert_without_proof_still_reports_a_refresh_miss() { + #[tokio::test(flavor = "current_thread")] + async fn snapshot_upsert_without_proof_still_reports_a_refresh_miss() { // No prior proof of existence: a `NotRead` snapshot that refreshes into // a miss must stay `Missing` so finalization can run orphan recovery. let graph = KvIdentityGraph::in_memory("empty-store"); let ec_id = snapshot_ec_id(); let updates = [PartnerIdUpdate::new("ssp_x", "uid-1")]; - let outcome = - graph.upsert_partner_ids_from_snapshot(&ec_id, &updates, EcKvSnapshot::NotRead); + let outcome = graph + .upsert_partner_ids_from_snapshot(&ec_id, &updates, EcKvSnapshot::NotRead) + .await; assert!( matches!(outcome, EcKvSnapshot::Missing { .. }), @@ -4507,8 +4911,8 @@ mod tests { ); } - #[test] - fn a_store_error_never_carries_the_whole_identifier() { + #[tokio::test(flavor = "current_thread")] + async fn a_store_error_never_carries_the_whole_identifier() { // Every message in this module goes through `log_id`, so a report that // reaches a log cannot disclose the identifier it is about. let kv = KvIdentityGraph::failing("test_store"); @@ -4516,6 +4920,7 @@ mod tests { let report = kv .create(&ec_id, &live_entry()) + .await .expect_err("the failing store should error"); let rendered = format!("{report:?}"); @@ -4525,8 +4930,8 @@ mod tests { ); } - #[test] - fn a_locally_built_error_never_carries_the_whole_identifier() { + #[tokio::test(flavor = "current_thread")] + async fn a_locally_built_error_never_carries_the_whole_identifier() { // The injected-failure case above covers errors the backend produces. // These are built in this module from the identifier itself, on every // path a request can reach: a duplicate create, single upserts naming a @@ -4534,22 +4939,28 @@ mod tests { // terminal errors. let kv = KvIdentityGraph::in_memory("test_store"); let ec_id = format!("{}.ABC123", "a".repeat(64)); - kv.create(&ec_id, &live_entry()).expect("should create"); + kv.create(&ec_id, &live_entry()) + .await + .expect("should create"); let duplicate = kv .create(&ec_id, &live_entry()) + .await .expect_err("a second create should be refused"); let missing = kv .upsert_partner_id(&format!("{}.ZZZ999", "b".repeat(64)), "partner", "uid") + .await .expect_err("an upsert on a missing key should be refused"); let withdrawn = { assert_eq!( kv.write_withdrawal_tombstone(&ec_id, drop) + .await .expect("should tombstone"), TombstoneOutcome::Written, "should tombstone the seeded identity" ); kv.upsert_partner_id(&ec_id, "partner", "uid") + .await .expect_err("an upsert on a withdrawn key should be refused") }; @@ -4557,23 +4968,26 @@ mod tests { // and a store that never lets a write land is the only way to reach them. let cas_revive = { let store = ConflictInjectingEcKv::new(MAX_CAS_RETRIES + 1, false); - store.seed_tombstone(&ec_id); + store.seed_tombstone(&ec_id).await; KvIdentityGraph::new(store) .create_or_revive(&ec_id, &live_entry()) + .await .expect_err("should exhaust CAS retries") }; let cas_upsert = { let store = ConflictInjectingEcKv::new(MAX_CAS_RETRIES + 1, false); - store.seed_live(&ec_id); + store.seed_live(&ec_id).await; KvIdentityGraph::new(store) .upsert_partner_id(&ec_id, "partner", "uid") + .await .expect_err("should exhaust CAS retries") }; let cas_if_exists = { let store = ConflictInjectingEcKv::new(MAX_CAS_RETRIES + 1, false); - store.seed_live(&ec_id); + store.seed_live(&ec_id).await; KvIdentityGraph::new(store) .upsert_partner_id_if_exists(&ec_id, "partner", "uid") + .await .expect_err("should exhaust CAS retries") }; @@ -4593,25 +5007,26 @@ mod tests { } } - #[test] - fn write_withdrawal_tombstone_ignores_an_identity_the_store_does_not_hold() { + #[tokio::test(flavor = "current_thread")] + async fn write_withdrawal_tombstone_ignores_an_identity_the_store_does_not_hold() { let kv = KvIdentityGraph::in_memory("test_store"); let ec_id = format!("{}.ABC123", "b".repeat(64)); assert_eq!( kv.write_withdrawal_tombstone(&ec_id, drop) + .await .expect("should resolve the withdrawal"), TombstoneOutcome::UnknownIdentity, "an identity that was never issued has nothing to withdraw" ); assert!( - kv.get(&ec_id).expect("should read back").is_none(), + kv.get(&ec_id).await.expect("should read back").is_none(), "should not create a row for an identity the store never held" ); } - #[test] - fn withdrawing_many_unheld_identities_creates_no_rows() { + #[tokio::test(flavor = "current_thread")] + async fn withdrawing_many_unheld_identities_creates_no_rows() { let kv = KvIdentityGraph::in_memory("test_store"); let hash = "c".repeat(64); @@ -4620,6 +5035,7 @@ mod tests { for suffix in ["aaaaaa", "bbbbbb", "cccccc", "dddddd"] { assert_eq!( kv.write_withdrawal_tombstone(&format!("{hash}.{suffix}"), drop) + .await .expect("should resolve the withdrawal"), TombstoneOutcome::UnknownIdentity, "suffix `{suffix}` was never issued" @@ -4628,41 +5044,46 @@ mod tests { assert_eq!( kv.count_hash_prefix_keys(&hash) + .await .expect("should count the prefix"), 0, "should hold no rows under a hash nothing was issued for" ); } - #[test] - fn withdrawal_does_not_lose_a_new_identity_to_lookup_lag() { + #[tokio::test(flavor = "current_thread")] + async fn withdrawal_does_not_lose_a_new_identity_to_lookup_lag() { let operations = Arc::new(RecordedEcKvOperations::default()); let kv = KvIdentityGraph::new(RecordingEcKv::with_lagging_live_lookups(operations)); let ec_id = format!("{}.ABC123", "a".repeat(64)); kv.create(&ec_id, &live_entry()) + .await .expect("should issue an identity"); assert!( kv.lookup_raw(&ec_id) + .await .expect("should read lagging replica") .is_none(), "should model the issuance replication gap" ); assert_eq!( kv.write_withdrawal_tombstone(&ec_id, drop) + .await .expect("should process withdrawal"), TombstoneOutcome::Written, "should tombstone an issued identity despite a lagging lookup" ); assert_eq!( kv.upsert_partner_id_if_exists(&ec_id, "partner", "uid") + .await .expect("should check batch-sync eligibility"), UpsertResult::ConsentWithdrawn, "should not revive the withdrawn identity through batch sync" ); } - #[test] - fn a_withdrawal_checks_strong_existence_once_without_eventual_lookup() { + #[tokio::test(flavor = "current_thread")] + async fn a_withdrawal_checks_strong_existence_once_without_eventual_lookup() { // One existence operation may page at the backend, but must not be // followed by an eventually consistent lookup. let operations = Arc::new(RecordedEcKvOperations::default()); @@ -4671,6 +5092,7 @@ mod tests { let absent = format!("{}.ABC123", "e".repeat(64)); assert_eq!( kv.write_withdrawal_tombstone(&absent, drop) + .await .expect("should resolve the withdrawal"), TombstoneOutcome::UnknownIdentity, "an absent identity is not held" @@ -4682,10 +5104,13 @@ mod tests { ); let held = format!("{}.ABC123", "a".repeat(64)); - kv.create(&held, &live_entry()).expect("should create"); + kv.create(&held, &live_entry()) + .await + .expect("should create"); let before = operations.exact_check_count(); assert_eq!( kv.write_withdrawal_tombstone(&held, drop) + .await .expect("should resolve the withdrawal"), TombstoneOutcome::Written, "a held identity is tombstoned" @@ -4702,20 +5127,23 @@ mod tests { ); } - #[test] - fn write_withdrawal_tombstone_refuses_an_empty_identifier() { + #[tokio::test(flavor = "current_thread")] + async fn write_withdrawal_tombstone_refuses_an_empty_identifier() { let kv = KvIdentityGraph::in_memory("test_store"); kv.create(&format!("{}.ABC123", "f".repeat(64)), &live_entry()) + .await .expect("should create"); assert_eq!( kv.write_withdrawal_tombstone("", drop) + .await .expect("should resolve the withdrawal"), TombstoneOutcome::UnknownIdentity, "an empty identifier names no key and must not withdraw anything" ); let (held, _) = kv .get(&format!("{}.ABC123", "f".repeat(64))) + .await .expect("should read back") .expect("should still hold the identity"); assert!( @@ -4737,47 +5165,51 @@ mod tests { } } + #[async_trait::async_trait(?Send)] impl EcKvStore for ReadFailingEcKv { fn store_name(&self) -> &str { self.inner.store_name() } - fn lookup(&self, _key: &str) -> Result, Report> { + async fn lookup( + &self, + _key: &str, + ) -> Result, Report> { Err(Report::new(TrustedServerError::KvStore { store_name: "test_store".to_owned(), message: "reads unavailable".to_owned(), })) } - fn key_exists(&self, key: &str) -> Result> { - self.lookup(key).map(|entry| entry.is_some()) + async fn key_exists(&self, key: &str) -> Result> { + self.lookup(key).await.map(|entry| entry.is_some()) } - fn insert( + async fn insert( &self, key: &str, write: EcKvWrite<'_>, ) -> Result> { - self.inner.insert(key, write) + self.inner.insert(key, write).await } // Left working so a test can prove no row was created without going // through the read path it just made fail. - fn list_keys_with_prefix( + async fn list_keys_with_prefix( &self, prefix: &str, limit: u32, ) -> Result, Report> { - self.inner.list_keys_with_prefix(prefix, limit) + self.inner.list_keys_with_prefix(prefix, limit).await } - fn delete(&self, key: &str) -> Result<(), Report> { - self.inner.delete(key) + async fn delete(&self, key: &str) -> Result<(), Report> { + self.inner.delete(key).await } } - #[test] - fn a_failing_check_is_not_a_way_to_write_for_an_identity_that_was_never_issued() { + #[tokio::test(flavor = "current_thread")] + async fn a_failing_check_is_not_a_way_to_write_for_an_identity_that_was_never_issued() { // The caller controls the identifier and can drive load, so a store // failure must not become a route to the write this gate exists to // prevent. @@ -4786,20 +5218,21 @@ mod tests { let ec_id = format!("{hash}.ABC123"); assert!( - kv.write_withdrawal_tombstone(&ec_id, drop).is_err(), + kv.write_withdrawal_tombstone(&ec_id, drop).await.is_err(), "a check that cannot answer is a fault, so a caller inspecting only \ the error case still reports it" ); assert_eq!( kv.count_hash_prefix_keys(&hash) + .await .expect("should count the prefix"), 0, "should not create a row while the store is degraded" ); } - #[test] - fn a_caller_that_only_inspects_the_error_case_still_sees_a_failed_check() { + #[tokio::test(flavor = "current_thread")] + async fn a_caller_that_only_inspects_the_error_case_still_sees_a_failed_check() { // The withdrawal call site is edited by more than one branch. Reporting // a failed check through `Err` means the common // `if let Err(..) = ...` shape cannot discard it, where a third `Ok` @@ -4808,51 +5241,57 @@ mod tests { let ec_id = format!("{}.ABC123", "6".repeat(64)); let mut reported = false; - if let Err(_err) = kv.write_withdrawal_tombstone(&ec_id, drop) { + if let Err(_err) = kv.write_withdrawal_tombstone(&ec_id, drop).await { reported = true; } assert!(reported, "a failed check must reach an error-only caller"); } - #[test] - fn a_longer_key_does_not_answer_for_the_identity_it_starts_with() { + #[tokio::test(flavor = "current_thread")] + async fn a_longer_key_does_not_answer_for_the_identity_it_starts_with() { let kv = KvIdentityGraph::in_memory("test_store"); let ec_id = format!("{}.ABC123", "7".repeat(64)); // Only a longer key exists. The identity itself was never issued, so a // check that matched by prefix would report it as held and tombstone it. kv.create(&format!("{ec_id}trailing"), &live_entry()) + .await .expect("should create"); assert!( - !kv.key_exists_confirmed(&ec_id).expect("should check"), + !kv.key_exists_confirmed(&ec_id).await.expect("should check"), "a longer key is a different identity" ); assert_eq!( kv.write_withdrawal_tombstone(&ec_id, drop) + .await .expect("should resolve the withdrawal"), TombstoneOutcome::UnknownIdentity, "should not tombstone an identity the store never held" ); assert!( - kv.get(&ec_id).expect("should read back").is_none(), + kv.get(&ec_id).await.expect("should read back").is_none(), "should not create a row via a prefix match" ); } - #[test] - fn key_exists_confirmed_distinguishes_held_identities() { + #[tokio::test(flavor = "current_thread")] + async fn key_exists_confirmed_distinguishes_held_identities() { let kv = KvIdentityGraph::in_memory("test_store"); let held = format!("{}.ABC123", "d".repeat(64)); let sibling = format!("{}.ZZZ999", "d".repeat(64)); - kv.create(&held, &live_entry()).expect("should create"); + kv.create(&held, &live_entry()) + .await + .expect("should create"); assert!( - kv.key_exists_confirmed(&held).expect("should check"), + kv.key_exists_confirmed(&held).await.expect("should check"), "should confirm a held identity" ); assert!( - !kv.key_exists_confirmed(&sibling).expect("should check"), + !kv.key_exists_confirmed(&sibling) + .await + .expect("should check"), "a different suffix under the same hash is a different identity" ); } diff --git a/crates/trusted-server-core/src/ec/kv_backend.rs b/crates/trusted-server-core/src/ec/kv_backend.rs index bfe74ca6f..28fa6c653 100644 --- a/crates/trusted-server-core/src/ec/kv_backend.rs +++ b/crates/trusted-server-core/src/ec/kv_backend.rs @@ -70,6 +70,10 @@ pub enum EcKvWriteOutcome { /// Infrastructure failures are reported as [`TrustedServerError::KvStore`]; /// write precondition failures are part of the normal control flow and are /// returned as [`EcKvWriteOutcome::PreconditionFailed`] instead of errors. +/// +/// I/O returns local futures so native backends can yield without blocking the +/// request runtime. Neither stores nor their futures need to be `Send` or `Sync`. +#[async_trait::async_trait(?Send)] pub trait EcKvStore { /// Returns the platform store name, used in log and error messages. fn store_name(&self) -> &str; @@ -81,7 +85,7 @@ pub trait EcKvStore { /// # Errors /// /// Returns [`TrustedServerError::KvStore`] on store open or read failure. - fn lookup(&self, key: &str) -> Result, Report>; + async fn lookup(&self, key: &str) -> Result, Report>; /// Checks exact-key existence against strongly consistent store state. /// @@ -108,7 +112,7 @@ pub trait EcKvStore { /// Returns [`TrustedServerError::KvStore`] on store failure or when a /// bounded check cannot determine existence. Never falls back to an /// eventually consistent read. - fn key_exists(&self, key: &str) -> Result>; + async fn key_exists(&self, key: &str) -> Result>; /// Writes an entry according to the requested precondition mode. /// @@ -117,7 +121,7 @@ pub trait EcKvStore { /// Returns [`TrustedServerError::KvStore`] on store open or write /// failure. Precondition failures are reported through the /// [`EcKvWriteOutcome`] instead. - fn insert( + async fn insert( &self, key: &str, write: EcKvWrite<'_>, @@ -138,7 +142,7 @@ pub trait EcKvStore { /// # Errors /// /// Returns [`TrustedServerError::KvStore`] on store open or list failure. - fn list_keys_with_prefix( + async fn list_keys_with_prefix( &self, prefix: &str, limit: u32, @@ -153,12 +157,12 @@ pub trait EcKvStore { /// # Errors /// /// Returns [`TrustedServerError::KvStore`] on store open or list failure. - fn count_keys_with_prefix( + async fn count_keys_with_prefix( &self, prefix: &str, limit: u32, ) -> Result> { - let count = self.list_keys_with_prefix(prefix, limit)?.len(); + let count = self.list_keys_with_prefix(prefix, limit).await?.len(); #[allow(clippy::cast_possible_truncation)] Ok(count as u32) } @@ -168,7 +172,7 @@ pub trait EcKvStore { /// # Errors /// /// Returns [`TrustedServerError::KvStore`] on store open or delete failure. - fn delete(&self, key: &str) -> Result<(), Report>; + async fn delete(&self, key: &str) -> Result<(), Report>; } #[cfg(test)] @@ -197,39 +201,43 @@ pub(crate) mod test_support { } } + #[async_trait::async_trait(?Send)] impl EcKvStore for CountingEcKv { fn store_name(&self) -> &str { self.inner.store_name() } - fn lookup(&self, key: &str) -> Result, Report> { + async fn lookup( + &self, + key: &str, + ) -> Result, Report> { self.lookups .fetch_add(1, std::sync::atomic::Ordering::Relaxed); - self.inner.lookup(key) + self.inner.lookup(key).await } - fn key_exists(&self, key: &str) -> Result> { - self.inner.key_exists(key) + async fn key_exists(&self, key: &str) -> Result> { + self.inner.key_exists(key).await } - fn insert( + async fn insert( &self, key: &str, write: EcKvWrite<'_>, ) -> Result> { - self.inner.insert(key, write) + self.inner.insert(key, write).await } - fn list_keys_with_prefix( + async fn list_keys_with_prefix( &self, prefix: &str, limit: u32, ) -> Result, Report> { - self.inner.list_keys_with_prefix(prefix, limit) + self.inner.list_keys_with_prefix(prefix, limit).await } - fn delete(&self, key: &str) -> Result<(), Report> { - self.inner.delete(key) + async fn delete(&self, key: &str) -> Result<(), Report> { + self.inner.delete(key).await } } @@ -258,42 +266,48 @@ pub(crate) mod test_support { } } + #[async_trait::async_trait(?Send)] impl EcKvStore for StaleLookupEcKv { fn store_name(&self) -> &str { self.inner.store_name() } - fn lookup(&self, key: &str) -> Result, Report> { - let mut remaining = self - .stale_lookups_remaining - .lock() - .expect("should lock stale-lookup counter"); - if *remaining > 0 { - *remaining -= 1; - return Ok(None); + async fn lookup( + &self, + key: &str, + ) -> Result, Report> { + { + let mut remaining = self + .stale_lookups_remaining + .lock() + .expect("should lock stale-lookup counter"); + if *remaining > 0 { + *remaining -= 1; + return Ok(None); + } } - self.inner.lookup(key) + self.inner.lookup(key).await } - fn key_exists(&self, key: &str) -> Result> { + async fn key_exists(&self, key: &str) -> Result> { if self.list_fails { return Err(Report::new(TrustedServerError::KvStore { store_name: self.inner.store_name().to_owned(), message: "existence check unavailable".to_owned(), })); } - self.inner.key_exists(key) + self.inner.key_exists(key).await } - fn insert( + async fn insert( &self, key: &str, write: EcKvWrite<'_>, ) -> Result> { - self.inner.insert(key, write) + self.inner.insert(key, write).await } - fn list_keys_with_prefix( + async fn list_keys_with_prefix( &self, prefix: &str, limit: u32, @@ -304,11 +318,11 @@ pub(crate) mod test_support { message: "list unavailable".to_owned(), })); } - self.inner.list_keys_with_prefix(prefix, limit) + self.inner.list_keys_with_prefix(prefix, limit).await } - fn delete(&self, key: &str) -> Result<(), Report> { - self.inner.delete(key) + async fn delete(&self, key: &str) -> Result<(), Report> { + self.inner.delete(key).await } } @@ -333,12 +347,16 @@ pub(crate) mod test_support { } } + #[async_trait::async_trait(?Send)] impl EcKvStore for InMemoryEcKv { fn store_name(&self) -> &str { &self.name } - fn lookup(&self, key: &str) -> Result, Report> { + async fn lookup( + &self, + key: &str, + ) -> Result, Report> { let entries = self.entries.lock().expect("should lock in-memory store"); Ok(entries.get(key).map(|stored| EcKvLookup { body: stored.body.clone(), @@ -347,12 +365,12 @@ pub(crate) mod test_support { })) } - fn key_exists(&self, key: &str) -> Result> { + async fn key_exists(&self, key: &str) -> Result> { let entries = self.entries.lock().expect("should lock in-memory store"); Ok(entries.contains_key(key)) } - fn insert( + async fn insert( &self, key: &str, write: EcKvWrite<'_>, @@ -385,7 +403,7 @@ pub(crate) mod test_support { Ok(EcKvWriteOutcome::Written) } - fn list_keys_with_prefix( + async fn list_keys_with_prefix( &self, prefix: &str, limit: u32, @@ -399,7 +417,7 @@ pub(crate) mod test_support { .collect()) } - fn delete(&self, key: &str) -> Result<(), Report> { + async fn delete(&self, key: &str) -> Result<(), Report> { let mut entries = self.entries.lock().expect("should lock in-memory store"); entries.remove(key); Ok(()) @@ -425,21 +443,25 @@ pub(crate) mod test_support { } } + #[async_trait::async_trait(?Send)] impl EcKvStore for FailingEcKv { fn store_name(&self) -> &str { &self.name } - fn lookup(&self, _key: &str) -> Result, Report> { + async fn lookup( + &self, + _key: &str, + ) -> Result, Report> { Err(self.error("lookup")) } - fn key_exists(&self, key: &str) -> Result> { + async fn key_exists(&self, key: &str) -> Result> { let _ = key; Err(self.error("key_exists")) } - fn insert( + async fn insert( &self, _key: &str, _write: EcKvWrite<'_>, @@ -447,7 +469,7 @@ pub(crate) mod test_support { Err(self.error("insert")) } - fn list_keys_with_prefix( + async fn list_keys_with_prefix( &self, _prefix: &str, _limit: u32, @@ -455,7 +477,7 @@ pub(crate) mod test_support { Err(self.error("list")) } - fn delete(&self, _key: &str) -> Result<(), Report> { + async fn delete(&self, _key: &str) -> Result<(), Report> { Err(self.error("delete")) } } diff --git a/crates/trusted-server-core/src/ec/mod.rs b/crates/trusted-server-core/src/ec/mod.rs index f24c489ed..e90623b2a 100644 --- a/crates/trusted-server-core/src/ec/mod.rs +++ b/crates/trusted-server-core/src/ec/mod.rs @@ -380,7 +380,7 @@ impl EcContext { /// /// Returns an error if the client IP is unavailable and generation is /// needed, or if HMAC generation fails. - pub fn generate_if_needed( + pub async fn generate_if_needed( &mut self, settings: &Settings, kv: Option<&KvIdentityGraph>, @@ -420,7 +420,7 @@ impl EcContext { .map(DeviceSignals::to_kv_device); if let Some(graph) = kv { - match graph.create_if_absent(&ec_id, &entry) { + match graph.create_if_absent(&ec_id, &entry).await { Ok(CreateIfAbsentOutcome::Written) => { self.kv_snapshot = EcKvSnapshot::Present { ec_id: ec_id.clone(), @@ -789,18 +789,22 @@ mod tests { } } + #[async_trait::async_trait(?Send)] impl EcKvStore for AddCollidingEcKv { fn store_name(&self) -> &str { self.inner.store_name() } - fn lookup(&self, key: &str) -> Result, Report> { - self.inner.lookup(key) + async fn lookup( + &self, + key: &str, + ) -> Result, Report> { + self.inner.lookup(key).await } - fn key_exists(&self, key: &str) -> Result> { - self.inner.key_exists(key) + async fn key_exists(&self, key: &str) -> Result> { + self.inner.key_exists(key).await } - fn insert( + async fn insert( &self, key: &str, write: EcKvWrite<'_>, @@ -815,17 +819,17 @@ mod tests { return Ok(EcKvWriteOutcome::PreconditionFailed); } } - self.inner.insert(key, write) + self.inner.insert(key, write).await } - fn list_keys_with_prefix( + async fn list_keys_with_prefix( &self, prefix: &str, limit: u32, ) -> Result, Report> { - self.inner.list_keys_with_prefix(prefix, limit) + self.inner.list_keys_with_prefix(prefix, limit).await } - fn delete(&self, key: &str) -> Result<(), Report> { - self.inner.delete(key) + async fn delete(&self, key: &str) -> Result<(), Report> { + self.inner.delete(key).await } } @@ -837,14 +841,15 @@ mod tests { } } - #[test] - fn generate_if_needed_retries_id_collision_then_persists() { + #[tokio::test(flavor = "current_thread")] + async fn generate_if_needed_retries_id_collision_then_persists() { let settings = create_test_settings(); let mut ec = EcContext::new_for_test_with_ip(None, granting_consent(), Some("192.0.2.5".to_owned())); let graph = KvIdentityGraph::new(AddCollidingEcKv::new(2)); ec.generate_if_needed(&settings, Some(&graph)) + .await .expect("should generate after bounded collisions"); assert!(ec.ec_value().is_some(), "should allocate a fresh EC ID"); @@ -855,15 +860,15 @@ mod tests { ); } - #[test] - fn generate_if_needed_errors_after_collision_exhaustion() { + #[tokio::test(flavor = "current_thread")] + async fn generate_if_needed_errors_after_collision_exhaustion() { let settings = create_test_settings(); let mut ec = EcContext::new_for_test_with_ip(None, granting_consent(), Some("192.0.2.6".to_owned())); // Collide on every attempt so the bounded retry is exhausted. let graph = KvIdentityGraph::new(AddCollidingEcKv::new(u32::MAX)); - let result = ec.generate_if_needed(&settings, Some(&graph)); + let result = ec.generate_if_needed(&settings, Some(&graph)).await; assert!(result.is_err(), "should fail after exhausting attempts"); assert!( @@ -1075,8 +1080,8 @@ mod tests { ); } - #[test] - fn generate_if_needed_skips_when_ec_exists() { + #[tokio::test(flavor = "current_thread")] + async fn generate_if_needed_skips_when_ec_exists() { let settings = create_test_settings(); let ec_id = valid_ec_id("d", "Exist1"); let cookie = format!("ts-ec={ec_id}"); @@ -1085,6 +1090,7 @@ mod tests { let mut ec = EcContext::read_from_request(&settings, &req, &noop_services()) .expect("should read EC context"); ec.generate_if_needed(&settings, None) + .await .expect("should not error when EC already exists"); assert_eq!( diff --git a/crates/trusted-server-core/src/ec/pull_sync.rs b/crates/trusted-server-core/src/ec/pull_sync.rs index dce473ab9..60131edf1 100644 --- a/crates/trusted-server-core/src/ec/pull_sync.rs +++ b/crates/trusted-server-core/src/ec/pull_sync.rs @@ -173,7 +173,7 @@ pub async fn dispatch_pull_sync( return; } - let live_snapshot = kv.load_snapshot(context.ec_id()); + let live_snapshot = kv.load_snapshot(context.ec_id()).await; let Some(live_entry) = live_snapshot.entry_for(context.ec_id()) else { log::warn!( "Pull sync: skipping dispatch for '{}' because the live identity row could not be \ @@ -261,8 +261,9 @@ pub async fn dispatch_pull_sync( // Write back from the revalidated snapshot, not the request one: its // generation is current, so the first CAS attempt is not spent losing a // conflict against the read that authorized this dispatch. - let outcome = - kv.upsert_partner_ids_from_snapshot(context.ec_id(), &updates, live_snapshot.clone()); + let outcome = kv + .upsert_partner_ids_from_snapshot(context.ec_id(), &updates, live_snapshot.clone()) + .await; if matches!(outcome, EcKvSnapshot::Failed { .. }) { log::warn!( "Pull sync: failed to persist partner updates for '{}'", @@ -524,8 +525,8 @@ mod tests { } } - #[test] - fn build_pull_sync_context_returns_context_when_valid() { + #[tokio::test(flavor = "current_thread")] + async fn build_pull_sync_context_returns_context_when_valid() { let consent = ConsentContext { jurisdiction: crate::consent::jurisdiction::Jurisdiction::NonRegulated, ..ConsentContext::default() @@ -533,7 +534,7 @@ mod tests { let ec_id = format!("{}.ABC123", "a".repeat(64)); let mut ec_context = EcContext::new_for_test(Some(ec_id.clone()), consent); let graph = KvIdentityGraph::in_memory("pull_store"); - ec_context.set_kv_snapshot(seed_present_snapshot(&graph, &ec_id)); + ec_context.set_kv_snapshot(seed_present_snapshot(&graph, &ec_id).await); let registry = PartnerRegistry::from_config(&[pull_enabled_ec_partner("ssp.example.com")]) .expect("should build registry"); @@ -563,8 +564,8 @@ mod tests { ); } - #[test] - fn build_pull_sync_context_skips_when_eid_persistence_is_denied() { + #[tokio::test(flavor = "current_thread")] + async fn build_pull_sync_context_skips_when_eid_persistence_is_denied() { // Purpose 1 granted keeps the EC, but Purpose 4 denied withholds EID // writes and leaves partner slots empty. Pull sync must not read those // empty slots as eligible and disclose the EC ID to partners. @@ -572,7 +573,7 @@ mod tests { let ec_id = format!("{}.ABC123", "a".repeat(64)); let mut ec_context = EcContext::new_for_test(Some(ec_id.clone()), consent); let graph = KvIdentityGraph::in_memory("pull_store"); - ec_context.set_kv_snapshot(seed_present_snapshot(&graph, &ec_id)); + ec_context.set_kv_snapshot(seed_present_snapshot(&graph, &ec_id).await); let registry = PartnerRegistry::from_config(&[pull_enabled_ec_partner("ssp.example.com")]) .expect("should build registry"); assert!( @@ -588,8 +589,8 @@ mod tests { ); } - #[test] - fn build_pull_sync_context_skips_empty_registry_and_complete_snapshot() { + #[tokio::test(flavor = "current_thread")] + async fn build_pull_sync_context_skips_empty_registry_and_complete_snapshot() { let consent = ConsentContext { jurisdiction: crate::consent::jurisdiction::Jurisdiction::NonRegulated, ..ConsentContext::default() @@ -597,7 +598,7 @@ mod tests { let ec_id = format!("{}.ABC123", "a".repeat(64)); let mut ec_context = EcContext::new_for_test(Some(ec_id.clone()), consent); let graph = KvIdentityGraph::in_memory("pull_store"); - let mut snapshot = seed_present_snapshot(&graph, &ec_id); + let mut snapshot = seed_present_snapshot(&graph, &ec_id).await; let registry = PartnerRegistry::from_config(&[pull_enabled_ec_partner("ssp.example.com")]) .expect("should build registry"); @@ -861,11 +862,14 @@ mod tests { format!("https://{source_domain}/sync?ec_id={}", snapshot_ec_id()) } - fn seed_present_snapshot(graph: &KvIdentityGraph, ec_id: &str) -> EcKvSnapshot { + async fn seed_present_snapshot(graph: &KvIdentityGraph, ec_id: &str) -> EcKvSnapshot { let mut entry = KvEntry::tombstone(1000); entry.consent.ok = true; - graph.create(ec_id, &entry).expect("should seed live entry"); - graph.load_snapshot(ec_id) + graph + .create(ec_id, &entry) + .await + .expect("should seed live entry"); + graph.load_snapshot(ec_id).await } /// Runs pull sync to completion on `clock`, advancing it as the requests wait. @@ -889,7 +893,57 @@ mod tests { } #[test] - fn dispatch_pull_sync_aggregates_batches_into_one_bulk_write() { + fn dispatch_pull_sync_runs_inside_the_fastly_executor_boundary() { + futures::executor::block_on(async { + let settings = create_test_settings(); + let registry = + PartnerRegistry::from_config(&[pull_enabled_ec_partner("alpha.example.com")]) + .expect("should build registry"); + let graph = KvIdentityGraph::in_memory("pull_store"); + let ec_id = snapshot_ec_id(); + let snapshot = seed_present_snapshot(&graph, &ec_id).await; + let clock = TestClock::new(); + let client = ScriptedHttpClient::new(&clock); + client.respond_to( + Method::GET, + &pull_url("alpha.example.com"), + ScriptedReply::status(StatusCode::OK).body(r#"{"uid":"synced-uid"}"#), + ); + let context = PullSyncContext { + ec_id: ec_id.clone(), + snapshot, + }; + + dispatch_pull_sync( + &settings, + &graph, + ®istry, + &AllowAllRateLimiter, + &context, + &client.http_client(), + &clock.monotonic_clock(), + ) + .await; + + assert_eq!( + client.requests().len(), + 1, + "should dispatch once through the direct client" + ); + let (entry, _) = graph + .get(&ec_id) + .await + .expect("should read row") + .expect("row should exist"); + assert_eq!( + entry.ids.get("alpha.example.com").map(|id| id.uid.as_str()), + Some("synced-uid"), + ); + }); + } + + #[tokio::test(flavor = "current_thread")] + async fn dispatch_pull_sync_aggregates_batches_into_one_bulk_write() { let mut settings = create_test_settings(); // Force one partner per concurrency batch so responses span batches. settings.ec.pull_sync_concurrency = 1; @@ -901,7 +955,7 @@ mod tests { let graph = KvIdentityGraph::in_memory("pull_store"); let ec_id = snapshot_ec_id(); - let snapshot = seed_present_snapshot(&graph, &ec_id); + let snapshot = seed_present_snapshot(&graph, &ec_id).await; let clock = TestClock::new(); let client = ScriptedHttpClient::new(&clock); @@ -922,6 +976,7 @@ mod tests { let (entry, generation) = graph .get(&ec_id) + .await .expect("should read store") .expect("should keep the entry"); assert_eq!( @@ -940,8 +995,8 @@ mod tests { ); } - #[test] - fn dispatch_pull_sync_sends_each_partner_its_bearer_token_and_transport() { + #[tokio::test(flavor = "current_thread")] + async fn dispatch_pull_sync_sends_each_partner_its_bearer_token_and_transport() { let mut settings = create_test_settings(); settings.ec.pull_sync_concurrency = 4; let registry = @@ -949,7 +1004,7 @@ mod tests { .expect("should build pull registry"); let graph = KvIdentityGraph::in_memory("pull_store"); let ec_id = snapshot_ec_id(); - let snapshot = seed_present_snapshot(&graph, &ec_id); + let snapshot = seed_present_snapshot(&graph, &ec_id).await; let clock = TestClock::new(); let client = ScriptedHttpClient::new(&clock); client.respond_to( @@ -982,8 +1037,8 @@ mod tests { }); } - #[test] - fn pull_sync_group_ends_at_fifteen_second_cutoff() { + #[tokio::test(flavor = "current_thread")] + async fn pull_sync_group_ends_at_fifteen_second_cutoff() { let mut settings = create_test_settings(); settings.ec.pull_sync_concurrency = 2; let registry = PartnerRegistry::from_config(&[ @@ -993,7 +1048,7 @@ mod tests { .expect("should build pull registry"); let graph = KvIdentityGraph::in_memory("pull_store"); let ec_id = snapshot_ec_id(); - let snapshot = seed_present_snapshot(&graph, &ec_id); + let snapshot = seed_present_snapshot(&graph, &ec_id).await; let clock = TestClock::new(); let started = clock.monotonic_clock().now(); let client = ScriptedHttpClient::new(&clock); @@ -1040,6 +1095,7 @@ mod tests { } let (entry, _) = graph .get(&ec_id) + .await .expect("should read store") .expect("should keep the entry"); assert_eq!( @@ -1053,8 +1109,8 @@ mod tests { ); } - #[test] - fn dispatch_pull_sync_drops_a_response_over_the_pull_cap() { + #[tokio::test(flavor = "current_thread")] + async fn dispatch_pull_sync_drops_a_response_over_the_pull_cap() { let mut settings = create_test_settings(); settings.ec.pull_sync_concurrency = 4; let registry = @@ -1062,7 +1118,7 @@ mod tests { .expect("should build pull registry"); let graph = KvIdentityGraph::in_memory("pull_store"); let ec_id = snapshot_ec_id(); - let snapshot = seed_present_snapshot(&graph, &ec_id); + let snapshot = seed_present_snapshot(&graph, &ec_id).await; let clock = TestClock::new(); let client = ScriptedHttpClient::new(&clock); let oversized = format!( @@ -1083,6 +1139,7 @@ mod tests { let (entry, generation) = graph .get(&ec_id) + .await .expect("should read store") .expect("should keep the entry"); assert!( @@ -1092,8 +1149,8 @@ mod tests { assert_eq!(generation, 1, "should not write the row back"); } - #[test] - fn dispatch_pull_sync_calls_and_persists_only_missing_partner() { + #[tokio::test(flavor = "current_thread")] + async fn dispatch_pull_sync_calls_and_persists_only_missing_partner() { let mut settings = create_test_settings(); settings.ec.pull_sync_concurrency = 4; let registry = PartnerRegistry::from_config(&[ @@ -1108,8 +1165,9 @@ mod tests { &ec_id, &KvEntry::minimal("alpha.example.com", "existing-alpha", 1_000), ) + .await .expect("should seed partial entry"); - let snapshot = graph.load_snapshot(&ec_id); + let snapshot = graph.load_snapshot(&ec_id).await; let clock = TestClock::new(); let client = ScriptedHttpClient::new(&clock); client.respond_to( @@ -1136,6 +1194,7 @@ mod tests { ); let (entry, generation) = graph .get(&ec_id) + .await .expect("should read store") .expect("should keep the entry"); assert_eq!( @@ -1157,8 +1216,8 @@ mod tests { assert_eq!(generation, 2, "missing UID should use one bulk CAS write"); } - #[test] - fn dispatch_pull_sync_skips_non_present_snapshots() { + #[tokio::test(flavor = "current_thread")] + async fn dispatch_pull_sync_skips_non_present_snapshots() { let mut settings = create_test_settings(); settings.ec.pull_sync_concurrency = 4; let registry = @@ -1187,13 +1246,17 @@ mod tests { client.assert_no_requests(); assert!( - graph.get(&ec_id).expect("should read store").is_none(), + graph + .get(&ec_id) + .await + .expect("should read store") + .is_none(), "no snapshot state should create a missing root" ); } - #[test] - fn dispatch_pull_sync_skips_tombstone_snapshot() { + #[tokio::test(flavor = "current_thread")] + async fn dispatch_pull_sync_skips_tombstone_snapshot() { let mut settings = create_test_settings(); settings.ec.pull_sync_concurrency = 4; let registry = @@ -1218,8 +1281,8 @@ mod tests { client.assert_no_requests(); } - #[test] - fn dispatch_pull_sync_skips_dispatch_when_ec_is_tombstoned_after_snapshot_capture() { + #[tokio::test(flavor = "current_thread")] + async fn dispatch_pull_sync_skips_dispatch_when_ec_is_tombstoned_after_snapshot_capture() { // The request snapshot authorized pull sync, then a concurrent request // completed a CMP withdrawal while the page response was in flight. // Pull sync runs post-send and discloses the raw `ec_id` to partners, so @@ -1232,12 +1295,13 @@ mod tests { .expect("should build registry"); let graph = KvIdentityGraph::in_memory("pull_store"); let ec_id = snapshot_ec_id(); - let snapshot = seed_present_snapshot(&graph, &ec_id); + let snapshot = seed_present_snapshot(&graph, &ec_id).await; // Concurrent withdrawal lands after the snapshot was captured. assert_eq!( graph .write_withdrawal_tombstone(&ec_id, drop) + .await .expect("should tombstone the row"), TombstoneOutcome::Written, "should tombstone the row the snapshot still reports as present" @@ -1255,6 +1319,7 @@ mod tests { client.assert_no_requests(); let (entry, _) = graph .get(&ec_id) + .await .expect("should read store") .expect("should keep the tombstone"); assert!( @@ -1263,8 +1328,8 @@ mod tests { ); } - #[test] - fn dispatch_pull_sync_skips_revalidation_read_when_no_partner_is_eligible() { + #[tokio::test(flavor = "current_thread")] + async fn dispatch_pull_sync_skips_revalidation_read_when_no_partner_is_eligible() { // Every pull-enabled partner already has a UID in the request snapshot, // so there is nothing to dispatch. The revalidation read exists to // authorize outbound calls; with no calls to authorize it must not cost @@ -1286,6 +1351,7 @@ mod tests { ); graph .create(&ec_id, &entry) + .await .expect("should seed live entry"); let snapshot = EcKvSnapshot::Present { ec_id: ec_id.clone(), diff --git a/crates/trusted-server-core/src/integrations/registry.rs b/crates/trusted-server-core/src/integrations/registry.rs index 2936fa8a1..a79a066b5 100644 --- a/crates/trusted-server-core/src/integrations/registry.rs +++ b/crates/trusted-server-core/src/integrations/registry.rs @@ -1039,7 +1039,7 @@ impl IntegrationRegistry { // Only generate for document navigations — subresource requests // may lack consent signals such as the Sec-GPC header. if is_navigation_request(&req) { - if let Err(err) = ec_context.generate_if_needed(settings, kv) { + if let Err(err) = ec_context.generate_if_needed(settings, kv).await { log::warn!("EC generation failed for integration proxy: {err:?}"); } } else { diff --git a/crates/trusted-server-core/src/publisher.rs b/crates/trusted-server-core/src/publisher.rs index c8849d650..b6f7b550d 100644 --- a/crates/trusted-server-core/src/publisher.rs +++ b/crates/trusted-server-core/src/publisher.rs @@ -5439,7 +5439,7 @@ pub async fn handle_publisher_request( })?); } if should_preload_ec && let (Some(graph), Some(active_ec_id)) = (kv, active_ec_id) { - let refreshed = graph.load_snapshot(active_ec_id); + let refreshed = graph.load_snapshot(active_ec_id).await; // Never downgrade an in-request Add-confirmed Present snapshot: a // freshly created row can read back Missing/Failed on an // eventually-consistent store, and rotating or suppressing that @@ -7927,7 +7927,7 @@ pub async fn handle_page_bids( // EIDs. Kill-switch, no-slot, bot/prefetch, and no-registry requests // never reach here, so they incur no billable KV read. let page_bids_kv_snapshot = match (kv, ec_id.as_deref(), auction.registry) { - (Some(graph), Some(ec_id), Some(_)) => graph.load_snapshot(ec_id), + (Some(graph), Some(ec_id), Some(_)) => graph.load_snapshot(ec_id).await, _ => crate::ec::EcKvSnapshot::NotRead, }; // Hand the loaded row to the request context so response @@ -8728,37 +8728,41 @@ mod tests { ok: bool, } + #[async_trait::async_trait(?Send)] impl EcKvStore for OrderRecordingKv { fn store_name(&self) -> &str { self.inner.store_name() } - fn lookup(&self, _key: &str) -> Result, Report> { + async fn lookup( + &self, + _key: &str, + ) -> Result, Report> { self.lookups.fetch_add(1, Ordering::SeqCst); self.http_calls_at_lookup .store(self.http.origin_requests().len(), Ordering::SeqCst); // Report a miss: the scheduling assertions only care about ordering. Ok(None) } - fn key_exists(&self, key: &str) -> Result> { - self.inner.key_exists(key) + async fn key_exists(&self, key: &str) -> Result> { + self.inner.key_exists(key).await } - fn insert( + async fn insert( &self, key: &str, write: EcKvWrite<'_>, ) -> Result> { - self.inner.insert(key, write) + self.inner.insert(key, write).await } - fn list_keys_with_prefix( + async fn list_keys_with_prefix( &self, prefix: &str, limit: u32, ) -> Result, Report> { - self.inner.list_keys_with_prefix(prefix, limit) + self.inner.list_keys_with_prefix(prefix, limit).await } - fn delete(&self, key: &str) -> Result<(), Report> { - self.inner.delete(key) + async fn delete(&self, key: &str) -> Result<(), Report> { + self.inner.delete(key).await } } @@ -15718,6 +15722,7 @@ mod tests { crate::ec::current_timestamp(), ), ) + .await .expect("should seed a live reader identity"); } let registry = IntegrationRegistry::new(&settings) @@ -15809,7 +15814,8 @@ mod tests { None, None, &mut response, - ); + ) + .await; assert_eq!( response.headers()[HEADER_X_TS_TEMPLATE_CACHE], if index == 0 { "miss-stored" } else { "hit" }, @@ -15891,6 +15897,7 @@ mod tests { for (index, identity) in identities.iter().enumerate() { let (entry, _) = graph .get(identity) + .await .expect("should read reader identity") .expect("should retain the live row or its tombstone"); assert_eq!(