Return the subscription's invalid_ts from wait_for_invalidation (#39457)

GitOrigin-RevId: 866df98fec22d82643a08030c6e01b0ed000277a

Author: goffrie

crates/application/src/api.rs

@@ -575,10 +575,10 @@ impl<RT: Runtime> SubscriptionClient for ApplicationSubscriptionClient<RT> {
     }
 }
 
-pub enum ExtendValidityResult {
-    /// The subscription's validity can be extended to the requested timestamp.
-    Extended,
-    /// The subscription may no longer be valid at the requested timestamp.
+pub enum SubscriptionValidity {
+    /// The subscription is valid.
+    Valid,
+    /// The subscription may no longer be valid.
     /// This result can be returned spuriously even if there were no conflicting
     /// writes.
     Invalid {
@@ -590,10 +590,14 @@ pub enum ExtendValidityResult {
 
 #[async_trait]
 pub trait SubscriptionTrait: Send + Sync {
-    fn wait_for_invalidation(&self) -> BoxFuture<'static, anyhow::Result<()>>;
-
-    /// See comments on [`ExtendValidityResult`]
-    async fn extend_validity(&self, new_ts: Timestamp) -> anyhow::Result<ExtendValidityResult>;
+    /// Returns a future that completes after the subscription becomes invalid.
+    /// The future yields the invalidation timestamp, if known; this is the same
+    /// thing as [`SubscriptionValidity`]'s `invalid_ts`.
+    fn wait_for_invalidation(&self) -> BoxFuture<'static, anyhow::Result<Option<Timestamp>>>;
+
+    /// Checks if the subscription is still valid as of `new_ts`. See comments
+    /// on [`SubscriptionValidity`].
+    async fn extend_validity(&self, new_ts: Timestamp) -> anyhow::Result<SubscriptionValidity>;
 }
 
 struct ApplicationSubscription {
@@ -608,15 +612,15 @@ struct ApplicationSubscription {
 
 #[async_trait]
 impl SubscriptionTrait for ApplicationSubscription {
-    fn wait_for_invalidation(&self) -> BoxFuture<'static, anyhow::Result<()>> {
+    fn wait_for_invalidation(&self) -> BoxFuture<'static, anyhow::Result<Option<Timestamp>>> {
         self.inner.wait_for_invalidation().map(Ok).boxed()
     }
 
     #[fastrace::trace]
-    async fn extend_validity(&self, new_ts: Timestamp) -> anyhow::Result<ExtendValidityResult> {
+    async fn extend_validity(&self, new_ts: Timestamp) -> anyhow::Result<SubscriptionValidity> {
         if new_ts < self.initial_ts {
             // new_ts is before the initial subscription timestamp.
-            return Ok(ExtendValidityResult::Invalid { invalid_ts: None });
+            return Ok(SubscriptionValidity::Invalid { invalid_ts: None });
         }
 
         // The inner subscription is periodically updated by the subscription
@@ -625,18 +629,18 @@ impl SubscriptionTrait for ApplicationSubscription {
             // Subscription is no longer valid. We could check validity from end_ts
             // to new_ts, but this is likely to fail and is potentially unbounded amount of
             // work, so we return false here. This is valid per the function contract.
-            return Ok(ExtendValidityResult::Invalid {
+            return Ok(SubscriptionValidity::Invalid {
                 invalid_ts: self.inner.invalid_ts(),
             });
         };
 
         let current_token = Token::new(self.reads.clone(), current_ts);
         Ok(match self.log.refresh_token(current_token, new_ts)? {
-            Ok(_new_token) => ExtendValidityResult::Extended,
+            Ok(_new_token) => SubscriptionValidity::Valid,
             Err(invalid_ts) => {
                 // Subscription validity can't be extended. Note that returning false
                 // here also doesn't mean there is a conflict.
-                ExtendValidityResult::Invalid { invalid_ts }
+                SubscriptionValidity::Invalid { invalid_ts }
             },
         })
     }

crates/database/src/subscription.rs

@@ -606,13 +606,15 @@ impl Subscription {
         self.validity.invalid_ts()
     }
 
-    pub fn wait_for_invalidation(&self) -> impl Future<Output = ()> {
+    pub fn wait_for_invalidation(&self) -> impl Future<Output = Option<Timestamp>> {
         let mut valid = self.valid.clone();
+        let validity = self.validity.clone();
         let span = fastrace::Span::enter_with_local_parent("wait_for_invalidation");
         async move {
             let _: Result<_, _> = valid
                 .wait_for(|state| matches!(state, SubscriptionState::Invalid))
                 .await;
+            validity.invalid_ts()
         }
         .in_span(span)
     }

crates/sync/src/state.rs

@@ -462,7 +462,7 @@ impl SyncState {
                 .as_ref()
                 .ok_or_else(|| anyhow::anyhow!("Missing subscription for {}", query_id))?
                 .wait_for_invalidation()
-                .map(move |r| r.map(move |()| query_id));
+                .map(move |r| r.map(move |_| query_id));
             let (future, handle) = future::abortable(future);
             sq.invalidation_future = Some(handle);
             self.invalidation_futures.push(future.boxed());

crates/sync/src/worker.rs

@@ -15,9 +15,9 @@ use application::{
     api::{
         ApplicationApi,
         ExecuteQueryTimestamp,
-        ExtendValidityResult,
         SubscriptionClient,
         SubscriptionTrait,
+        SubscriptionValidity,
     },
     redaction::{
         RedactedJsError,
@@ -821,8 +821,8 @@ impl<RT: Runtime> SyncWorker<RT> {
                         let new_subscription = match current_subscription {
                             Some(subscription) => {
                                 match subscription.extend_validity(new_ts).await? {
-                                    ExtendValidityResult::Extended => Some(subscription),
-                                    ExtendValidityResult::Invalid { invalid_ts } => {
+                                    SubscriptionValidity::Valid => Some(subscription),
+                                    SubscriptionValidity::Invalid { invalid_ts } => {
                                         metrics::log_query_invalidated(
                                             partition_id,
                                             invalid_ts,