diff --git a/pageserver/src/page_service.rs b/pageserver/src/page_service.rs index b3e18fed99..cf2276f65e 100644 --- a/pageserver/src/page_service.rs +++ b/pageserver/src/page_service.rs @@ -558,31 +558,31 @@ enum BatchedFeMessage { Exists { span: Span, timer: SmgrOpTimer, - shard: timeline::handle::Handle, + shard: timeline::handle::WeakHandle, req: models::PagestreamExistsRequest, }, Nblocks { span: Span, timer: SmgrOpTimer, - shard: timeline::handle::Handle, + shard: timeline::handle::WeakHandle, req: models::PagestreamNblocksRequest, }, GetPage { span: Span, - shard: timeline::handle::Handle, + shard: timeline::handle::WeakHandle, effective_request_lsn: Lsn, pages: smallvec::SmallVec<[BatchedGetPageRequest; 1]>, }, DbSize { span: Span, timer: SmgrOpTimer, - shard: timeline::handle::Handle, + shard: timeline::handle::WeakHandle, req: models::PagestreamDbSizeRequest, }, GetSlruSegment { span: Span, timer: SmgrOpTimer, - shard: timeline::handle::Handle, + shard: timeline::handle::WeakHandle, 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 for QueryError { } } +impl From 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( diff --git a/pageserver/src/tenant/timeline/handle.rs b/pageserver/src/tenant/timeline/handle.rs index 68ba5929e5..983eae89ca 100644 --- a/pageserver/src/tenant/timeline/handle.rs +++ b/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 { map: Map, } -type Map = HashMap>>; +type Map = HashMap>; impl Default for Cache { fn default() -> Self { @@ -171,11 +172,8 @@ pub(crate) struct ShardTimelineId { } /// See module-level comment. -pub(crate) struct Handle( - // We need sth with interior mutability + Send-ness. - // tokio sync Mutex gives that, and adds OwnedMutexGuard, which is handy. - Arc>>, -); +pub(crate) struct Handle(tokio::sync::OwnedMutexGuard>); +pub(crate) struct WeakHandle(Arc>>); enum HandleInner { KeepingTimelineGateOpen { timeline: T::Timeline, @@ -414,18 +412,37 @@ impl Cache { } } -impl Handle { - pub(crate) async fn upgrade(&mut self) -> Option> { - // This could probably simplified with a +pub(crate) enum HandleUpgradeError { + ShutDown, +} + +impl WeakHandle { + pub(crate) async fn upgrade(&self) -> Result, HandleUpgradeError> { let lock_guard = self.0.lock_owned().await; - let res = OwnedMutexGuard::try_map(lock_guard, |f: &mut HandleInner| 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 Deref for Handle { + type Target = T::Timeline; + + fn deref(&self) -> &Self::Target { + match &*self.0 { + HandleInner::KeepingTimelineGateOpen { timeline, .. } => timeline, + HandleInner::ShutDown => unreachable!(), + } + } +} + +impl Handle { + pub(crate) fn downgrade(&self) -> WeakHandle { + WeakHandle(Arc::clone(&self.0.mutex())) } }