From 53743991decd9f1d13fd5063a8e840a38cbda383 Mon Sep 17 00:00:00 2001
From: Joonas Koivunen <joonas@neon.tech>
Date: Tue, 6 Feb 2024 13:34:13 +0200
Subject: [PATCH] uploader: avoid cloning vecs just to get Bytes (#6645)

Fix cloning the serialized heatmap on every attempt by just turning it
into `bytes::Bytes` before clone so it will be a refcounted instead of
refcounting a vec clone later on.

Also fixes one cancellation token cloning I had missed in #6618.
Cc: #6096
---
 .../src/tenant/secondary/heatmap_uploader.rs       | 14 ++++++--------
 1 file changed, 6 insertions(+), 8 deletions(-)

diff --git a/pageserver/src/tenant/secondary/heatmap_uploader.rs b/pageserver/src/tenant/secondary/heatmap_uploader.rs
index fff29b2487..806e3fb0e8 100644
--- a/pageserver/src/tenant/secondary/heatmap_uploader.rs
+++ b/pageserver/src/tenant/secondary/heatmap_uploader.rs
@@ -371,8 +371,6 @@ async fn upload_tenant_heatmap(
     };
     let timelines = tenant.timelines.lock().unwrap().clone();
 
-    let tenant_cancel = tenant.cancel.clone();
-
     // Ensure that Tenant::shutdown waits for any upload in flight: this is needed because otherwise
     // when we delete a tenant, we might race with an upload in flight and end up leaving a heatmap behind
     // in remote storage.
@@ -401,6 +399,7 @@ async fn upload_tenant_heatmap(
 
     // Serialize the heatmap
     let bytes = serde_json::to_vec(&heatmap).map_err(|e| anyhow::anyhow!(e))?;
+    let bytes = bytes::Bytes::from(bytes);
     let size = bytes.len();
 
     // Drop out early if nothing changed since our last upload
@@ -411,13 +410,12 @@ async fn upload_tenant_heatmap(
 
     let path = remote_heatmap_path(tenant.get_tenant_shard_id());
 
-    // Write the heatmap.
+    let cancel = &tenant.cancel;
+
     tracing::debug!("Uploading {size} byte heatmap to {path}");
     if let Err(e) = backoff::retry(
         || async {
-            let bytes = futures::stream::once(futures::future::ready(Ok(bytes::Bytes::from(
-                bytes.clone(),
-            ))));
+            let bytes = futures::stream::once(futures::future::ready(Ok(bytes.clone())));
             remote_storage
                 .upload_storage_object(bytes, size, &path)
                 .await
@@ -426,13 +424,13 @@ async fn upload_tenant_heatmap(
         3,
         u32::MAX,
         "Uploading heatmap",
-        &tenant_cancel,
+        cancel,
     )
     .await
     .ok_or_else(|| anyhow::anyhow!("Shutting down"))
     .and_then(|x| x)
     {
-        if tenant_cancel.is_cancelled() {
+        if cancel.is_cancelled() {
             return Err(UploadHeatmapError::Cancelled);
         } else {
             return Err(e.into());