From fcc3acfd7eb3ac4cef4a189162cf3a893f2a682d Mon Sep 17 00:00:00 2001
From: Heikki Linnakangas <heikki@neon.tech>
Date: Mon, 8 May 2023 01:59:43 +0300
Subject: [PATCH] Add a tracing span to where we open connection to Postgres.

Local testing shows that this can take up to 100 ms, so it'd be nice
to include it explicitly in the trace.
---
 compute_tools/src/compute.rs | 37 +++++++++++++++++++-----------------
 1 file changed, 20 insertions(+), 17 deletions(-)

diff --git a/compute_tools/src/compute.rs b/compute_tools/src/compute.rs
index da5ad00da6..ca805f5e89 100644
--- a/compute_tools/src/compute.rs
+++ b/compute_tools/src/compute.rs
@@ -338,27 +338,30 @@ impl ComputeNode {
         // In this case we need to connect with old `zenith_admin` name
         // and create new user. We cannot simply rename connected user,
         // but we can create a new one and grant it all privileges.
-        let mut client = match Client::connect(self.connstr.as_str(), NoTls) {
-            Err(e) => {
-                info!(
-                    "cannot connect to postgres: {}, retrying with `zenith_admin` username",
-                    e
-                );
-                let mut zenith_admin_connstr = self.connstr.clone();
+        let mut client = {
+            let _span = tracing::info_span!("connect").entered();
+            match Client::connect(self.connstr.as_str(), NoTls) {
+                Err(e) => {
+                    info!(
+                        "cannot connect to postgres: {}, retrying with `zenith_admin` username",
+                        e
+                    );
+                    let mut zenith_admin_connstr = self.connstr.clone();
 
-                zenith_admin_connstr
-                    .set_username("zenith_admin")
-                    .map_err(|_| anyhow::anyhow!("invalid connstr"))?;
+                    zenith_admin_connstr
+                        .set_username("zenith_admin")
+                        .map_err(|_| anyhow::anyhow!("invalid connstr"))?;
 
-                let mut client = Client::connect(zenith_admin_connstr.as_str(), NoTls)?;
-                client.simple_query("CREATE USER cloud_admin WITH SUPERUSER")?;
-                client.simple_query("GRANT zenith_admin TO cloud_admin")?;
-                drop(client);
+                    let mut client = Client::connect(zenith_admin_connstr.as_str(), NoTls)?;
+                    client.simple_query("CREATE USER cloud_admin WITH SUPERUSER")?;
+                    client.simple_query("GRANT zenith_admin TO cloud_admin")?;
+                    drop(client);
 
-                // reconnect with connsting with expected name
-                Client::connect(self.connstr.as_str(), NoTls)?
+                    // reconnect with connsting with expected name
+                    Client::connect(self.connstr.as_str(), NoTls)?
+                }
+                Ok(client) => client,
             }
-            Ok(client) => client,
         };
 
         // Proceed with post-startup configuration. Note, that order of operations is important.