From fc3c7aaa6c4efc0a4c515f99ee43aab7cadae13d Mon Sep 17 00:00:00 2001
From: John Spray <john@neon.tech>
Date: Thu, 2 Nov 2023 15:22:18 +0000
Subject: [PATCH] pageserver: add TenantConf::enable_heatmap (default false)

---
 control_plane/src/pageserver.rs      | 10 ++++++++++
 libs/pageserver_api/src/models.rs    |  2 ++
 pageserver/src/http/openapi_spec.yml |  2 ++
 pageserver/src/tenant.rs             |  8 ++++++++
 pageserver/src/tenant/config.rs      | 13 +++++++++++++
 5 files changed, 35 insertions(+)

diff --git a/control_plane/src/pageserver.rs b/control_plane/src/pageserver.rs
index 0746dde4ef..f2785ea99c 100644
--- a/control_plane/src/pageserver.rs
+++ b/control_plane/src/pageserver.rs
@@ -389,6 +389,11 @@ impl PageServerNode {
                 .map(|x| x.parse::<bool>())
                 .transpose()
                 .context("Failed to parse 'gc_feedback' as bool")?,
+            enable_heatmap: settings
+                .remove("enable_heatmap")
+                .map(|x| x.parse::<bool>())
+                .transpose()
+                .context("Failed to parse 'enable_heatmap' as bool")?,
         };
 
         let request = models::TenantCreateRequest {
@@ -486,6 +491,11 @@ impl PageServerNode {
                     .map(|x| x.parse::<bool>())
                     .transpose()
                     .context("Failed to parse 'gc_feedback' as bool")?,
+                enable_heatmap: settings
+                    .remove("enable_heatmap")
+                    .map(|x| x.parse::<bool>())
+                    .transpose()
+                    .context("Failed to parse 'enable_heatmap' as bool")?,
             }
         };
 
diff --git a/libs/pageserver_api/src/models.rs b/libs/pageserver_api/src/models.rs
index cb99dc0a55..a7ab77ca2a 100644
--- a/libs/pageserver_api/src/models.rs
+++ b/libs/pageserver_api/src/models.rs
@@ -236,6 +236,7 @@ pub struct TenantConfig {
     pub min_resident_size_override: Option<u64>,
     pub evictions_low_residence_duration_metric_threshold: Option<String>,
     pub gc_feedback: Option<bool>,
+    pub enable_heatmap: Option<bool>,
 }
 
 /// A flattened analog of a `pagesever::tenant::LocationMode`, which
@@ -324,6 +325,7 @@ impl TenantConfigRequest {
             min_resident_size_override: None,
             evictions_low_residence_duration_metric_threshold: None,
             gc_feedback: None,
+            enable_heatmap: None,
         };
         TenantConfigRequest { tenant_id, config }
     }
diff --git a/pageserver/src/http/openapi_spec.yml b/pageserver/src/http/openapi_spec.yml
index 4d455243f0..a12a81e831 100644
--- a/pageserver/src/http/openapi_spec.yml
+++ b/pageserver/src/http/openapi_spec.yml
@@ -1303,6 +1303,8 @@ components:
           type: integer
         trace_read_requests:
           type: boolean
+        enable_heatmap:
+          type: boolean
     TenantConfigResponse:
       type: object
       properties:
diff --git a/pageserver/src/tenant.rs b/pageserver/src/tenant.rs
index 758f8b15a1..b72e15667b 100644
--- a/pageserver/src/tenant.rs
+++ b/pageserver/src/tenant.rs
@@ -2208,6 +2208,13 @@ impl Tenant {
             .or(self.conf.default_tenant_conf.min_resident_size_override)
     }
 
+    pub fn get_enable_heatmap(&self) -> bool {
+        let tenant_conf = self.tenant_conf.read().unwrap().tenant_conf;
+        tenant_conf
+            .enable_heatmap
+            .unwrap_or(self.conf.default_tenant_conf.enable_heatmap)
+    }
+
     pub fn set_new_tenant_config(&self, new_tenant_conf: TenantConfOpt) {
         self.tenant_conf.write().unwrap().tenant_conf = new_tenant_conf;
         // Don't hold self.timelines.lock() during the notifies.
@@ -3472,6 +3479,7 @@ pub(crate) mod harness {
                     tenant_conf.evictions_low_residence_duration_metric_threshold,
                 ),
                 gc_feedback: Some(tenant_conf.gc_feedback),
+                enable_heatmap: Some(tenant_conf.enable_heatmap),
             }
         }
     }
diff --git a/pageserver/src/tenant/config.rs b/pageserver/src/tenant/config.rs
index 5f8c7f6c59..5b52e5308c 100644
--- a/pageserver/src/tenant/config.rs
+++ b/pageserver/src/tenant/config.rs
@@ -305,6 +305,11 @@ pub struct TenantConf {
     #[serde(with = "humantime_serde")]
     pub evictions_low_residence_duration_metric_threshold: Duration,
     pub gc_feedback: bool,
+
+    /// Whether to upload a heatmap to remote storage, for use by secondary mode.
+    /// This may be left disabled if a Tenant will only every be attached to
+    /// one node.
+    pub enable_heatmap: bool,
 }
 
 /// Same as TenantConf, but this struct preserves the information about
@@ -385,6 +390,10 @@ pub struct TenantConfOpt {
     #[serde(skip_serializing_if = "Option::is_none")]
     #[serde(default)]
     pub gc_feedback: Option<bool>,
+
+    #[serde(skip_serializing_if = "Option::is_none")]
+    #[serde(default)]
+    pub enable_heatmap: Option<bool>,
 }
 
 #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
@@ -453,6 +462,7 @@ impl TenantConfOpt {
                 .evictions_low_residence_duration_metric_threshold
                 .unwrap_or(global_conf.evictions_low_residence_duration_metric_threshold),
             gc_feedback: self.gc_feedback.unwrap_or(global_conf.gc_feedback),
+            enable_heatmap: self.enable_heatmap.unwrap_or(global_conf.enable_heatmap),
         }
     }
 }
@@ -490,6 +500,7 @@ impl Default for TenantConf {
             )
             .expect("cannot parse default evictions_low_residence_duration_metric_threshold"),
             gc_feedback: false,
+            enable_heatmap: false,
         }
     }
 }
@@ -586,6 +597,8 @@ impl TryFrom<&'_ models::TenantConfig> for TenantConfOpt {
         }
         tenant_conf.gc_feedback = request_data.gc_feedback;
 
+        tenant_conf.enable_heatmap = request_data.enable_heatmap;
+
         Ok(tenant_conf)
     }
 }