Avoid no-op EC KV reads in post-send pull sync

crates/trusted-server-adapter-fastly/src/main.rs

@@ -21,7 +21,6 @@ use trusted_server_core::ec::registry::PartnerRegistry;
 use trusted_server_core::error::TrustedServerError;
 use trusted_server_core::integrations::RequestFilterEffects;
 use trusted_server_core::platform::PlatformGeo as _;
-use trusted_server_core::platform::RuntimeServices;
 use trusted_server_core::proxy::{stream_asset_body, AssetProxyCachePolicy};
 use trusted_server_core::settings::Settings;
 
@@ -312,11 +311,44 @@ fn run_edgezero_pull_sync_after_send(
     partner_registry: &PartnerRegistry,
     ec_state: &EcFinalizeState,
 ) {
-    if ec_state.is_real_browser {
-        if let Some(context) = build_pull_sync_context(&ec_state.ec_context) {
-            run_pull_sync_after_send(settings, partner_registry, &context, &ec_state.services);
-        }
+    if !ec_state.is_real_browser {
+        return;
     }
+
+    let prepared_context = build_pull_sync_context(&ec_state.ec_context, partner_registry);
+    let Some((context, kv)) =
+        prepare_pull_sync_after_send(prepared_context, || require_identity_graph(settings))
+    else {
+        return;
+    };
+
+    let limiter = FastlyRateLimiter::new(RATE_COUNTER_NAME);
+    dispatch_pull_sync(
+        settings,
+        &kv,
+        partner_registry,
+        &limiter,
+        &context,
+        &ec_state.services,
+    );
+}
+
+fn prepare_pull_sync_after_send<F>(
+    context: Option<PullSyncContext>,
+    graph_factory: F,
+) -> Option<(PullSyncContext, KvIdentityGraph)>
+where
+    F: FnOnce() -> Result<KvIdentityGraph, Report<TrustedServerError>>,
+{
+    let context = context?;
+    let kv = match graph_factory() {
+        Ok(kv) => kv,
+        Err(err) => {
+            log::debug!("Pull sync: identity graph unavailable, skipping: {err:?}");
+            return None;
+        }
+    };
+    Some((context, kv))
 }
 
 /// Sends a finalized `EdgeZero` response to the client.
@@ -421,24 +453,6 @@ pub(crate) fn maybe_identity_graph(settings: &Settings) -> Option<KvIdentityGrap
         .map(|store_name| KvIdentityGraph::new(FastlyEcKvStore::new(store_name)))
 }
 
-fn run_pull_sync_after_send(
-    settings: &Settings,
-    partner_registry: &PartnerRegistry,
-    context: &PullSyncContext,
-    services: &RuntimeServices,
-) {
-    let kv = match require_identity_graph(settings) {
-        Ok(kv) => kv,
-        Err(err) => {
-            log::debug!("Pull sync: identity graph unavailable, skipping: {err:?}");
-            return;
-        }
-    };
-
-    let limiter = FastlyRateLimiter::new(RATE_COUNTER_NAME);
-    dispatch_pull_sync(settings, &kv, partner_registry, &limiter, context, services);
-}
-
 /// Constructs a `KvIdentityGraph` from settings, or returns an error if the
 /// `ec_store` config is not set.
 pub(crate) fn require_identity_graph(
@@ -513,6 +527,26 @@ mod tests {
         .expect("should parse test settings")
     }
 
+    #[test]
+    fn pull_sync_noop_states_skip_post_send_graph_factory() {
+        for reason in ["no partners", "complete snapshot", "unread marker state"] {
+            let calls = std::cell::Cell::new(0);
+            let result = prepare_pull_sync_after_send(None, || {
+                calls.set(calls.get() + 1);
+                Err(Report::new(TrustedServerError::KvStore {
+                    store_name: "unexpected".to_owned(),
+                    message: "graph factory should not run".to_owned(),
+                }))
+            });
+            assert!(result.is_none(), "{reason} preparation should return none");
+            assert_eq!(
+                calls.get(),
+                0,
+                "{reason} should not invoke the graph factory"
+            );
+        }
+    }
+
     #[test]
     fn health_response_short_circuits_get_health() {
         let req = FastlyRequest::get("https://example.com/health");

crates/trusted-server-core/src/constants.rs

@@ -1,6 +1,8 @@
 use http::header::HeaderName;
 
 pub const COOKIE_TS_EC: &str = "ts-ec";
+/// Short-lived signed proof that the current EC row has every pull-partner UID.
+pub const COOKIE_TS_EC_PULL_COMPLETE: &str = "ts-ec-pull-complete";
 /// Cookie written by the Trusted Server JS SDK containing a standard-base64-encoded
 /// JSON array of Extended User IDs (`[{ source, uids }]`) from identity providers.
 pub const COOKIE_TS_EIDS: &str = "ts-eids";

crates/trusted-server-core/src/ec/finalize.rs

@@ -16,6 +16,7 @@ use super::generation::{generate_ec_id, is_valid_ec_id};
 use super::kv::{apply_partner_id_updates, CreateIfAbsentOutcome, KvIdentityGraph};
 use super::kv_types::KvEntry;
 use super::prebid_eids::collect_eid_cookie_updates;
+use super::pull_sync_marker::{expire_marker, reconcile_marker};
 use super::registry::PartnerRegistry;
 use super::{current_timestamp, EcContext, EcKvSnapshot};
 
@@ -47,6 +48,7 @@ pub fn ec_finalize_response(
     sharedid_cookie: Option<&str>,
     response: &mut Response<EdgeBody>,
 ) {
+    ec_context.validate_pull_sync_marker(settings, registry);
     let consent_allows_ec = ec_consent_granted(ec_context.consent());
     let consent_withdrawn = ec_consent_withdrawn(ec_context.consent());
 
@@ -57,8 +59,12 @@ pub fn ec_finalize_response(
         // consent input.
         clear_ec_headers_on_response(response, Some(registry));
 
-        // Only expire the browser cookie and tombstone the identity-graph row
-        // when the request carries an explicit withdrawal signal.
+        if consent_withdrawn {
+            expire_marker(ec_context.pull_sync_marker_mut(), response);
+        }
+
+        // Only expire the EC cookie and tombstone the identity-graph row when
+        // explicit withdrawal accompanies an EC cookie.
         if consent_withdrawn && ec_context.cookie_was_present() {
             expire_ec_cookie(settings, response);
 
@@ -102,6 +108,8 @@ pub fn ec_finalize_response(
             }
         }
 
+        reconcile_pull_sync_marker(settings, registry, ec_context, response);
+
         // Ordinary returning-user page views no longer refresh the browser
         // cookie, emit the EC header, or update KV TTL.
         return;
@@ -113,6 +121,7 @@ pub fn ec_finalize_response(
     if ec_context.ec_generated() {
         let (Some(graph), Some(ec_id)) = (kv, ec_context.ec_value().map(str::to_owned)) else {
             log::info!("Skipping generated EC response write because KV graph is unavailable");
+            reconcile_pull_sync_marker(settings, registry, ec_context, response);
             return;
         };
 
@@ -129,6 +138,26 @@ pub fn ec_finalize_response(
             log::warn!("Skipping generated EC cookie because backing row is not authoritative");
         }
     }
+
+    reconcile_pull_sync_marker(settings, registry, ec_context, response);
+}
+
+fn reconcile_pull_sync_marker(
+    settings: &Settings,
+    registry: &PartnerRegistry,
+    ec_context: &mut EcContext,
+    response: &mut Response<EdgeBody>,
+) {
+    let ec_id = ec_context.ec_value().map(str::to_owned);
+    let snapshot = ec_context.kv_snapshot().clone();
+    reconcile_marker(
+        settings,
+        registry,
+        ec_id.as_deref(),
+        &snapshot,
+        ec_context.pull_sync_marker_mut(),
+        response,
+    );
 }
 
 fn recover_orphaned_ec(
@@ -982,6 +1011,104 @@ mod tests {
         );
     }
 
+    #[test]
+    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");
+        partner.pull_sync_enabled = true;
+        partner.pull_sync_url = Some("https://sync.example.com/pull".to_owned());
+        partner.pull_sync_allowed_domains = vec!["sync.example.com".to_owned()];
+        partner.ts_pull_token = Some(Redacted::new("pull-token".to_owned()));
+        let registry = PartnerRegistry::from_config(&[partner]).expect("should build registry");
+        let mut ec_context = make_context(
+            Some(&ec_id),
+            Some(&ec_id),
+            true,
+            false,
+            Jurisdiction::NonRegulated,
+        );
+        let mut entry = live_entry();
+        entry.ids.insert(
+            "ssp.example.com".to_owned(),
+            crate::ec::kv_types::KvPartnerId {
+                uid: "partner-uid".to_owned(),
+            },
+        );
+        ec_context.set_kv_snapshot(EcKvSnapshot::Present {
+            ec_id,
+            entry: Box::new(entry),
+            generation: Some(1),
+        });
+        let mut response = empty_response();
+
+        ec_finalize_response(
+            &settings,
+            &mut ec_context,
+            None,
+            &registry,
+            None,
+            None,
+            &mut response,
+        );
+
+        let cookies = response
+            .headers()
+            .get_all(http::header::SET_COOKIE)
+            .iter()
+            .filter_map(|value| value.to_str().ok())
+            .collect::<Vec<_>>();
+        assert!(
+            cookies
+                .iter()
+                .any(|cookie| cookie.starts_with("ts-ec-pull-complete=v1.")),
+            "complete snapshot should issue the marker"
+        );
+    }
+
+    #[test]
+    fn explicit_withdrawal_expires_marker_without_ec_cookie() {
+        let settings = create_test_settings();
+        let consent = ConsentContext {
+            jurisdiction: Jurisdiction::UsState("CA".to_owned()),
+            gpc: true,
+            source: ConsentSource::Cookie,
+            ..Default::default()
+        };
+        let mut ec_context = make_context_with_consent(None, None, false, false, consent);
+        ec_context.set_pull_sync_marker_for_test(
+            crate::ec::pull_sync_marker::PullSyncMarkerState::Invalid,
+        );
+        let mut response = empty_response();
+
+        ec_finalize_response(
+            &settings,
+            &mut ec_context,
+            None,
+            &PartnerRegistry::empty(),
+            None,
+            None,
+            &mut response,
+        );
+
+        let cookies = response
+            .headers()
+            .get_all(http::header::SET_COOKIE)
+            .iter()
+            .filter_map(|value| value.to_str().ok())
+            .collect::<Vec<_>>();
+        assert!(
+            cookies.iter().any(|cookie| {
+                cookie.starts_with("ts-ec-pull-complete=;") && cookie.contains("Max-Age=0")
+            }),
+            "withdrawal should expire the marker independently of EC cookie state"
+        );
+        assert!(
+            cookies.iter().all(|cookie| !cookie.starts_with("ts-ec=;")),
+            "missing EC cookie should not add an EC-cookie expiry"
+        );
+    }
+
     fn live_entry() -> KvEntry {
         let mut entry = KvEntry::tombstone(1000);
         entry.consent.ok = true;