handle downgrade during batching

pageserver/src/page_service.rs

@@ -558,31 +558,31 @@ enum BatchedFeMessage {
     Exists {
         span: Span,
         timer: SmgrOpTimer,
-        shard: timeline::handle::Handle<TenantManagerTypes>,
+        shard: timeline::handle::WeakHandle<TenantManagerTypes>,
         req: models::PagestreamExistsRequest,
     },
     Nblocks {
         span: Span,
         timer: SmgrOpTimer,
-        shard: timeline::handle::Handle<TenantManagerTypes>,
+        shard: timeline::handle::WeakHandle<TenantManagerTypes>,
         req: models::PagestreamNblocksRequest,
     },
     GetPage {
         span: Span,
-        shard: timeline::handle::Handle<TenantManagerTypes>,
+        shard: timeline::handle::WeakHandle<TenantManagerTypes>,
         effective_request_lsn: Lsn,
         pages: smallvec::SmallVec<[BatchedGetPageRequest; 1]>,
     },
     DbSize {
         span: Span,
         timer: SmgrOpTimer,
-        shard: timeline::handle::Handle<TenantManagerTypes>,
+        shard: timeline::handle::WeakHandle<TenantManagerTypes>,
         req: models::PagestreamDbSizeRequest,
     },
     GetSlruSegment {
         span: Span,
         timer: SmgrOpTimer,
-        shard: timeline::handle::Handle<TenantManagerTypes>,
+        shard: timeline::handle::WeakHandle<TenantManagerTypes>,
         req: models::PagestreamGetSlruSegmentRequest,
     },
     RespondError {
@@ -734,7 +734,7 @@ impl PageServerHandler {
                 BatchedFeMessage::Exists {
                     span,
                     timer,
-                    shard,
+                    shard: shard.downgrade(),
                     req,
                 }
             }
@@ -753,7 +753,7 @@ impl PageServerHandler {
                 BatchedFeMessage::Nblocks {
                     span,
                     timer,
-                    shard,
+                    shard: shard.downgrade(),
                     req,
                 }
             }
@@ -772,7 +772,7 @@ impl PageServerHandler {
                 BatchedFeMessage::DbSize {
                     span,
                     timer,
-                    shard,
+                    shard: shard.downgrade(),
                     req,
                 }
             }
@@ -791,7 +791,7 @@ impl PageServerHandler {
                 BatchedFeMessage::GetSlruSegment {
                     span,
                     timer,
-                    shard,
+                    shard: shard.downgrade(),
                     req,
                 }
             }
@@ -843,6 +843,7 @@ impl PageServerHandler {
                 )
                 .await?;
 
+                // We're holding the Handle
                 let effective_request_lsn = match Self::wait_or_get_last_lsn(
                     &shard,
                     req.hdr.request_lsn,
@@ -860,7 +861,7 @@ impl PageServerHandler {
                 };
                 BatchedFeMessage::GetPage {
                     span,
-                    shard,
+                    shard: shard.downgrade(),
                     effective_request_lsn,
                     pages: smallvec::smallvec![BatchedGetPageRequest { req, timer }],
                 }
@@ -968,7 +969,14 @@ impl PageServerHandler {
                 fail::fail_point!("ps::handle-pagerequest-message::exists");
                 (
                     vec![self
-                        .handle_get_rel_exists_request(&shard, &req, ctx)
+                        .handle_get_rel_exists_request(
+                            &*shard
+                                .upgrade()
+                                // TODO: need sensitive to cancel()
+                                .await?,
+                            &req,
+                            ctx,
+                        )
                         .instrument(span.clone())
                         .await
                         .map(|msg| (msg, timer))
@@ -985,7 +993,7 @@ impl PageServerHandler {
                 fail::fail_point!("ps::handle-pagerequest-message::nblocks");
                 (
                     vec![self
-                        .handle_get_nblocks_request(&shard, &req, ctx)
+                        .handle_get_nblocks_request(&*shard.upgrade().await?, &req, ctx)
                         .instrument(span.clone())
                         .await
                         .map(|msg| (msg, timer))
@@ -1006,7 +1014,7 @@ impl PageServerHandler {
                         trace!(npages, "handling getpage request");
                         let res = self
                             .handle_get_page_at_lsn_request_batched(
-                                &shard,
+                                &*shard.upgrade().await?,
                                 effective_request_lsn,
                                 pages,
                                 ctx,
@@ -1028,7 +1036,7 @@ impl PageServerHandler {
                 fail::fail_point!("ps::handle-pagerequest-message::dbsize");
                 (
                     vec![self
-                        .handle_db_size_request(&shard, &req, ctx)
+                        .handle_db_size_request(&*shard.upgrade().await?, &req, ctx)
                         .instrument(span.clone())
                         .await
                         .map(|msg| (msg, timer))
@@ -1045,7 +1053,7 @@ impl PageServerHandler {
                 fail::fail_point!("ps::handle-pagerequest-message::slrusegment");
                 (
                     vec![self
-                        .handle_get_slru_segment_request(&shard, &req, ctx)
+                        .handle_get_slru_segment_request(&*shard.upgrade().await?, &req, ctx)
                         .instrument(span.clone())
                         .await
                         .map(|msg| (msg, timer))
@@ -2391,6 +2399,14 @@ impl From<GetActiveTimelineError> for QueryError {
     }
 }
 
+impl From<crate::tenant::timeline::handle::HandleUpgradeError> for QueryError {
+    fn from(e: crate::tenant::timeline::handle::HandleUpgradeError) -> Self {
+        match e {
+            crate::tenant::timeline::handle::HandleUpgradeError::ShutDown => QueryError::Shutdown,
+        }
+    }
+}
+
 fn set_tracing_field_shard_id(timeline: &Timeline) {
     debug_assert_current_span_has_tenant_and_timeline_id_no_shard_id();
     tracing::Span::current().record(

pageserver/src/tenant/timeline/handle.rs

@@ -121,6 +121,7 @@ use utils::id::TimelineId;
 use utils::shard::ShardIndex;
 use utils::shard::ShardNumber;
 
+use crate::page_service::GetActiveTimelineError;
 use crate::tenant::mgr::ShardSelector;
 
 /// The requirement for Debug is so that #[derive(Debug)] works in some places.
@@ -153,7 +154,7 @@ pub(crate) struct Cache<T: Types> {
     map: Map<T>,
 }
 
-type Map<T> = HashMap<ShardTimelineId, Weak<HandleInner<T>>>;
+type Map<T> = HashMap<ShardTimelineId, WeakHandle<T>>;
 
 impl<T: Types> Default for Cache<T> {
     fn default() -> Self {
@@ -171,11 +172,8 @@ pub(crate) struct ShardTimelineId {
 }
 
 /// See module-level comment.
-pub(crate) struct Handle<T: Types>(
-    // We need sth with interior mutability + Send-ness.
-    // tokio sync Mutex gives that, and adds OwnedMutexGuard, which is handy.
-    Arc<tokio::sync::Mutex<HandleInner<T>>>,
-);
+pub(crate) struct Handle<T: Types>(tokio::sync::OwnedMutexGuard<HandleInner<T>>);
+pub(crate) struct WeakHandle<T: Types>(Arc<tokio::sync::Mutex<HandleInner<T>>>);
 enum HandleInner<T: Types> {
     KeepingTimelineGateOpen {
         timeline: T::Timeline,
@@ -414,18 +412,37 @@ impl<T: Types> Cache<T> {
     }
 }
 
-impl<T: Types> Handle<T> {
-    pub(crate) async fn upgrade(&mut self) -> Option<impl Deref<Target = T::Timeline>> {
-        // This could probably simplified with a
+pub(crate) enum HandleUpgradeError {
+    ShutDown,
+}
+
+impl<T: Types> WeakHandle<T> {
+    pub(crate) async fn upgrade(&self) -> Result<Handle<T>, HandleUpgradeError> {
         let lock_guard = self.0.lock_owned().await;
-        let res = OwnedMutexGuard::try_map(lock_guard, |f: &mut HandleInner<T>| match f {
-            HandleInner::KeepingTimelineGateOpen { timeline, .. } => Some(&mut timeline),
-            HandleInner::ShutDown => None,
-        });
-        let Ok(mapped_lock_guard) = res else {
-            return None;
-        };
-        Some(mapped_lock_guard)
+        match &*lock_guard {
+            HandleInner::KeepingTimelineGateOpen {
+                timeline,
+                gate_guard,
+            } => Ok(Handle(lock_guard)),
+            HandleInner::ShutDown => Err(HandleUpgradeError::ShutDown),
+        }
+    }
+}
+
+impl<T: Types> Deref for Handle<T> {
+    type Target = T::Timeline;
+
+    fn deref(&self) -> &Self::Target {
+        match &*self.0 {
+            HandleInner::KeepingTimelineGateOpen { timeline, .. } => timeline,
+            HandleInner::ShutDown => unreachable!(),
+        }
+    }
+}
+
+impl<T: Types> Handle<T> {
+    pub(crate) fn downgrade(&self) -> WeakHandle<T> {
+        WeakHandle(Arc::clone(&self.0.mutex()))
     }
 }