From 896d51367ecb17773677b5f845803dc3c6aa2a70 Mon Sep 17 00:00:00 2001
From: Anna Khanova <32508607+khanova@users.noreply.github.com>
Date: Tue, 27 Feb 2024 19:53:02 +0400
Subject: [PATCH] proxy: introdice is cold start for analytics (#6902)

## Problem

Data team cannot distinguish between cold start and not cold start.

## Summary of changes

Report `is_cold_start` to analytics.

---------

Co-authored-by: Conrad Ludgate <conrad@neon.tech>
---
 proxy/src/console/messages.rs |  1 +
 proxy/src/context.rs          |  3 ++
 proxy/src/context/parquet.rs  | 62 +++++++++++++++++++----------------
 3 files changed, 37 insertions(+), 29 deletions(-)

diff --git a/proxy/src/console/messages.rs b/proxy/src/console/messages.rs
index 4e5920436f..1f94059f1e 100644
--- a/proxy/src/console/messages.rs
+++ b/proxy/src/console/messages.rs
@@ -98,6 +98,7 @@ pub struct MetricsAuxInfo {
     pub endpoint_id: EndpointId,
     pub project_id: ProjectId,
     pub branch_id: BranchId,
+    pub is_cold_start: Option<bool>,
 }
 
 #[cfg(test)]
diff --git a/proxy/src/context.rs b/proxy/src/context.rs
index e5caa5bd59..4d8ced6f8f 100644
--- a/proxy/src/context.rs
+++ b/proxy/src/context.rs
@@ -40,6 +40,7 @@ pub struct RequestMonitoring {
     error_kind: Option<ErrorKind>,
     pub(crate) auth_method: Option<AuthMethod>,
     success: bool,
+    is_cold_start: Option<bool>,
 
     // extra
     // This sender is here to keep the request monitoring channel open while requests are taking place.
@@ -79,6 +80,7 @@ impl RequestMonitoring {
             error_kind: None,
             auth_method: None,
             success: false,
+            is_cold_start: None,
 
             sender: LOG_CHAN.get().and_then(|tx| tx.upgrade()),
             latency_timer: LatencyTimer::new(protocol),
@@ -102,6 +104,7 @@ impl RequestMonitoring {
         self.branch = Some(x.branch_id);
         self.endpoint_id = Some(x.endpoint_id);
         self.project = Some(x.project_id);
+        self.is_cold_start = x.is_cold_start;
     }
 
     pub fn set_project_id(&mut self, project_id: ProjectId) {
diff --git a/proxy/src/context/parquet.rs b/proxy/src/context/parquet.rs
index d941445c2d..54f51604bf 100644
--- a/proxy/src/context/parquet.rs
+++ b/proxy/src/context/parquet.rs
@@ -92,6 +92,8 @@ struct RequestData {
     /// Success is counted if we form a HTTP response with sql rows inside
     /// Or if we make it to proxy_pass
     success: bool,
+    /// Indicates if the cplane started the new compute node for this request.
+    is_cold_start: Option<bool>,
     /// Tracks time from session start (HTTP request/libpq TCP handshake)
     /// Through to success/failure
     duration_us: u64,
@@ -119,6 +121,7 @@ impl From<RequestMonitoring> for RequestData {
             region: value.region,
             error: value.error_kind.as_ref().map(|e| e.to_metric_label()),
             success: value.success,
+            is_cold_start: value.is_cold_start,
             duration_us: SystemTime::from(value.first_packet)
                 .elapsed()
                 .unwrap_or_default()
@@ -452,6 +455,7 @@ mod tests {
             region: "us-east-1",
             error: None,
             success: rng.gen(),
+            is_cold_start: Some(true),
             duration_us: rng.gen_range(0..30_000_000),
         }
     }
@@ -521,15 +525,15 @@ mod tests {
         assert_eq!(
             file_stats,
             [
-                (1313727, 3, 6000),
-                (1313720, 3, 6000),
-                (1313780, 3, 6000),
-                (1313737, 3, 6000),
-                (1313867, 3, 6000),
-                (1313709, 3, 6000),
-                (1313501, 3, 6000),
-                (1313737, 3, 6000),
-                (438118, 1, 2000)
+                (1315032, 3, 6000),
+                (1315025, 3, 6000),
+                (1315085, 3, 6000),
+                (1315042, 3, 6000),
+                (1315172, 3, 6000),
+                (1315014, 3, 6000),
+                (1314806, 3, 6000),
+                (1315042, 3, 6000),
+                (438563, 1, 2000)
             ],
         );
 
@@ -559,11 +563,11 @@ mod tests {
         assert_eq!(
             file_stats,
             [
-                (1219459, 5, 10000),
-                (1225609, 5, 10000),
-                (1227403, 5, 10000),
-                (1226765, 5, 10000),
-                (1218043, 5, 10000)
+                (1220433, 5, 10000),
+                (1226583, 5, 10000),
+                (1228377, 5, 10000),
+                (1227739, 5, 10000),
+                (1219017, 5, 10000)
             ],
         );
 
@@ -595,11 +599,11 @@ mod tests {
         assert_eq!(
             file_stats,
             [
-                (1205106, 5, 10000),
-                (1204837, 5, 10000),
-                (1205130, 5, 10000),
-                (1205118, 5, 10000),
-                (1205373, 5, 10000)
+                (1206080, 5, 10000),
+                (1205811, 5, 10000),
+                (1206104, 5, 10000),
+                (1206092, 5, 10000),
+                (1206347, 5, 10000)
             ],
         );
 
@@ -624,15 +628,15 @@ mod tests {
         assert_eq!(
             file_stats,
             [
-                (1313727, 3, 6000),
-                (1313720, 3, 6000),
-                (1313780, 3, 6000),
-                (1313737, 3, 6000),
-                (1313867, 3, 6000),
-                (1313709, 3, 6000),
-                (1313501, 3, 6000),
-                (1313737, 3, 6000),
-                (438118, 1, 2000)
+                (1315032, 3, 6000),
+                (1315025, 3, 6000),
+                (1315085, 3, 6000),
+                (1315042, 3, 6000),
+                (1315172, 3, 6000),
+                (1315014, 3, 6000),
+                (1314806, 3, 6000),
+                (1315042, 3, 6000),
+                (438563, 1, 2000)
             ],
         );
 
@@ -669,7 +673,7 @@ mod tests {
         // files are smaller than the size threshold, but they took too long to fill so were flushed early
         assert_eq!(
             file_stats,
-            [(658383, 2, 3001), (658097, 2, 3000), (657893, 2, 2999)],
+            [(659129, 2, 3001), (658842, 2, 3000), (658638, 2, 2999)],
         );
 
         tmpdir.close().unwrap();