From 2052b8b98dae90cf7eeab76dee1eb34cf3fab750 Mon Sep 17 00:00:00 2001
From: Christian Schwarz <christian@neon.tech>
Date: Fri, 7 Feb 2025 03:11:03 +0100
Subject: [PATCH] make shard available during flush

---
 pageserver/src/page_service.rs | 10 ++++++++--
 1 file changed, 8 insertions(+), 2 deletions(-)

diff --git a/pageserver/src/page_service.rs b/pageserver/src/page_service.rs
index 880d0f5692..1711c19184 100644
--- a/pageserver/src/page_service.rs
+++ b/pageserver/src/page_service.rs
@@ -913,9 +913,10 @@ impl PageServerHandler {
         IO: AsyncRead + AsyncWrite + Send + Sync + Unpin,
     {
         // invoke handler function
-        let (handler_results, span): (
+        let (handler_results, span, shard): (
             Vec<Result<(PagestreamBeMessage, SmgrOpTimer), PageStreamError>>,
             _,
+            _,
         ) = match batch {
             BatchedFeMessage::Exists {
                 span,
@@ -931,6 +932,7 @@ impl PageServerHandler {
                         .await
                         .map(|msg| (msg, timer))],
                     span,
+                    Some(shard),
                 )
             }
             BatchedFeMessage::Nblocks {
@@ -947,6 +949,7 @@ impl PageServerHandler {
                         .await
                         .map(|msg| (msg, timer))],
                     span,
+                    Some(shard),
                 )
             }
             BatchedFeMessage::GetPage {
@@ -973,6 +976,7 @@ impl PageServerHandler {
                         res
                     },
                     span,
+                    Some(shard),
                 )
             }
             BatchedFeMessage::DbSize {
@@ -989,6 +993,7 @@ impl PageServerHandler {
                         .await
                         .map(|msg| (msg, timer))],
                     span,
+                    Some(shard),
                 )
             }
             BatchedFeMessage::GetSlruSegment {
@@ -1005,12 +1010,13 @@ impl PageServerHandler {
                         .await
                         .map(|msg| (msg, timer))],
                     span,
+                    Some(shard),
                 )
             }
             BatchedFeMessage::RespondError { span, error } => {
                 // We've already decided to respond with an error, so we don't need to
                 // call the handler.
-                (vec![Err(error)], span)
+                (vec![Err(error)], span, None)
             }
         };