From df2806e7a09dc6da310cbdd0db601d41f6d2306f Mon Sep 17 00:00:00 2001
From: Erik Grinaker <erik@neon.tech>
Date: Wed, 30 Apr 2025 15:00:16 +0200
Subject: [PATCH] page_api: add `GetPageRequest::id`

---
 pageserver/page_api/proto/page_service.proto      | 15 +++++++++++----
 pageserver/page_api/src/model.rs                  |  3 +++
 .../pagebench/src/cmd/getpage_latest_lsn.rs       |  2 ++
 pageserver/src/compute_service_grpc.rs            |  7 +++++--
 4 files changed, 21 insertions(+), 6 deletions(-)

diff --git a/pageserver/page_api/proto/page_service.proto b/pageserver/page_api/proto/page_service.proto
index a8514792e2..0afdcb3899 100644
--- a/pageserver/page_api/proto/page_service.proto
+++ b/pageserver/page_api/proto/page_service.proto
@@ -81,13 +81,20 @@ message RelSizeResponse {
 }
 
 message GetPageRequest {
-  RequestCommon common = 1;
-  RelTag rel = 2;
-  uint32 block_number = 3;
+  // A request ID. Will be included in the response. Should be unique for
+  // in-flight requests on the stream.
+  uint64 id = 1;
+  RequestCommon common = 2;
+  RelTag rel = 3;
+  uint32 block_number = 4;
 }
 
+// TODO: should this include page metadata, like reltag, LSN, and block number?
 message GetPageResponse {
-  bytes page_image = 1;
+  // The original request's ID.
+  uint64 id = 1;
+  // The 8KB page image.
+  bytes page_image = 2;
 }
 
 message DbSizeRequest {
diff --git a/pageserver/page_api/src/model.rs b/pageserver/page_api/src/model.rs
index 1d028e5075..6a112f5de2 100644
--- a/pageserver/page_api/src/model.rs
+++ b/pageserver/page_api/src/model.rs
@@ -49,6 +49,7 @@ pub struct RelSizeResponse {
 
 #[derive(Clone, Debug)]
 pub struct GetPageRequest {
+    pub id: u64,
     pub common: RequestCommon,
     pub rel: RelTag,
     pub block_number: u32,
@@ -188,6 +189,7 @@ impl TryFrom<&proto::RelSizeRequest> for RelSizeRequest {
 impl From<&GetPageRequest> for proto::GetPageRequest {
     fn from(value: &GetPageRequest) -> proto::GetPageRequest {
         proto::GetPageRequest {
+            id: value.id,
             common: Some((&value.common).into()),
             rel: Some((&value.rel).into()),
             block_number: value.block_number,
@@ -199,6 +201,7 @@ impl TryFrom<&proto::GetPageRequest> for GetPageRequest {
 
     fn try_from(value: &proto::GetPageRequest) -> Result<GetPageRequest, ProtocolError> {
         Ok(GetPageRequest {
+            id: value.id,
             common: (&value.common.ok_or(ProtocolError::Missing("common"))?).into(),
             rel: (&value.rel.ok_or(ProtocolError::Missing("rel"))?).try_into()?,
             block_number: value.block_number,
diff --git a/pageserver/pagebench/src/cmd/getpage_latest_lsn.rs b/pageserver/pagebench/src/cmd/getpage_latest_lsn.rs
index 422844cae4..e8550eaa9e 100644
--- a/pageserver/pagebench/src/cmd/getpage_latest_lsn.rs
+++ b/pageserver/pagebench/src/cmd/getpage_latest_lsn.rs
@@ -500,6 +500,7 @@ async fn client_grpc(
                     .to_rel_block()
                     .expect("we filter non-rel-block keys out above");
                 pageserver_page_api::model::GetPageRequest {
+                    id: 0, // TODO
                     common: pageserver_page_api::model::RequestCommon {
                         request_lsn: if rng.gen_bool(args.req_latest_probability) {
                             Lsn::MAX
@@ -605,6 +606,7 @@ async fn client_grpc_stream(
                     .to_rel_block()
                     .expect("we filter non-rel-block keys out above");
                 pageserver_page_api::model::GetPageRequest {
+                    id: 0, // TODO
                     common: pageserver_page_api::model::RequestCommon {
                         request_lsn: if rng.gen_bool(args.req_latest_probability) {
                             Lsn::MAX
diff --git a/pageserver/src/compute_service_grpc.rs b/pageserver/src/compute_service_grpc.rs
index 7bc5727af3..4305d289f8 100644
--- a/pageserver/src/compute_service_grpc.rs
+++ b/pageserver/src/compute_service_grpc.rs
@@ -261,7 +261,10 @@ impl PageService for PageServiceService {
                 )
                 .await?;
 
-            Ok(tonic::Response::new(proto::GetPageResponse { page_image }))
+            Ok(tonic::Response::new(proto::GetPageResponse {
+                id: req.id,
+                page_image,
+            }))
         }
         .instrument(span)
         .await
@@ -308,7 +311,7 @@ impl PageService for PageServiceService {
                     )
                     .await?;
 
-                yield proto::GetPageResponse { page_image };
+                yield proto::GetPageResponse { id: request.id, page_image };
             }
         };