From 17e2bff4d5e2f31e0f69bce59462cd7f46d9024d Mon Sep 17 00:00:00 2001
From: Vlad Lazar <vlad@neon.tech>
Date: Mon, 10 Feb 2025 23:27:20 +0100
Subject: [PATCH] pageserver: add perf span utilities to request context

---
 pageserver/src/context.rs | 108 ++++++++++++++++++++++++++++++++++++++
 1 file changed, 108 insertions(+)

diff --git a/pageserver/src/context.rs b/pageserver/src/context.rs
index 072f8081c3..81ed86507d 100644
--- a/pageserver/src/context.rs
+++ b/pageserver/src/context.rs
@@ -89,6 +89,13 @@
 //! [`RequestContext`] argument. Functions in the middle of the call chain
 //! only need to pass it on.
 
+use futures::FutureExt;
+use futures::future::BoxFuture;
+use std::future::Future;
+use tracing_utils::perf_span::{PerfInstrument, PerfSpan};
+
+use tracing::{Dispatch, Span};
+
 use crate::task_mgr::TaskKind;
 
 // The main structure of this module, see module-level comment.
@@ -99,6 +106,21 @@ pub struct RequestContext {
     access_stats_behavior: AccessStatsBehavior,
     page_content_kind: PageContentKind,
     read_path_debug: bool,
+    perf_span: Option<PerfSpan>,
+    perf_span_dispatch: Option<Dispatch>,
+}
+
+impl std::fmt::Debug for RequestContext {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("RequestContext")
+            .field("task_kind", &self.task_kind)
+            .field("download_behavior", &self.download_behavior)
+            .field("access_stats_behavior", &self.access_stats_behavior)
+            .field("page_content_kind", &self.page_content_kind)
+            .field("read_path_debug", &self.read_path_debug)
+            // perf_span and perf_span_dispatch are omitted on purpose
+            .finish()
+    }
 }
 
 /// The kind of access to the page cache.
@@ -157,6 +179,8 @@ impl RequestContextBuilder {
                 access_stats_behavior: AccessStatsBehavior::Update,
                 page_content_kind: PageContentKind::Unknown,
                 read_path_debug: false,
+                perf_span: None,
+                perf_span_dispatch: None,
             },
         }
     }
@@ -196,6 +220,43 @@ impl RequestContextBuilder {
         self
     }
 
+    pub(crate) fn perf_span_dispatch(mut self, dispatch: Option<Dispatch>) -> Self {
+        self.inner.perf_span_dispatch = dispatch;
+        self
+    }
+
+    pub fn root_perf_span<Fn>(mut self, make_span: Fn) -> Self
+    where
+        Fn: FnOnce() -> Span,
+    {
+        assert!(self.inner.perf_span.is_none());
+        assert!(self.inner.perf_span_dispatch.is_some());
+
+        let dispatcher = self.inner.perf_span_dispatch.as_ref().unwrap();
+        let new_span = tracing::dispatcher::with_default(dispatcher, make_span);
+
+        self.inner.perf_span = Some(PerfSpan::new(new_span, dispatcher.clone()));
+
+        self
+    }
+
+    pub fn perf_span<Fn>(mut self, make_span: Fn) -> Self
+    where
+        Fn: FnOnce(&Span) -> Span,
+    {
+        if let Some(ref perf_span) = self.inner.perf_span {
+            assert!(self.inner.perf_span_dispatch.is_some());
+            let dispatcher = self.inner.perf_span_dispatch.as_ref().unwrap();
+
+            let new_span =
+                tracing::dispatcher::with_default(dispatcher, || make_span(perf_span.inner()));
+
+            self.inner.perf_span = Some(PerfSpan::new(new_span, dispatcher.clone()));
+        }
+
+        self
+    }
+
     pub fn root(self) -> RequestContext {
         self.inner
     }
@@ -307,4 +368,51 @@ impl RequestContext {
     pub(crate) fn read_path_debug(&self) -> bool {
         self.read_path_debug
     }
+
+    pub(crate) fn perf_follows_from(&self, from: &RequestContext) {
+        if let (Some(span), Some(from_span)) = (&self.perf_span, &from.perf_span) {
+            span.inner().follows_from(from_span.inner());
+        }
+    }
+
+    pub(crate) fn maybe_instrument<'a, Fut, Fn>(
+        &self,
+        future: Fut,
+        make_span: Fn,
+    ) -> BoxFuture<'a, Fut::Output>
+    where
+        Fut: Future + Send + 'a,
+        Fn: FnOnce(&Span) -> Span,
+    {
+        match &self.perf_span {
+            Some(perf_span) => {
+                assert!(self.perf_span_dispatch.is_some());
+                let dispatcher = self.perf_span_dispatch.as_ref().unwrap();
+
+                let new_span =
+                    tracing::dispatcher::with_default(dispatcher, || make_span(perf_span.inner()));
+
+                let new_perf_span = PerfSpan::new(new_span, dispatcher.clone());
+                future.instrument(new_perf_span).boxed()
+            }
+            None => future.boxed(),
+        }
+    }
+
+    pub(crate) fn perf_span_record<
+        Q: tracing::field::AsField + ?Sized,
+        V: tracing::field::Value,
+    >(
+        &self,
+        field: &Q,
+        value: V,
+    ) {
+        if let Some(span) = &self.perf_span {
+            span.record(field, value);
+        }
+    }
+
+    pub(crate) fn has_perf_span(&self) -> bool {
+        self.perf_span.is_some()
+    }
 }