fix close order

## Problem

A binary Protobuf schema descriptor can be used to expose an API
reflection service, which in turn allows convenient usage of e.g.
`grpcurl` against the gRPC server.

Touches #11728.

## Summary of changes

* Generate a binary schema descriptor as
`pageserver_page_api::proto::FILE_DESCRIPTOR_SET`.
* Opportunistically rename the Protobuf package from `page_service` to
`page_api`.

.github/workflows/build_and_test.yml

@@ -963,7 +963,7 @@ jobs:
           fi
 
       - name: Verify docker-compose example and test extensions
-        timeout-minutes: 20
+        timeout-minutes: 60
         env:
           TAG: >-
             ${{

Cargo.lock

@@ -4434,6 +4434,16 @@ dependencies = [
  "workspace_hack",
 ]
 
+[[package]]
+name = "pageserver_page_api"
+version = "0.1.0"
+dependencies = [
+ "prost 0.13.3",
+ "tonic",
+ "tonic-build",
+ "workspace_hack",
+]
+
 [[package]]
 name = "papaya"
 version = "0.2.1"

Cargo.toml

@@ -9,6 +9,7 @@ members = [
     "pageserver/ctl",
     "pageserver/client",
     "pageserver/pagebench",
+    "pageserver/page_api",
     "proxy",
     "safekeeper",
     "safekeeper/client",
@@ -252,6 +253,7 @@ pageserver = { path = "./pageserver" }
 pageserver_api = { version = "0.1", path = "./libs/pageserver_api/" }
 pageserver_client = { path = "./pageserver/client" }
 pageserver_compaction = { version = "0.1", path = "./pageserver/compaction/" }
+pageserver_page_api = { path = "./pageserver/page_api" }
 postgres_backend = { version = "0.1", path = "./libs/postgres_backend/" }
 postgres_connection = { version = "0.1", path = "./libs/postgres_connection/" }
 postgres_ffi = { version = "0.1", path = "./libs/postgres_ffi/" }

compute/patches/rum.patch

@@ -7,7 +7,7 @@ index 255e616..1c6edb7 100644
  			 RelationGetRelationName(index));
  
 +#ifdef NEON_SMGR
-+	smgr_start_unlogged_build(index->rd_smgr);
++	smgr_start_unlogged_build(RelationGetSmgr(index));
 +#endif
 +
  	initRumState(&buildstate.rumstate, index);
@@ -18,7 +18,7 @@ index 255e616..1c6edb7 100644
  	rumUpdateStats(index, &buildstate.buildStats, buildstate.rumstate.isBuild);
  
 +#ifdef NEON_SMGR
-+	smgr_finish_unlogged_build_phase_1(index->rd_smgr);
++	smgr_finish_unlogged_build_phase_1(RelationGetSmgr(index));
 +#endif
 +
  	/*
@@ -29,7 +29,7 @@ index 255e616..1c6edb7 100644
  	}
  
 +#ifdef NEON_SMGR
-+	smgr_end_unlogged_build(index->rd_smgr);
++	smgr_end_unlogged_build(RelationGetSmgr(index));
 +#endif
 +
  	/*

endpoint_storage/src/app.rs

@@ -462,6 +462,8 @@ MC4CAQAwBQYDK2VwBCIEID/Drmc1AA6U/znNRWpF3zEGegOATQxfkdWxitcOMsIH
         if var(REAL_S3_ENV).is_ok() {
             assert!(body.contains("remote_storage_s3_deleted_objects_total"));
         }
+
+        #[cfg(target_os = "linux")]
         assert!(body.contains("process_threads"));
     }
 

pageserver/page_api/Cargo.toml

@@ -0,0 +1,13 @@
+[package]
+name = "pageserver_page_api"
+version = "0.1.0"
+edition.workspace = true
+license.workspace = true
+
+[dependencies]
+prost.workspace = true
+tonic.workspace = true
+workspace_hack.workspace = true
+
+[build-dependencies]
+tonic-build.workspace = true

pageserver/page_api/build.rs

@@ -0,0 +1,13 @@
+use std::env;
+use std::path::PathBuf;
+
+/// Generates Rust code from .proto Protobuf schemas, along with a binary file
+/// descriptor set for Protobuf schema reflection.
+fn main() -> Result<(), Box<dyn std::error::Error>> {
+    let out_dir = PathBuf::from(env::var("OUT_DIR")?);
+    tonic_build::configure()
+        .bytes(["."])
+        .file_descriptor_set_path(out_dir.join("page_api_descriptor.bin"))
+        .compile_protos(&["proto/page_service.proto"], &["proto"])
+        .map_err(|err| err.into())
+}

pageserver/page_api/proto/page_service.proto

@@ -0,0 +1,233 @@
+// Page service, presented by pageservers for computes.
+//
+// This is the compute read path. It primarily serves page versions at given
+// LSNs, but also base backups, SLRU segments, and relation metadata.
+//
+// EXPERIMENTAL: this is still under development and subject to change.
+//
+// Request metadata headers:
+// - authorization: JWT token ("Bearer <token>"), if auth is enabled
+// - neon-tenant-id: tenant ID ("7c4a1f9e3bd6470c8f3e21a65bd2e980")
+// - neon-shard-id: shard ID, as <number><count> in hex ("0b10" = shard 11 of 16, 0-based)
+// - neon-timeline-id: timeline ID ("f08c4e9a2d5f76b1e3a7c2d8910f4b3e")
+//
+// The service can be accessed via e.g. grpcurl:
+//
+//    ```
+//    grpcurl \
+//      -plaintext \
+//      -H "neon-tenant-id: 7c4a1f9e3bd6470c8f3e21a65bd2e980" \
+//      -H "neon-shard-id: 0b10" \
+//      -H "neon-timeline-id: f08c4e9a2d5f76b1e3a7c2d8910f4b3e" \
+//      -H "authorization: Bearer $JWT" \
+//      -d '{"read_lsn": {"request_lsn": 1234567890}, "rel": {"spc_oid": 1663, "db_oid": 1234, "rel_number": 5678, "fork_number": 0}}'
+//      localhost:51051 page_api.PageService/CheckRelExists
+//    ```
+//
+// TODO: consider adding neon-compute-mode ("primary", "static", "replica").
+// However, this will require reconnecting when changing modes.
+//
+// TODO: write implementation guidance on
+// - Health checks
+// - Tracing, OpenTelemetry
+// - Compression
+
+syntax = "proto3";
+package page_api;
+
+service PageService {
+  // Returns whether a relation exists.
+  rpc CheckRelExists(CheckRelExistsRequest) returns (CheckRelExistsResponse);
+
+  // Fetches a base backup.
+  rpc GetBaseBackup (GetBaseBackupRequest) returns (stream GetBaseBackupResponseChunk);
+
+  // Returns the total size of a database, as # of bytes.
+  rpc GetDbSize (GetDbSizeRequest) returns (GetDbSizeResponse);
+
+  // Fetches pages.
+  //
+  // This is implemented as a bidirectional streaming RPC for performance. Unary
+  // requests incur costs for e.g. HTTP/2 stream setup, header parsing,
+  // authentication, and so on -- with streaming, we only pay these costs during
+  // the initial stream setup. This ~doubles throughput in benchmarks. Other
+  // RPCs use regular unary requests, since they are not as frequent and
+  // performance-critical, and this simplifies implementation.
+  //
+  // NB: a status response (e.g. errors) will terminate the stream. The stream
+  // may be shared by e.g. multiple Postgres backends, so we should avoid this.
+  // Most errors are therefore sent as GetPageResponse.status instead.
+  rpc GetPages (stream GetPageRequest) returns (stream GetPageResponse);
+
+  // Returns the size of a relation, as # of blocks.
+  rpc GetRelSize (GetRelSizeRequest) returns (GetRelSizeResponse);
+
+  // Fetches an SLRU segment.
+  rpc GetSlruSegment (GetSlruSegmentRequest) returns (GetSlruSegmentResponse);
+}
+
+// The LSN a request should read at.
+message ReadLsn {
+  // The request's read LSN. Required.
+  uint64 request_lsn = 1;
+  // If given, the caller guarantees that the page has not been modified since
+  // this LSN. Must be smaller than or equal to request_lsn. This allows the
+  // Pageserver to serve an old page without waiting for the request LSN to
+  // arrive. Valid for all request types.
+  //
+  // It is undefined behaviour to make a request such that the page was, in
+  // fact, modified between request_lsn and not_modified_since_lsn. The
+  // Pageserver might detect it and return an error, or it might return the old
+  // page version or the new page version. Setting not_modified_since_lsn equal
+  // to request_lsn is always safe, but can lead to unnecessary waiting.
+  uint64 not_modified_since_lsn = 2;
+}
+
+// A relation identifier.
+message RelTag {
+    uint32 spc_oid = 1;
+    uint32 db_oid = 2;
+    uint32 rel_number = 3;
+    uint32 fork_number = 4;
+}
+
+// Checks whether a relation exists, at the given LSN. Only valid on shard 0,
+// other shards will error.
+message CheckRelExistsRequest {
+  ReadLsn read_lsn = 1;
+  RelTag rel = 2;
+}
+
+message CheckRelExistsResponse {
+  bool exists = 1;
+}
+
+// Requests a base backup at a given LSN.
+message GetBaseBackupRequest {
+  // The LSN to fetch a base backup at.
+  ReadLsn read_lsn = 1;
+  // If true, logical replication slots will not be created.
+  bool replica = 2;
+}
+
+// Base backup response chunk, returned as an ordered stream.
+message GetBaseBackupResponseChunk {
+  // A basebackup data chunk. The size is undefined, but bounded by the 4 MB
+  // gRPC message size limit.
+  bytes chunk = 1;
+}
+
+// Requests the size of a database, as # of bytes. Only valid on shard 0, other
+// shards will error.
+message GetDbSizeRequest {
+  ReadLsn read_lsn = 1;
+  uint32 db_oid = 2;
+}
+
+message GetDbSizeResponse {
+  uint64 num_bytes = 1;
+}
+
+// Requests one or more pages.
+message GetPageRequest {
+  // A request ID. Will be included in the response. Should be unique for
+  // in-flight requests on the stream.
+  uint64 request_id = 1;
+  // The request class.
+  GetPageClass request_class = 2;
+  // The LSN to read at.
+  ReadLsn read_lsn = 3;
+  // The relation to read from.
+  RelTag rel = 4;
+  // Page numbers to read. Must belong to the remote shard.
+  //
+  // Multiple pages will be executed as a single batch by the Pageserver,
+  // amortizing layer access costs and parallelizing them. This may increase the
+  // latency of any individual request, but improves the overall latency and
+  // throughput of the batch as a whole.
+  //
+  // TODO: this causes an allocation in the common single-block case. The sender
+  // can use a SmallVec to stack-allocate it, but Prost will always deserialize
+  // into a heap-allocated Vec. Consider optimizing this.
+  //
+  // TODO: we might be able to avoid a sort or something if we mandate that these
+  // are always in order. But we can't currenly rely on this on the server, because
+  // of compatibility with the libpq protocol handler.
+  repeated uint32 block_number = 5;
+}
+
+// A GetPageRequest class. Primarily intended for observability, but may also be
+// used for prioritization in the future.
+enum GetPageClass {
+  // Unknown class. For forwards compatibility: used when the client sends a
+  // class that the server doesn't know about.
+  GET_PAGE_CLASS_UNKNOWN = 0;
+  // A normal request. This is the default.
+  GET_PAGE_CLASS_NORMAL = 1;
+  // A prefetch request. NB: can only be classified on pg < 18.
+  GET_PAGE_CLASS_PREFETCH = 2;
+  // A background request (e.g. vacuum).
+  GET_PAGE_CLASS_BACKGROUND = 3;
+}
+
+// A GetPage response.
+//
+// A batch response will contain all of the requested pages. We could eagerly
+// emit individual pages as soon as they are ready, but on a readv() Postgres
+// holds buffer pool locks on all pages in the batch and we'll only return once
+// the entire batch is ready, so no one can make use of the individual pages.
+message GetPageResponse {
+  // The original request's ID.
+  uint64 request_id = 1;
+  // The response status code.
+  GetPageStatus status = 2;
+  // A string describing the status, if any.
+  string reason = 3;
+  // The 8KB page images, in the same order as the request. Empty if status != OK.
+  repeated bytes page_image = 4;
+}
+
+// A GetPageResponse status code. Since we use a bidirectional stream, we don't
+// want to send errors as gRPC statuses, since this would terminate the stream.
+enum GetPageStatus {
+  // Unknown status. For forwards compatibility: used when the server sends a
+  // status code that the client doesn't know about.
+  GET_PAGE_STATUS_UNKNOWN = 0;
+  // The request was successful.
+  GET_PAGE_STATUS_OK = 1;
+  // The page did not exist. The tenant/timeline/shard has already been
+  // validated during stream setup.
+  GET_PAGE_STATUS_NOT_FOUND = 2;
+  // The request was invalid.
+  GET_PAGE_STATUS_INVALID = 3;
+  // The tenant is rate limited. Slow down and retry later.
+  GET_PAGE_STATUS_SLOW_DOWN = 4;
+  // TODO: consider adding a GET_PAGE_STATUS_LAYER_DOWNLOAD in the case of a
+  // layer download. This could free up the server task to process other
+  // requests while the layer download is in progress.
+}
+
+// Fetches the size of a relation at a given LSN, as # of blocks. Only valid on
+// shard 0, other shards will error.
+message GetRelSizeRequest {
+  ReadLsn read_lsn = 1;
+  RelTag rel = 2;
+}
+
+message GetRelSizeResponse {
+  uint32 num_blocks = 1;
+}
+
+// Requests an SLRU segment. Only valid on shard 0, other shards will error.
+message GetSlruSegmentRequest {
+  ReadLsn read_lsn = 1;
+  uint32 kind = 2;
+  uint32 segno = 3;
+}
+
+// Returns an SLRU segment.
+//
+// These are up 32 pages (256 KB), so we can send them as a single response.
+message GetSlruSegmentResponse {
+  bytes segment = 1;
+}

pageserver/page_api/src/lib.rs

@@ -0,0 +1,19 @@
+//! This crate provides the Pageserver's page API. It contains:
+//!
+//! * proto/page_service.proto: the Protobuf schema for the page API.
+//! * proto: auto-generated Protobuf types for gRPC.
+//!
+//! This crate is used by both the client and the server. Try to keep it slim.
+
+// Code generated by protobuf.
+pub mod proto {
+    tonic::include_proto!("page_api");
+
+    /// File descriptor set for Protobuf schema reflection. This allows using
+    /// e.g. grpcurl with the API.
+    pub const FILE_DESCRIPTOR_SET: &[u8] =
+        tonic::include_file_descriptor_set!("page_api_descriptor");
+
+    pub use page_service_client::PageServiceClient;
+    pub use page_service_server::{PageService, PageServiceServer};
+}

pageserver/src/page_service.rs

@@ -1278,7 +1278,7 @@ impl PageServerHandler {
     }
 
     #[instrument(level = tracing::Level::DEBUG, skip_all)]
-    async fn pagesteam_handle_batched_message<IO>(
+    async fn pagestream_handle_batched_message<IO>(
         &mut self,
         pgb_writer: &mut PostgresBackend<IO>,
         batch: BatchedFeMessage,
@@ -1733,7 +1733,7 @@ impl PageServerHandler {
             };
 
             let result = self
-                .pagesteam_handle_batched_message(
+                .pagestream_handle_batched_message(
                     pgb_writer,
                     msg,
                     io_concurrency.clone(),
@@ -1909,7 +1909,7 @@ impl PageServerHandler {
                             return Err(e);
                         }
                     };
-                    self.pagesteam_handle_batched_message(
+                    self.pagestream_handle_batched_message(
                         pgb_writer,
                         batch,
                         io_concurrency.clone(),

pgxn/neon/file_cache.c

@@ -936,6 +936,44 @@ lfc_prewarm_main(Datum main_arg)
 	lfc_ctl->prewarm_workers[worker_id].completed = GetCurrentTimestamp();
 }
 
+void
+lfc_invalidate(NRelFileInfo rinfo, ForkNumber forkNum, BlockNumber nblocks)
+{
+	BufferTag	tag;
+	FileCacheEntry *entry;
+	uint32		hash;
+
+	if (lfc_maybe_disabled())	/* fast exit if file cache is disabled */
+		return;
+
+	CopyNRelFileInfoToBufTag(tag, rinfo);
+	tag.forkNum = forkNum;
+
+	CriticalAssert(BufTagGetRelNumber(&tag) != InvalidRelFileNumber);
+
+	LWLockAcquire(lfc_lock, LW_EXCLUSIVE);
+	if (LFC_ENABLED())
+	{
+		for (BlockNumber blkno = 0; blkno < nblocks; blkno += lfc_blocks_per_chunk)
+		{
+			tag.blockNum = blkno;
+			hash = get_hash_value(lfc_hash, &tag);
+			entry = hash_search_with_hash_value(lfc_hash, &tag, hash, HASH_FIND, NULL);
+			if (entry != NULL)
+			{
+				for (int i = 0; i < lfc_blocks_per_chunk; i++)
+				{
+					if (GET_STATE(entry, i) == AVAILABLE)
+					{
+						lfc_ctl->used_pages -= 1;
+						SET_STATE(entry, i, UNAVAILABLE);
+					}
+				}
+			}
+		}
+	}
+	LWLockRelease(lfc_lock);
+}
 
 /*
  * Check if page is present in the cache.

pgxn/neon/file_cache.h

@@ -28,6 +28,7 @@ typedef struct FileCacheState
 extern bool lfc_store_prefetch_result;
 
 /* functions for local file cache */
+extern void lfc_invalidate(NRelFileInfo rinfo, ForkNumber forkNum, BlockNumber nblocks);
 extern void lfc_writev(NRelFileInfo rinfo, ForkNumber forkNum,
 					   BlockNumber blkno, const void *const *buffers,
 					   BlockNumber nblocks);

pgxn/neon/neon_pgversioncompat.h

@@ -86,7 +86,7 @@ InitBufferTag(BufferTag *tag, const RelFileNode *rnode,
 
 #define InvalidRelFileNumber InvalidOid
 
-#define SMgrRelGetRelInfo(reln) \
+#define SMgrRelGetRelInfo(reln)				\
 	(reln->smgr_rnode.node)
 
 #define DropRelationAllLocalBuffers DropRelFileNodeAllLocalBuffers
@@ -148,6 +148,12 @@ InitBufferTag(BufferTag *tag, const RelFileNode *rnode,
 #define DropRelationAllLocalBuffers DropRelationAllLocalBuffers
 #endif
 
+#define NRelFileInfoInvalidate(rinfo) do { \
+		NInfoGetSpcOid(rinfo) = InvalidOid; \
+		NInfoGetDbOid(rinfo) = InvalidOid; \
+		NInfoGetRelNumber(rinfo) = InvalidRelFileNumber; \
+	} while (0)
+
 #if PG_MAJORVERSION_NUM < 17
 #define ProcNumber BackendId
 #define INVALID_PROC_NUMBER InvalidBackendId

pgxn/neon/pagestore_smgr.c

@@ -108,7 +108,7 @@ typedef enum
 	UNLOGGED_BUILD_NOT_PERMANENT
 } UnloggedBuildPhase;
 
-static SMgrRelation unlogged_build_rel = NULL;
+static NRelFileInfo unlogged_build_rel_info;
 static UnloggedBuildPhase unlogged_build_phase = UNLOGGED_BUILD_NOT_IN_PROGRESS;
 
 static bool neon_redo_read_buffer_filter(XLogReaderState *record, uint8 block_id);
@@ -912,16 +912,19 @@ neon_extend(SMgrRelation reln, ForkNumber forkNum, BlockNumber blkno,
 	{
 		case 0:
 			neon_log(ERROR, "cannot call smgrextend() on rel with unknown persistence");
+			break;
 
 		case RELPERSISTENCE_PERMANENT:
+			if (RelFileInfoEquals(unlogged_build_rel_info, InfoFromSMgrRel(reln)))
+			{
+				mdextend(reln, forkNum, blkno, buffer, skipFsync);
+				return;
+			}
 			break;
 
 		case RELPERSISTENCE_TEMP:
 		case RELPERSISTENCE_UNLOGGED:
 			mdextend(reln, forkNum, blkno, buffer, skipFsync);
-			/* Update LFC in case of unlogged index build */
-			if (reln == unlogged_build_rel && unlogged_build_phase == UNLOGGED_BUILD_PHASE_2)
-				lfc_write(InfoFromSMgrRel(reln), forkNum, blkno, buffer);
 			return;
 
 		default:
@@ -1003,21 +1006,19 @@ neon_zeroextend(SMgrRelation reln, ForkNumber forkNum, BlockNumber blocknum,
 	{
 		case 0:
 			neon_log(ERROR, "cannot call smgrextend() on rel with unknown persistence");
+			break;
 
 		case RELPERSISTENCE_PERMANENT:
+			if (RelFileInfoEquals(unlogged_build_rel_info, InfoFromSMgrRel(reln)))
+			{
+				mdzeroextend(reln, forkNum, blocknum, nblocks, skipFsync);
+				return;
+			}
 			break;
 
 		case RELPERSISTENCE_TEMP:
 		case RELPERSISTENCE_UNLOGGED:
 			mdzeroextend(reln, forkNum, blocknum, nblocks, skipFsync);
-			/* Update LFC in case of unlogged index build */
-			if (reln == unlogged_build_rel && unlogged_build_phase == UNLOGGED_BUILD_PHASE_2)
-			{
-				for (int i = 0; i < nblocks; i++)
-				{
-					lfc_write(InfoFromSMgrRel(reln), forkNum, blocknum + i, buffer.data);
-				}
-			}
 			return;
 
 		default:
@@ -1387,8 +1388,14 @@ neon_read(SMgrRelation reln, ForkNumber forkNum, BlockNumber blkno, void *buffer
 	{
 		case 0:
 			neon_log(ERROR, "cannot call smgrread() on rel with unknown persistence");
+			break;
 
 		case RELPERSISTENCE_PERMANENT:
+			if (RelFileInfoEquals(unlogged_build_rel_info, InfoFromSMgrRel(reln)))
+			{
+				mdread(reln, forkNum, blkno, buffer);
+				return;
+			}
 			break;
 
 		case RELPERSISTENCE_TEMP:
@@ -1474,8 +1481,14 @@ neon_readv(SMgrRelation reln, ForkNumber forknum, BlockNumber blocknum,
 	{
 		case 0:
 			neon_log(ERROR, "cannot call smgrread() on rel with unknown persistence");
+			break;
 
 		case RELPERSISTENCE_PERMANENT:
+			if (RelFileInfoEquals(unlogged_build_rel_info, InfoFromSMgrRel(reln)))
+			{
+				mdreadv(reln, forknum, blocknum, buffers, nblocks);
+				return;
+			}
 			break;
 
 		case RELPERSISTENCE_TEMP:
@@ -1608,6 +1621,15 @@ neon_write(SMgrRelation reln, ForkNumber forknum, BlockNumber blocknum, const vo
 			break;
 
 		case RELPERSISTENCE_PERMANENT:
+			if (RelFileInfoEquals(unlogged_build_rel_info, InfoFromSMgrRel(reln)))
+			{
+#if PG_MAJORVERSION_NUM >= 17
+				mdwritev(reln, forknum, blocknum, &buffer, 1, skipFsync);
+#else
+				mdwrite(reln, forknum, blocknum, buffer, skipFsync);
+#endif
+				return;
+			}
 			break;
 
 		case RELPERSISTENCE_TEMP:
@@ -1617,9 +1639,6 @@ neon_write(SMgrRelation reln, ForkNumber forknum, BlockNumber blocknum, const vo
 			#else
 			mdwrite(reln, forknum, blocknum, buffer, skipFsync);
 			#endif
-			/* Update LFC in case of unlogged index build */
-			if (reln == unlogged_build_rel && unlogged_build_phase == UNLOGGED_BUILD_PHASE_2)
-				lfc_write(InfoFromSMgrRel(reln), forknum, blocknum, buffer);
 			return;
 		default:
 			neon_log(ERROR, "unknown relpersistence '%c'", reln->smgr_relpersistence);
@@ -1680,14 +1699,16 @@ neon_writev(SMgrRelation reln, ForkNumber forknum, BlockNumber blkno,
 			break;
 
 		case RELPERSISTENCE_PERMANENT:
+			if (RelFileInfoEquals(unlogged_build_rel_info, InfoFromSMgrRel(reln)))
+			{
+				mdwritev(reln, forknum, blkno, buffers, nblocks, skipFsync);
+				return;
+			}
 			break;
 
 		case RELPERSISTENCE_TEMP:
 		case RELPERSISTENCE_UNLOGGED:
 			mdwritev(reln, forknum, blkno, buffers, nblocks, skipFsync);
-			/* Update LFC in case of unlogged index build */
-			if (reln == unlogged_build_rel && unlogged_build_phase == UNLOGGED_BUILD_PHASE_2)
-				lfc_writev(InfoFromSMgrRel(reln), forknum, blkno, buffers, nblocks);
 			return;
 		default:
 			neon_log(ERROR, "unknown relpersistence '%c'", reln->smgr_relpersistence);
@@ -1723,6 +1744,10 @@ neon_nblocks(SMgrRelation reln, ForkNumber forknum)
 			break;
 
 		case RELPERSISTENCE_PERMANENT:
+			if (RelFileInfoEquals(unlogged_build_rel_info, InfoFromSMgrRel(reln)))
+			{
+				return mdnblocks(reln, forknum);
+			}
 			break;
 
 		case RELPERSISTENCE_TEMP:
@@ -1792,6 +1817,11 @@ neon_truncate(SMgrRelation reln, ForkNumber forknum, BlockNumber old_blocks, Blo
 			break;
 
 		case RELPERSISTENCE_PERMANENT:
+			if (RelFileInfoEquals(unlogged_build_rel_info, InfoFromSMgrRel(reln)))
+			{
+				mdtruncate(reln, forknum, old_blocks, nblocks);
+				return;
+			}
 			break;
 
 		case RELPERSISTENCE_TEMP:
@@ -1930,7 +1960,6 @@ neon_start_unlogged_build(SMgrRelation reln)
 	 */
 	if (unlogged_build_phase != UNLOGGED_BUILD_NOT_IN_PROGRESS)
 		neon_log(ERROR, "unlogged relation build is already in progress");
-	Assert(unlogged_build_rel == NULL);
 
 	ereport(SmgrTrace,
 			(errmsg(NEON_TAG "starting unlogged build of relation %u/%u/%u",
@@ -1947,7 +1976,7 @@ neon_start_unlogged_build(SMgrRelation reln)
 
 		case RELPERSISTENCE_TEMP:
 		case RELPERSISTENCE_UNLOGGED:
-			unlogged_build_rel = reln;
+			unlogged_build_rel_info = InfoFromSMgrRel(reln);
 			unlogged_build_phase = UNLOGGED_BUILD_NOT_PERMANENT;
 #ifdef DEBUG_COMPARE_LOCAL
 			if (!IsParallelWorker())
@@ -1968,12 +1997,9 @@ neon_start_unlogged_build(SMgrRelation reln)
 		neon_log(ERROR, "cannot perform unlogged index build, index is not empty ");
 #endif
 
-	unlogged_build_rel = reln;
+	unlogged_build_rel_info = InfoFromSMgrRel(reln);
 	unlogged_build_phase = UNLOGGED_BUILD_PHASE_1;
 
-	/* Make the relation look like it's unlogged */
-	reln->smgr_relpersistence = RELPERSISTENCE_UNLOGGED;
-
 	/*
 	 * Create the local file. In a parallel build, the leader is expected to
 	 * call this first and do it.
@@ -2000,17 +2026,16 @@ neon_start_unlogged_build(SMgrRelation reln)
 static void
 neon_finish_unlogged_build_phase_1(SMgrRelation reln)
 {
-	Assert(unlogged_build_rel == reln);
+	Assert(RelFileInfoEquals(unlogged_build_rel_info, InfoFromSMgrRel(reln)));
 
 	ereport(SmgrTrace,
 			(errmsg(NEON_TAG "finishing phase 1 of unlogged build of relation %u/%u/%u",
-					RelFileInfoFmt(InfoFromSMgrRel(reln)))));
+					RelFileInfoFmt((unlogged_build_rel_info)))));
 
 	if (unlogged_build_phase == UNLOGGED_BUILD_NOT_PERMANENT)
 		return;
 
 	Assert(unlogged_build_phase == UNLOGGED_BUILD_PHASE_1);
-	Assert(reln->smgr_relpersistence == RELPERSISTENCE_UNLOGGED);
 
 	/*
 	 * In a parallel build, (only) the leader process performs the 2nd
@@ -2018,7 +2043,7 @@ neon_finish_unlogged_build_phase_1(SMgrRelation reln)
 	 */
 	if (IsParallelWorker())
 	{
-		unlogged_build_rel = NULL;
+		NRelFileInfoInvalidate(unlogged_build_rel_info);
 		unlogged_build_phase = UNLOGGED_BUILD_NOT_IN_PROGRESS;
 	}
 	else
@@ -2039,11 +2064,11 @@ neon_end_unlogged_build(SMgrRelation reln)
 {
 	NRelFileInfoBackend rinfob = InfoBFromSMgrRel(reln);
 
-	Assert(unlogged_build_rel == reln);
+	Assert(RelFileInfoEquals(unlogged_build_rel_info, InfoFromSMgrRel(reln)));
 
 	ereport(SmgrTrace,
 			(errmsg(NEON_TAG "ending unlogged build of relation %u/%u/%u",
-					RelFileInfoFmt(InfoFromNInfoB(rinfob)))));
+					RelFileInfoFmt(unlogged_build_rel_info))));
 
 	if (unlogged_build_phase != UNLOGGED_BUILD_NOT_PERMANENT)
 	{
@@ -2051,7 +2076,6 @@ neon_end_unlogged_build(SMgrRelation reln)
 		BlockNumber nblocks;
 
 		Assert(unlogged_build_phase == UNLOGGED_BUILD_PHASE_2);
-		Assert(reln->smgr_relpersistence == RELPERSISTENCE_UNLOGGED);
 
 		/*
 		 * Update the last-written LSN cache.
@@ -2072,9 +2096,6 @@ neon_end_unlogged_build(SMgrRelation reln)
 								InfoFromNInfoB(rinfob),
 								MAIN_FORKNUM);
 
-		/* Make the relation look permanent again */
-		reln->smgr_relpersistence = RELPERSISTENCE_PERMANENT;
-
 		/* Remove local copy */
 		for (int forknum = 0; forknum <= MAX_FORKNUM; forknum++)
 		{
@@ -2083,6 +2104,8 @@ neon_end_unlogged_build(SMgrRelation reln)
 				 forknum);
 
 			forget_cached_relsize(InfoFromNInfoB(rinfob), forknum);
+			lfc_invalidate(InfoFromNInfoB(rinfob), forknum, nblocks);
+
 			mdclose(reln, forknum);
 #ifndef DEBUG_COMPARE_LOCAL
 			/* use isRedo == true, so that we drop it immediately */
@@ -2093,7 +2116,7 @@ neon_end_unlogged_build(SMgrRelation reln)
 		mdunlink(rinfob, INIT_FORKNUM, true);
 #endif
 	}
-	unlogged_build_rel = NULL;
+	NRelFileInfoInvalidate(unlogged_build_rel_info);
 	unlogged_build_phase = UNLOGGED_BUILD_NOT_IN_PROGRESS;
 }
 
@@ -2166,7 +2189,7 @@ AtEOXact_neon(XactEvent event, void *arg)
 			 * Forget about any build we might have had in progress. The local
 			 * file will be unlinked by smgrDoPendingDeletes()
 			 */
-			unlogged_build_rel = NULL;
+			NRelFileInfoInvalidate(unlogged_build_rel_info);
 			unlogged_build_phase = UNLOGGED_BUILD_NOT_IN_PROGRESS;
 			break;
 
@@ -2178,7 +2201,7 @@ AtEOXact_neon(XactEvent event, void *arg)
 		case XACT_EVENT_PRE_PREPARE:
 			if (unlogged_build_phase != UNLOGGED_BUILD_NOT_IN_PROGRESS)
 			{
-				unlogged_build_rel = NULL;
+				NRelFileInfoInvalidate(unlogged_build_rel_info);
 				unlogged_build_phase = UNLOGGED_BUILD_NOT_IN_PROGRESS;
 				ereport(ERROR,
 						(errcode(ERRCODE_INTERNAL_ERROR),

poetry.lock

@@ -1,4 +1,4 @@
-# This file is automatically @generated by Poetry 2.1.2 and should not be changed by hand.
+# This file is automatically @generated by Poetry 2.1.1 and should not be changed by hand.
 
 [[package]]
 name = "aiohappyeyeballs"
@@ -1145,18 +1145,19 @@ dotenv = ["python-dotenv"]
 
 [[package]]
 name = "flask-cors"
-version = "5.0.0"
-description = "A Flask extension adding a decorator for CORS support"
+version = "6.0.0"
+description = "A Flask extension simplifying CORS support"
 optional = false
-python-versions = "*"
+python-versions = "<4.0,>=3.9"
 groups = ["main"]
 files = [
-    {file = "Flask_Cors-5.0.0-py2.py3-none-any.whl", hash = "sha256:b9e307d082a9261c100d8fb0ba909eec6a228ed1b60a8315fd85f783d61910bc"},
-    {file = "flask_cors-5.0.0.tar.gz", hash = "sha256:5aadb4b950c4e93745034594d9f3ea6591f734bb3662e16e255ffbf5e89c88ef"},
+    {file = "flask_cors-6.0.0-py3-none-any.whl", hash = "sha256:6332073356452343a8ccddbfec7befdc3fdd040141fe776ec9b94c262f058657"},
+    {file = "flask_cors-6.0.0.tar.gz", hash = "sha256:4592c1570246bf7beee96b74bc0adbbfcb1b0318f6ba05c412e8909eceec3393"},
 ]
 
 [package.dependencies]
-Flask = ">=0.9"
+flask = ">=0.9"
+Werkzeug = ">=0.7"
 
 [[package]]
 name = "frozenlist"

proxy/src/binary/pg_sni_router.rs

@@ -383,12 +383,19 @@ async fn handle_client(
     info!("performing the proxy pass...");
 
     let res = match client {
-        Connection::Raw(mut c) => copy_bidirectional_client_compute(&mut tls_stream, &mut c).await,
-        Connection::Tls(mut c) => copy_bidirectional_client_compute(&mut tls_stream, &mut c).await,
+        Connection::Raw(mut c) => {
+            copy_bidirectional_client_compute(&mut tls_stream, &mut c, |_, _| {}).await
+        }
+        Connection::Tls(mut c) => {
+            copy_bidirectional_client_compute(&mut tls_stream, &mut c, |_, _| {}).await
+        }
     };
 
     match res {
-        Ok(_) => Ok(()),
+        Ok(()) => Ok(()),
+        Err(ErrorSource::Timeout(_)) => Err(anyhow!(
+            "timed out while gracefully shutting down the connection"
+        )),
         Err(ErrorSource::Client(err)) => Err(err).context("client"),
         Err(ErrorSource::Compute(err)) => Err(err).context("compute"),
     }

proxy/src/binary/proxy.rs

@@ -161,8 +161,11 @@ struct ProxyCliArgs {
     #[clap(long, default_values_t = RateBucketInfo::DEFAULT_REDIS_SET)]
     redis_rps_limit: Vec<RateBucketInfo>,
     /// Cancellation channel size (max queue size for redis kv client)
-    #[clap(long, default_value = "1024")]
+    #[clap(long, default_value_t = 1024)]
     cancellation_ch_size: usize,
+    /// Cancellation ops batch size for redis
+    #[clap(long, default_value_t = 8)]
+    cancellation_batch_size: usize,
     /// cache for `allowed_ips` (use `size=0` to disable)
     #[clap(long, default_value = config::CacheOptions::CACHE_DEFAULT_OPTIONS)]
     allowed_ips_cache: String,
@@ -542,7 +545,12 @@ pub async fn run() -> anyhow::Result<()> {
             if let Some(mut redis_kv_client) = redis_kv_client {
                 maintenance_tasks.spawn(async move {
                     redis_kv_client.try_connect().await?;
-                    handle_cancel_messages(&mut redis_kv_client, rx_cancel).await?;
+                    handle_cancel_messages(
+                        &mut redis_kv_client,
+                        rx_cancel,
+                        args.cancellation_batch_size,
+                    )
+                    .await?;
 
                     drop(redis_kv_client);
 

proxy/src/cancellation.rs

@@ -30,8 +30,6 @@ use crate::tls::postgres_rustls::MakeRustlsConnect;
 type IpSubnetKey = IpNet;
 
 const CANCEL_KEY_TTL: i64 = 1_209_600; // 2 weeks cancellation key expire time
-const REDIS_SEND_TIMEOUT: std::time::Duration = std::time::Duration::from_millis(10);
-const BATCH_SIZE: usize = 8;
 
 // Message types for sending through mpsc channel
 pub enum CancelKeyOp {
@@ -231,12 +229,13 @@ impl CancelReplyOp {
 pub async fn handle_cancel_messages(
     client: &mut RedisKVClient,
     mut rx: mpsc::Receiver<CancelKeyOp>,
+    batch_size: usize,
 ) -> anyhow::Result<()> {
-    let mut batch = Vec::with_capacity(BATCH_SIZE);
-    let mut pipeline = Pipeline::with_capacity(BATCH_SIZE);
+    let mut batch = Vec::with_capacity(batch_size);
+    let mut pipeline = Pipeline::with_capacity(batch_size);
 
     loop {
-        if rx.recv_many(&mut batch, BATCH_SIZE).await == 0 {
+        if rx.recv_many(&mut batch, batch_size).await == 0 {
             warn!("shutting down cancellation queue");
             break Ok(());
         }
@@ -367,8 +366,7 @@ impl CancellationHandler {
             return Err(CancelError::InternalError);
         };
 
-        tx.send_timeout(op, REDIS_SEND_TIMEOUT)
-            .await
+        tx.try_send(op)
             .map_err(|e| {
                 tracing::warn!("failed to send GetCancelData for {key}: {e}");
             })
@@ -570,7 +568,7 @@ impl Session {
     }
 
     // Send the store key op to the cancellation handler and set TTL for the key
-    pub(crate) async fn write_cancel_key(
+    pub(crate) fn write_cancel_key(
         &self,
         cancel_closure: CancelClosure,
     ) -> Result<(), CancelError> {
@@ -596,14 +594,14 @@ impl Session {
             expire: CANCEL_KEY_TTL,
         };
 
-        let _ = tx.send_timeout(op, REDIS_SEND_TIMEOUT).await.map_err(|e| {
+        let _ = tx.try_send(op).map_err(|e| {
             let key = self.key;
             tracing::warn!("failed to send StoreCancelKey for {key}: {e}");
         });
         Ok(())
     }
 
-    pub(crate) async fn remove_cancel_key(&self) -> Result<(), CancelError> {
+    pub(crate) fn remove_cancel_key(&self) -> Result<(), CancelError> {
         let Some(tx) = &self.cancellation_handler.tx else {
             tracing::warn!("cancellation handler is not available");
             return Err(CancelError::InternalError);
@@ -619,7 +617,7 @@ impl Session {
                 .guard(RedisMsgKind::HDel),
         };
 
-        let _ = tx.send_timeout(op, REDIS_SEND_TIMEOUT).await.map_err(|e| {
+        let _ = tx.try_send(op).map_err(|e| {
             let key = self.key;
             tracing::warn!("failed to send RemoveCancelKey for {key}: {e}");
         });

proxy/src/console_redirect_proxy.rs

@@ -129,6 +129,12 @@ pub async fn task_main(
                     let _disconnect = ctx.log_connect();
                     match p.proxy_pass(&config.connect_to_compute).await {
                         Ok(()) => {}
+                        Err(ErrorSource::Timeout(_)) => {
+                            info!(
+                                ?session_id,
+                                "per-client task timed out while gracefully shutting down the connection"
+                            );
+                        }
                         Err(ErrorSource::Client(e)) => {
                             error!(
                                 ?session_id,
@@ -244,9 +250,7 @@ pub(crate) async fn handle_client<S: AsyncRead + AsyncWrite + Unpin>(
     let cancellation_handler_clone = Arc::clone(&cancellation_handler);
     let session = cancellation_handler_clone.get_key();
 
-    session
-        .write_cancel_key(node.cancel_closure.clone())
-        .await?;
+    session.write_cancel_key(node.cancel_closure.clone())?;
 
     prepare_client_connection(&node, *session.key(), &mut stream).await?;
 

proxy/src/metrics.rs

@@ -200,8 +200,10 @@ pub enum HttpDirection {
 #[derive(FixedCardinalityLabel, Copy, Clone)]
 #[label(singleton = "direction")]
 pub enum Direction {
-    Tx,
-    Rx,
+    #[label(rename = "tx")]
+    ComputeToClient,
+    #[label(rename = "rx")]
+    ClientToCompute,
 }
 
 #[derive(FixedCardinalityLabel, Clone, Copy, Debug)]

proxy/src/proxy/copy_bidirectional.rs

@@ -1,313 +1,394 @@
 use std::future::poll_fn;
 use std::io;
+use std::ops::Range;
 use std::pin::Pin;
 use std::task::{Context, Poll, ready};
+use std::time::Duration;
 
 use tokio::io::{AsyncRead, AsyncWrite, ReadBuf};
 use tracing::info;
 
-#[derive(Debug)]
-enum TransferState {
-    Running(CopyBuffer),
-    ShuttingDown(u64),
-    Done(u64),
-}
+use crate::metrics::Direction;
 
-#[derive(Debug)]
-pub(crate) enum ErrorDirection {
-    Read(io::Error),
-    Write(io::Error),
-}
+const DISCONNECT_TIMEOUT: Duration = Duration::from_secs(10);
 
-impl ErrorSource {
-    fn from_client(err: ErrorDirection) -> ErrorSource {
-        match err {
-            ErrorDirection::Read(client) => Self::Client(client),
-            ErrorDirection::Write(compute) => Self::Compute(compute),
-        }
-    }
-    fn from_compute(err: ErrorDirection) -> ErrorSource {
-        match err {
-            ErrorDirection::Write(client) => Self::Client(client),
-            ErrorDirection::Read(compute) => Self::Compute(compute),
-        }
-    }
+/// Mark a value as being unlikely.
+#[cold]
+#[inline(always)]
+fn cold<I>(i: I) -> I {
+    i
 }
 
 #[derive(Debug)]
 pub enum ErrorSource {
     Client(io::Error),
     Compute(io::Error),
+    Timeout(tokio::time::error::Elapsed),
 }
 
-fn transfer_one_direction<A, B>(
-    cx: &mut Context<'_>,
-    state: &mut TransferState,
-    r: &mut A,
-    w: &mut B,
-) -> Poll<Result<u64, ErrorDirection>>
-where
-    A: AsyncRead + AsyncWrite + Unpin + ?Sized,
-    B: AsyncRead + AsyncWrite + Unpin + ?Sized,
-{
-    let mut r = Pin::new(r);
-    let mut w = Pin::new(w);
-    loop {
-        match state {
-            TransferState::Running(buf) => {
-                let count = ready!(buf.poll_copy(cx, r.as_mut(), w.as_mut()))?;
-                *state = TransferState::ShuttingDown(count);
-            }
-            TransferState::ShuttingDown(count) => {
-                ready!(w.as_mut().poll_shutdown(cx)).map_err(ErrorDirection::Write)?;
-                *state = TransferState::Done(*count);
-            }
-            TransferState::Done(count) => return Poll::Ready(Ok(*count)),
+impl ErrorSource {
+    fn read(dir: Direction, err: io::Error) -> Self {
+        match dir {
+            Direction::ComputeToClient => ErrorSource::Compute(err),
+            Direction::ClientToCompute => ErrorSource::Client(err),
+        }
+    }
+
+    fn write(dir: Direction, err: io::Error) -> Self {
+        match dir {
+            Direction::ComputeToClient => ErrorSource::Client(err),
+            Direction::ClientToCompute => ErrorSource::Compute(err),
         }
     }
 }
 
-#[tracing::instrument(skip_all)]
 pub async fn copy_bidirectional_client_compute<Client, Compute>(
     client: &mut Client,
     compute: &mut Compute,
-) -> Result<(u64, u64), ErrorSource>
+    mut f: impl for<'a> FnMut(Direction, &'a [u8]),
+) -> Result<(), ErrorSource>
 where
     Client: AsyncRead + AsyncWrite + Unpin + ?Sized,
     Compute: AsyncRead + AsyncWrite + Unpin + ?Sized,
 {
-    let mut client_to_compute = TransferState::Running(CopyBuffer::new());
-    let mut compute_to_client = TransferState::Running(CopyBuffer::new());
+    let f = &mut f;
+    let mut client_to_compute = CopyBuffer::new(Direction::ClientToCompute);
+    let mut compute_to_client = CopyBuffer::new(Direction::ComputeToClient);
 
-    poll_fn(|cx| {
-        let mut client_to_compute_result =
-            transfer_one_direction(cx, &mut client_to_compute, client, compute)
-                .map_err(ErrorSource::from_client)?;
-        let mut compute_to_client_result =
-            transfer_one_direction(cx, &mut compute_to_client, compute, client)
-                .map_err(ErrorSource::from_compute)?;
+    let mut client = Pin::new(client);
+    let mut compute = Pin::new(compute);
 
-        // TODO: 1 info log, with a enum label for close direction.
+    // Initial copy hot path
+    let close_dir = poll_fn(|cx| -> Poll<Result<_, ErrorSource>> {
+        let copy1 = client_to_compute.poll_copy(cx, f, client.as_mut(), compute.as_mut())?;
+        let copy2 = compute_to_client.poll_copy(cx, f, compute.as_mut(), client.as_mut())?;
 
-        // Early termination checks from compute to client.
-        if let TransferState::Done(_) = compute_to_client {
-            if let TransferState::Running(buf) = &client_to_compute {
-                info!("Compute is done, terminate client");
-                // Initiate shutdown
-                client_to_compute = TransferState::ShuttingDown(buf.amt);
-                client_to_compute_result =
-                    transfer_one_direction(cx, &mut client_to_compute, client, compute)
-                        .map_err(ErrorSource::from_client)?;
-            }
+        match (copy1, copy2) {
+            (Poll::Pending, Poll::Pending) => Poll::Pending,
+            (Poll::Ready(()), _) => Poll::Ready(Ok(client_to_compute.dir)),
+            (_, Poll::Ready(())) => Poll::Ready(Ok(compute_to_client.dir)),
         }
-
-        // Early termination checks from client to compute.
-        if let TransferState::Done(_) = client_to_compute {
-            if let TransferState::Running(buf) = &compute_to_client {
-                info!("Client is done, terminate compute");
-                // Initiate shutdown
-                compute_to_client = TransferState::ShuttingDown(buf.amt);
-                compute_to_client_result =
-                    transfer_one_direction(cx, &mut compute_to_client, compute, client)
-                        .map_err(ErrorSource::from_compute)?;
-            }
-        }
-
-        // It is not a problem if ready! returns early ... (comment remains the same)
-        let client_to_compute = ready!(client_to_compute_result);
-        let compute_to_client = ready!(compute_to_client_result);
-
-        Poll::Ready(Ok((client_to_compute, compute_to_client)))
     })
-    .await
+    .await?;
+
+    // initiate shutdown.
+    match close_dir {
+        Direction::ClientToCompute => {
+            info!("Client is done, terminate compute");
+
+            // we will never write anymore data to the client.
+            compute_to_client.filled = 0..0;
+
+            // make sure to shutdown the client conn.
+            compute_to_client.need_flush = true;
+        }
+        Direction::ComputeToClient => {
+            info!("Compute is done, terminate client");
+
+            // we will never write anymore data to the compute.
+            client_to_compute.filled = 0..0;
+
+            // make sure to shutdown the compute conn.
+            client_to_compute.need_flush = true;
+        }
+    }
+
+    // Finish sending the rest of the data to client/compute before shutting it down.
+    //
+    // Edge case:
+    // * peer has filled the TCP buffers and is blocking on a `write()`,
+    // * proxy has filled the TCP buffers and is waiting on a `write()`.
+    // Since no side is reading from the buffers, no progress will be made.
+    let shutdown = poll_fn(|cx| {
+        let res1 = client_to_compute.poll_empty(cx, compute.as_mut())?;
+        let res2 = compute_to_client.poll_empty(cx, client.as_mut())?;
+
+        if res1.is_ready() && res2.is_ready() {
+            Poll::Ready(Ok(()))
+        } else {
+            Poll::Pending
+        }
+    });
+
+    // We assume most peers will have enough buffer space so this issue doesn't arise, but we apply
+    // a timeout just in case.
+    //
+    // We could also update `poll_empty` to try and read the data, but I think this is not an edge case
+    // worth overcomplicating.
+    let res = tokio::time::timeout(DISCONNECT_TIMEOUT, shutdown).await;
+
+    match res {
+        Ok(res) => res,
+        Err(timeout) => Err(ErrorSource::Timeout(timeout)),
+    }
 }
 
-#[derive(Debug)]
-pub(super) struct CopyBuffer {
-    read_done: bool,
-    need_flush: bool,
-    pos: usize,
-    cap: usize,
-    amt: u64,
-    buf: Box<[u8]>,
-}
 const DEFAULT_BUF_SIZE: usize = 1024;
+pub(super) struct CopyBuffer {
+    dir: Direction,
+    need_flush: bool,
+    filled: Range<usize>,
+    buf: [u8; DEFAULT_BUF_SIZE],
+}
 
 impl CopyBuffer {
-    pub(super) fn new() -> Self {
+    pub(super) const fn new(dir: Direction) -> Self {
         Self {
-            read_done: false,
+            dir,
             need_flush: false,
-            pos: 0,
-            cap: 0,
-            amt: 0,
-            buf: vec![0; DEFAULT_BUF_SIZE].into_boxed_slice(),
+            filled: 0..0,
+            buf: [0; DEFAULT_BUF_SIZE],
         }
     }
 
-    fn poll_fill_buf<R>(
+    /// Returns Ready(Ok(())) when no more writes could progress, and the buffer has space to read.
+    #[inline(always)]
+    fn poll_write_loop<W>(
         &mut self,
         cx: &mut Context<'_>,
-        reader: Pin<&mut R>,
-    ) -> Poll<io::Result<()>>
-    where
-        R: AsyncRead + ?Sized,
-    {
-        let me = &mut *self;
-        let mut buf = ReadBuf::new(&mut me.buf);
-        buf.set_filled(me.cap);
-
-        let res = reader.poll_read(cx, &mut buf);
-        if let Poll::Ready(Ok(())) = res {
-            let filled_len = buf.filled().len();
-            me.read_done = me.cap == filled_len;
-            me.cap = filled_len;
-        }
-        res
-    }
-
-    fn poll_write_buf<R, W>(
-        &mut self,
-        cx: &mut Context<'_>,
-        mut reader: Pin<&mut R>,
         mut writer: Pin<&mut W>,
-    ) -> Poll<Result<usize, ErrorDirection>>
+    ) -> Poll<Result<(), ErrorSource>>
     where
-        R: AsyncRead + ?Sized,
         W: AsyncWrite + ?Sized,
     {
-        let me = &mut *self;
-        match writer.as_mut().poll_write(cx, &me.buf[me.pos..me.cap]) {
-            Poll::Pending => {
-                // Top up the buffer towards full if we can read a bit more
-                // data - this should improve the chances of a large write
-                if !me.read_done && me.cap < me.buf.len() {
-                    ready!(me.poll_fill_buf(cx, reader.as_mut())).map_err(ErrorDirection::Read)?;
-                }
-                Poll::Pending
-            }
-            res @ Poll::Ready(_) => res.map_err(ErrorDirection::Write),
-        }
-    }
+        debug_assert!(!self.filled.is_empty());
 
-    pub(super) fn poll_copy<R, W>(
-        &mut self,
-        cx: &mut Context<'_>,
-        mut reader: Pin<&mut R>,
-        mut writer: Pin<&mut W>,
-    ) -> Poll<Result<u64, ErrorDirection>>
-    where
-        R: AsyncRead + ?Sized,
-        W: AsyncWrite + ?Sized,
-    {
         loop {
-            // If there is some space left in our buffer, then we try to read some
-            // data to continue, thus maximizing the chances of a large write.
-            if self.cap < self.buf.len() && !self.read_done {
-                match self.poll_fill_buf(cx, reader.as_mut()) {
-                    Poll::Ready(Ok(())) => (),
-                    Poll::Ready(Err(err)) => return Poll::Ready(Err(ErrorDirection::Read(err))),
-                    Poll::Pending => {
-                        // Ignore pending reads when our buffer is not empty, because
-                        // we can try to write data immediately.
-                        if self.pos == self.cap {
-                            // Try flushing when the reader has no progress to avoid deadlock
-                            // when the reader depends on buffered writer.
-                            if self.need_flush {
-                                ready!(writer.as_mut().poll_flush(cx))
-                                    .map_err(ErrorDirection::Write)?;
-                                self.need_flush = false;
-                            }
+            let filled_buf = &self.buf[self.filled.clone()];
+            match writer.as_mut().poll_write(cx, filled_buf) {
+                Poll::Ready(Err(err)) => {
+                    return Poll::Ready(Err(ErrorSource::write(self.dir, cold(err))));
+                }
+                Poll::Ready(Ok(0)) => {
+                    let err =
+                        io::Error::new(io::ErrorKind::WriteZero, "write zero byte into writer");
+                    return Poll::Ready(Err(ErrorSource::write(self.dir, cold(err))));
+                }
+                Poll::Ready(Ok(i)) => {
+                    // update the write head.
+                    self.filled.start += i;
+                    self.need_flush = true;
 
-                            return Poll::Pending;
-                        }
+                    // we wrote some data, but the filled buffer might not be fully empty yet.
+                    if !self.filled.is_empty() {
+                        continue;
                     }
+
+                    // the buffer is definitely empty. reset positions.
+                    self.filled = 0..0;
+                    break;
                 }
-            }
-
-            // If our buffer has some data, let's write it out!
-            while self.pos < self.cap {
-                let i = ready!(self.poll_write_buf(cx, reader.as_mut(), writer.as_mut()))?;
-                if i == 0 {
-                    return Poll::Ready(Err(ErrorDirection::Write(io::Error::new(
-                        io::ErrorKind::WriteZero,
-                        "write zero byte into writer",
-                    ))));
-                }
-                self.pos += i;
-                self.amt += i as u64;
-                self.need_flush = true;
-            }
-
-            // If pos larger than cap, this loop will never stop.
-            // In particular, user's wrong poll_write implementation returning
-            // incorrect written length may lead to thread blocking.
-            debug_assert!(
-                self.pos <= self.cap,
-                "writer returned length larger than input slice"
-            );
-
-            // All data has been written, the buffer can be considered empty again
-            self.pos = 0;
-            self.cap = 0;
-
-            // If we've written all the data and we've seen EOF, flush out the
-            // data and finish the transfer.
-            if self.read_done {
-                ready!(writer.as_mut().poll_flush(cx)).map_err(ErrorDirection::Write)?;
-                return Poll::Ready(Ok(self.amt));
+                // While we couldn't write, we might be able to read.
+                Poll::Pending if self.filled.end < self.buf.len() => break,
+                // We couldn't write, and have no space to read. Just exit.
+                Poll::Pending => return Poll::Pending,
             }
         }
+
+        Poll::Ready(Ok(()))
+    }
+
+    /// Returns Ready(Ok((true))) when read returns EOF.
+    /// Returns Ready(Ok((false))) when read returns data.
+    #[inline(always)]
+    fn poll_read_once<R>(
+        &mut self,
+        cx: &mut Context<'_>,
+        f: &mut impl for<'a> FnMut(Direction, &'a [u8]),
+        reader: Pin<&mut R>,
+    ) -> Poll<Result<bool, ErrorSource>>
+    where
+        R: AsyncRead + ?Sized,
+    {
+        debug_assert!(self.filled.end < self.buf.len());
+        let mut buf = ReadBuf::new(&mut self.buf[self.filled.end..]);
+
+        match reader.poll_read(cx, &mut buf) {
+            Poll::Ready(Ok(())) => {
+                let filled = buf.filled();
+                // no more data to read, switch to shutdown mode.
+                if filled.is_empty() {
+                    self.need_flush = true;
+                    return Poll::Ready(Ok(true));
+                }
+
+                // run our inspection callback.
+                f(self.dir, filled);
+
+                // update the read head.
+                self.filled.end += filled.len();
+
+                // read more data
+                Poll::Ready(Ok(false))
+            }
+            // cannot continue on error.
+            Poll::Ready(Err(e)) => Poll::Ready(Err(ErrorSource::read(self.dir, cold(e)))),
+            // No more data to read, and no more data to write.
+            Poll::Pending => Poll::Pending,
+        }
+    }
+
+    /// Returns Ready(Ok(())) when read returns EOF.
+    fn poll_copy<R, W>(
+        &mut self,
+        cx: &mut Context<'_>,
+        f: &mut impl for<'a> FnMut(Direction, &'a [u8]),
+        mut reader: Pin<&mut R>,
+        mut writer: Pin<&mut W>,
+    ) -> Poll<Result<(), ErrorSource>>
+    where
+        R: AsyncRead + ?Sized,
+        W: AsyncWrite + ?Sized,
+    {
+        // this register eliminates a branch in the hot loop.
+        let mut empty = self.filled.is_empty();
+
+        // write then read hot loop
+        loop {
+            if !empty {
+                ready!(self.poll_write_loop(cx, writer.as_mut())?);
+            }
+
+            // If empty is true, there is guaranteed space to read.
+            // If empty is false, and the write loop returned ready, then we know there's space for more reads.
+            match self.poll_read_once(cx, f, reader.as_mut())? {
+                // EOF
+                Poll::Ready(true) => return Poll::Ready(Ok(())),
+                // Needs write.
+                Poll::Ready(false) => empty = false,
+                // Cannot read. The peer might not send us anything until
+                // they receive data from us, so let's switch to flushing.
+                Poll::Pending => break,
+            }
+        }
+
+        if self.need_flush {
+            let flush = writer.as_mut().poll_flush(cx);
+            ready!(flush.map_err(|e| ErrorSource::write(self.dir, e))?);
+            self.need_flush = false;
+        }
+
+        // there might be more data still to read.
+        Poll::Pending
+    }
+
+    /// Returns Ready(Ok(())) when the conn is fully shutdown.
+    pub(super) fn poll_empty<W>(
+        &mut self,
+        cx: &mut Context<'_>,
+        mut writer: Pin<&mut W>,
+    ) -> Poll<Result<(), ErrorSource>>
+    where
+        W: AsyncWrite + ?Sized,
+    {
+        if !self.filled.is_empty() {
+            ready!(self.poll_write_loop(cx, writer.as_mut())?);
+
+            if !self.filled.is_empty() {
+                // still some data to write
+                return Poll::Pending;
+            }
+        }
+
+        if self.need_flush {
+            let res = writer.poll_shutdown(cx);
+            ready!(res.map_err(|e| ErrorSource::write(self.dir, e))?);
+            self.need_flush = false;
+        }
+
+        // no data to read, no data to write.
+        Poll::Ready(Ok(()))
     }
 }
 
 #[cfg(test)]
 mod tests {
-    use tokio::io::AsyncWriteExt;
+    use tokio::io::{AsyncReadExt, AsyncWriteExt};
 
     use super::*;
 
     #[tokio::test]
     async fn test_client_to_compute() {
-        let (mut client_client, mut client_proxy) = tokio::io::duplex(8); // Create a mock duplex stream
-        let (mut compute_proxy, mut compute_client) = tokio::io::duplex(32); // Create a mock duplex stream
+        let (mut client, mut client_proxy) = tokio::io::duplex(32); // Create a mock duplex stream
+        let (mut proxy_compute, mut compute) = tokio::io::duplex(16); // Create a mock duplex stream
 
-        // Simulate 'a' finishing while there's still data for 'b'
-        client_client.write_all(b"hello").await.unwrap();
-        client_client.shutdown().await.unwrap();
-        compute_client.write_all(b"Neon").await.unwrap();
-        compute_client.shutdown().await.unwrap();
+        client.write_all(b"Neon Serverless Postgres").await.unwrap();
+        compute.write_all(b"is amazing").await.unwrap();
 
-        let result = copy_bidirectional_client_compute(&mut client_proxy, &mut compute_proxy)
-            .await
-            .unwrap();
+        client.shutdown().await.unwrap();
+
+        let copy = tokio::spawn(async move {
+            copy_bidirectional_client_compute(&mut client_proxy, &mut proxy_compute, |_, _| {})
+                .await
+                .unwrap();
+        });
 
         // Assert correct transferred amounts
-        let (client_to_compute_count, compute_to_client_count) = result;
-        assert_eq!(client_to_compute_count, 5); // 'hello' was transferred
-        assert_eq!(compute_to_client_count, 4); // response only partially transferred or not at all
+        let mut client_recv = String::new();
+        let mut compute_recv = String::new();
+
+        client.read_to_string(&mut client_recv).await.unwrap();
+        compute.read_to_string(&mut compute_recv).await.unwrap();
+
+        assert_eq!(compute_recv, "Neon Serverless Postgres");
+        assert_eq!(client_recv, "is amazing");
+
+        copy.await.unwrap();
     }
 
     #[tokio::test]
     async fn test_compute_to_client() {
-        let (mut client_client, mut client_proxy) = tokio::io::duplex(32); // Create a mock duplex stream
-        let (mut compute_proxy, mut compute_client) = tokio::io::duplex(8); // Create a mock duplex stream
+        let (mut client, mut client_proxy) = tokio::io::duplex(32); // Create a mock duplex stream
+        let (mut proxy_compute, mut compute) = tokio::io::duplex(16); // Create a mock duplex stream
 
-        // Simulate 'a' finishing while there's still data for 'b'
-        compute_client.write_all(b"hello").await.unwrap();
-        compute_client.shutdown().await.unwrap();
-        client_client
-            .write_all(b"Neon Serverless Postgres")
-            .await
-            .unwrap();
+        client.write_all(b"Neon Serverless Postgres").await.unwrap();
+        compute.write_all(b"is amazing").await.unwrap();
 
-        let result = copy_bidirectional_client_compute(&mut client_proxy, &mut compute_proxy)
-            .await
-            .unwrap();
+        compute.shutdown().await.unwrap();
+
+        let copy = tokio::spawn(async move {
+            copy_bidirectional_client_compute(&mut client_proxy, &mut proxy_compute, |_, _| {})
+                .await
+                .unwrap();
+        });
 
         // Assert correct transferred amounts
-        let (client_to_compute_count, compute_to_client_count) = result;
-        assert_eq!(compute_to_client_count, 5); // 'hello' was transferred
-        assert!(client_to_compute_count <= 8); // response only partially transferred or not at all
+        let mut client_recv = String::new();
+        let mut compute_recv = String::new();
+
+        client.read_to_string(&mut client_recv).await.unwrap();
+        compute.read_to_string(&mut compute_recv).await.unwrap();
+
+        assert_eq!(compute_recv, "Neon Serverless ");
+        assert_eq!(client_recv, "is amazing");
+
+        copy.await.unwrap();
+    }
+
+    #[tokio::test(start_paused = true)]
+    async fn test_timeout() {
+        let (mut client, mut client_proxy) = tokio::io::duplex(32); // Create a mock duplex stream
+        let (mut proxy_compute, mut compute) = tokio::io::duplex(16); // Create a mock duplex stream
+
+        // Try to send 24 bytes to compute, but compute only has space for 16 bytes.
+        // Writes will not succeed.
+        client.write_all(b"Neon Serverless Postgres").await.unwrap();
+        client.shutdown().await.unwrap();
+
+        let copy = tokio::spawn(async move {
+            copy_bidirectional_client_compute(&mut client_proxy, &mut proxy_compute, |_, _| {})
+                .await
+                .unwrap_err()
+        });
+
+        tokio::time::advance(DISCONNECT_TIMEOUT).await;
+
+        let res = copy.await.unwrap();
+        assert!(matches!(res, ErrorSource::Timeout(_)));
+
+        // Assert correct transferred amounts
+        let mut compute_recv = String::new();
+        compute.read_to_string(&mut compute_recv).await.unwrap();
+        assert_eq!(compute_recv, "Neon Serverless ");
     }
 }

proxy/src/proxy/mod.rs

@@ -167,6 +167,12 @@ pub async fn task_main(
                     let _disconnect = ctx.log_connect();
                     match p.proxy_pass(&config.connect_to_compute).await {
                         Ok(()) => {}
+                        Err(ErrorSource::Timeout(_)) => {
+                            info!(
+                                ?session_id,
+                                "per-client task timed out while gracefully shutting down the connection"
+                            );
+                        }
                         Err(ErrorSource::Client(e)) => {
                             warn!(
                                 ?session_id,
@@ -383,9 +389,7 @@ pub(crate) async fn handle_client<S: AsyncRead + AsyncWrite + Unpin>(
     let cancellation_handler_clone = Arc::clone(&cancellation_handler);
     let session = cancellation_handler_clone.get_key();
 
-    session
-        .write_cancel_key(node.cancel_closure.clone())
-        .await?;
+    session.write_cancel_key(node.cancel_closure.clone())?;
 
     prepare_client_connection(&node, *session.key(), &mut stream).await?;
 

proxy/src/proxy/passthrough.rs

@@ -1,7 +1,6 @@
 use smol_str::SmolStr;
 use tokio::io::{AsyncRead, AsyncWrite};
 use tracing::debug;
-use utils::measured_stream::MeasuredStream;
 
 use super::copy_bidirectional::ErrorSource;
 use crate::cancellation;
@@ -9,14 +8,15 @@ use crate::compute::PostgresConnection;
 use crate::config::ComputeConfig;
 use crate::control_plane::messages::MetricsAuxInfo;
 use crate::metrics::{Direction, Metrics, NumClientConnectionsGuard, NumConnectionRequestsGuard};
+use crate::proxy::copy_bidirectional_client_compute;
 use crate::stream::Stream;
 use crate::usage_metrics::{Ids, MetricCounterRecorder, USAGE_METRICS};
 
 /// Forward bytes in both directions (client <-> compute).
-#[tracing::instrument(skip_all)]
+#[tracing::instrument(level = "debug", skip_all)]
 pub(crate) async fn proxy_pass(
-    client: impl AsyncRead + AsyncWrite + Unpin,
-    compute: impl AsyncRead + AsyncWrite + Unpin,
+    mut client: Stream<impl AsyncRead + AsyncWrite + Unpin>,
+    mut compute: impl AsyncRead + AsyncWrite + Unpin,
     aux: MetricsAuxInfo,
     private_link_id: Option<SmolStr>,
 ) -> Result<(), ErrorSource> {
@@ -28,37 +28,30 @@ pub(crate) async fn proxy_pass(
     });
 
     let metrics = &Metrics::get().proxy.io_bytes;
-    let m_sent = metrics.with_labels(Direction::Tx);
-    let mut client = MeasuredStream::new(
-        client,
-        |_| {},
-        |cnt| {
-            // Number of bytes we sent to the client (outbound).
-            metrics.get_metric(m_sent).inc_by(cnt as u64);
-            usage_tx.record_egress(cnt as u64);
-        },
-    );
+    let m_sent = metrics.with_labels(Direction::ComputeToClient);
+    let m_recv = metrics.with_labels(Direction::ClientToCompute);
 
-    let m_recv = metrics.with_labels(Direction::Rx);
-    let mut compute = MeasuredStream::new(
-        compute,
-        |_| {},
-        |cnt| {
-            // Number of bytes the client sent to the compute node (inbound).
-            metrics.get_metric(m_recv).inc_by(cnt as u64);
-            usage_tx.record_ingress(cnt as u64);
-        },
-    );
+    let inspect = |direction, bytes: &[u8]| match direction {
+        Direction::ComputeToClient => {
+            metrics.get_metric(m_sent).inc_by(bytes.len() as u64);
+            usage_tx.record_egress(bytes.len() as u64);
+        }
+        Direction::ClientToCompute => {
+            metrics.get_metric(m_recv).inc_by(bytes.len() as u64);
+            usage_tx.record_ingress(bytes.len() as u64);
+        }
+    };
 
     // Starting from here we only proxy the client's traffic.
     debug!("performing the proxy pass...");
-    let _ = crate::proxy::copy_bidirectional::copy_bidirectional_client_compute(
-        &mut client,
-        &mut compute,
-    )
-    .await?;
 
-    Ok(())
+    // reduce branching internal to the hot path.
+    match &mut client {
+        Stream::Raw { raw } => copy_bidirectional_client_compute(raw, &mut compute, inspect).await,
+        Stream::Tls { tls, .. } => {
+            copy_bidirectional_client_compute(&mut *tls, &mut compute, inspect).await
+        }
+    }
 }
 
 pub(crate) struct ProxyPassthrough<S> {
@@ -94,7 +87,7 @@ impl<S: AsyncRead + AsyncWrite + Unpin> ProxyPassthrough<S> {
             tracing::warn!(session_id = ?self.session_id, ?err, "could not cancel the query in the database");
         }
 
-        drop(self.cancel.remove_cancel_key().await); // we don't need a result. If the queue is full, we just log the error
+        drop(self.cancel.remove_cancel_key()); // we don't need a result. If the queue is full, we just log the error
 
         res
     }

proxy/src/serverless/websocket.rs

@@ -2,7 +2,7 @@ use std::pin::Pin;
 use std::sync::Arc;
 use std::task::{Context, Poll, ready};
 
-use anyhow::Context as _;
+use anyhow::{Context as _, anyhow};
 use bytes::{Buf, BufMut, Bytes, BytesMut};
 use framed_websockets::{Frame, OpCode, WebSocketServer};
 use futures::{Sink, Stream};
@@ -169,6 +169,9 @@ pub(crate) async fn serve_websocket(
             ctx.log_connect();
             match p.proxy_pass(&config.connect_to_compute).await {
                 Ok(()) => Ok(()),
+                Err(ErrorSource::Timeout(_)) => Err(anyhow!(
+                    "timed out while gracefully shutting down the connection"
+                )),
                 Err(ErrorSource::Client(err)) => Err(err).context("client"),
                 Err(ErrorSource::Compute(err)) => Err(err).context("compute"),
             }

test_runner/fixtures/neon_cli.py

@@ -103,7 +103,7 @@ class AbstractNeonCli:
             else:
                 stdout = ""
 
-            log.warn(f"CLI timeout: stderr={stderr}, stdout={stdout}")
+            log.warning(f"CLI timeout: stderr={stderr}, stdout={stdout}")
             raise
 
         indent = "  "

test_runner/regress/test_pageserver_secondary.py

@@ -510,7 +510,7 @@ def list_elegible_layers(
         except KeyError:
             # Unexpected: tests should call this when pageservers are in a quiet state such that the layer map
             # matches what's on disk.
-            log.warn(f"Lookup {layer_file_name} from {list(visible_map.keys())}")
+            log.warning(f"Lookup {layer_file_name} from {list(visible_map.keys())}")
             raise
 
     return list(c for c in candidates if is_visible(c))
@@ -636,7 +636,7 @@ def test_secondary_downloads(neon_env_builder: NeonEnvBuilder):
     except:
         # On assertion failures, log some details to help with debugging
         heatmap = env.pageserver_remote_storage.heatmap_content(tenant_id)
-        log.warn(f"heatmap contents: {json.dumps(heatmap, indent=2)}")
+        log.warning(f"heatmap contents: {json.dumps(heatmap, indent=2)}")
         raise
 
     # Scrub the remote storage

vendor/revisions.json

@@ -1,18 +1,18 @@
 {
   "v17": [
     "17.5",
-    "e5374b72997b0afc8374137674e873f7a558120a"
+    "8be779fd3ab9e87206da96a7e4842ef1abf04f44"
   ],
   "v16": [
     "16.9",
-    "15710a76b7d07912110fcbbaf0c8ad6d7e5a9fbc"
+    "0bf96bd6d70301a0b43b0b3457bb3cf8fb43c198"
   ],
   "v15": [
     "15.13",
-    "daa81cffcf063c54b29a9aabdb6604625f675ad0"
+    "de7640f55da07512834d5cc40c4b3fb376b5f04f"
   ],
   "v14": [
     "14.18",
-    "4cca6f8083483dda9e12eae292cf788d45bd561f"
+    "55c0d45abe6467c02084c2192bca117eda6ce1e7"
   ]
 }