jwt support; debug spans in basebackup

pageserver/src/basebackup.rs

@@ -166,71 +166,111 @@ where
             }
         }
 
-        // Gather non-relational files from object storage pages.
+        debug!("Gather non-relational files from object storage pages");
         for kind in [
             SlruKind::Clog,
             SlruKind::MultiXactOffsets,
             SlruKind::MultiXactMembers,
         ] {
-            for segno in self
-                .timeline
-                .list_slru_segments(kind, self.lsn, self.ctx)
-                .await?
-            {
-                self.add_slru_segment(kind, segno).await?;
+            async {
+                debug!("list slru segments");
+                for segno in self
+                    .timeline
+                    .list_slru_segments(kind, self.lsn, self.ctx)
+                    .await?
+                {
+                    async {
+                        debug!("add slru segment");
+                        self.add_slru_segment(kind, segno).await?;
+                        anyhow::Ok(())
+                    }
+                    .instrument(debug_span!("slru segment", ?segno))
+                    .await?;
+                }
+                anyhow::Ok(())
             }
+            .instrument(debug_span!("non-rel file", ?kind))
+            .await?;
         }
 
         let mut min_restart_lsn: Lsn = Lsn::MAX;
-        // Create tablespace directories
+        debug!("Create tablespace directories");
         for ((spcnode, dbnode), has_relmap_file) in
             self.timeline.list_dbdirs(self.lsn, self.ctx).await?
         {
-            self.add_dbdir(spcnode, dbnode, has_relmap_file).await?;
+            async {
+                debug!("iter");
+                self.add_dbdir(spcnode, dbnode, has_relmap_file).await?;
 
-            // If full backup is requested, include all relation files.
-            // Otherwise only include init forks of unlogged relations.
-            let rels = self
-                .timeline
-                .list_rels(spcnode, dbnode, self.lsn, self.ctx)
-                .await?;
-            for &rel in rels.iter() {
-                // Send init fork as main fork to provide well formed empty
-                // contents of UNLOGGED relations. Postgres copies it in
-                // `reinit.c` during recovery.
-                if rel.forknum == INIT_FORKNUM {
-                    // I doubt we need _init fork itself, but having it at least
-                    // serves as a marker relation is unlogged.
-                    self.add_rel(rel, rel).await?;
-                    self.add_rel(rel, rel.with_forknum(MAIN_FORKNUM)).await?;
-                    continue;
-                }
+                // If full backup is requested, include all relation files.
+                // Otherwise only include init forks of unlogged relations.
+                debug!("list rels");
+                let rels = self
+                    .timeline
+                    .list_rels(spcnode, dbnode, self.lsn, self.ctx)
+                    .await?;
+                for &rel in rels.iter() {
+                    async {
+                        debug!("iter");
+                        // Send init fork as main fork to provide well formed empty
+                        // contents of UNLOGGED relations. Postgres copies it in
+                        // `reinit.c` during recovery.
+                        if rel.forknum == INIT_FORKNUM {
+                            // I doubt we need _init fork itself, but having it at least
+                            // serves as a marker relation is unlogged.
+                            self.add_rel(rel, rel).await?;
+                            self.add_rel(rel, rel.with_forknum(MAIN_FORKNUM)).await?;
+                            return Ok(());
+                        }
 
-                if self.full_backup {
-                    if rel.forknum == MAIN_FORKNUM && rels.contains(&rel.with_forknum(INIT_FORKNUM))
-                    {
-                        // skip this, will include it when we reach the init fork
-                        continue;
+                        if self.full_backup {
+                            if rel.forknum == MAIN_FORKNUM
+                                && rels.contains(&rel.with_forknum(INIT_FORKNUM))
+                            {
+                                // skip this, will include it when we reach the init fork
+                                return Ok(());
+                            }
+                            self.add_rel(rel, rel).await?;
+                        }
+                        anyhow::Ok(())
                     }
-                    self.add_rel(rel, rel).await?;
+                    .instrument(debug_span!("process rel", ?rel))
+                    .await?;
                 }
-            }
 
-            for (path, content) in self.timeline.list_aux_files(self.lsn, self.ctx).await? {
-                if path.starts_with("pg_replslot") {
-                    let offs = pg_constants::REPL_SLOT_ON_DISK_OFFSETOF_RESTART_LSN;
-                    let restart_lsn = Lsn(u64::from_le_bytes(
-                        content[offs..offs + 8].try_into().unwrap(),
-                    ));
-                    info!("Replication slot {} restart LSN={}", path, restart_lsn);
-                    min_restart_lsn = Lsn::min(min_restart_lsn, restart_lsn);
+                debug!("list aux files");
+                for (path, content) in self.timeline.list_aux_files(self.lsn, self.ctx).await? {
+                    async {
+                        debug!("iter");
+                        if path.starts_with("pg_replslot") {
+                            let offs = pg_constants::REPL_SLOT_ON_DISK_OFFSETOF_RESTART_LSN;
+                            let restart_lsn = Lsn(u64::from_le_bytes(
+                                content[offs..offs + 8].try_into().unwrap(),
+                            ));
+                            info!("Replication slot {} restart LSN={}", path, restart_lsn);
+                            min_restart_lsn = Lsn::min(min_restart_lsn, restart_lsn);
+                        }
+                        let header = new_tar_header(&path, content.len() as u64)?;
+                        self.ar
+                            .append(&header, &*content)
+                            .await
+                            .context("could not add aux file to basebackup tarball")?;
+                        anyhow::Ok(())
+                    }
+                    .instrument(debug_span!("process aux file", ?path))
+                    .await?;
                 }
-                let header = new_tar_header(&path, content.len() as u64)?;
-                self.ar
-                    .append(&header, &*content)
-                    .await
-                    .context("could not add aux file to basebackup tarball")?;
+
+                debug!("done");
+
+                anyhow::Ok(())
             }
+            .instrument(debug_span!(
+                "process tablespace directory",
+                ?spcnode,
+                ?dbnode
+            ))
+            .await?;
         }
         if min_restart_lsn != Lsn::MAX {
             info!(
@@ -244,19 +284,25 @@ where
                 .await
                 .context("could not add restart.lsn file to basebackup tarball")?;
         }
+        debug!("list twophase files");
         for xid in self
             .timeline
             .list_twophase_files(self.lsn, self.ctx)
             .await?
         {
-            self.add_twophase_file(xid).await?;
+            async {
+                self.add_twophase_file(xid).await?;
+                anyhow::Ok(())
+            }
+            .instrument(debug_span!("process twophase file", ?xid))
+            .await?;
         }
 
         fail_point!("basebackup-before-control-file", |_| {
             bail!("failpoint basebackup-before-control-file")
         });
 
-        // Generate pg_control and bootstrap WAL segment.
+        debug!("Generate pg_control and bootstrap WAL segment.");
         self.add_pgcontrol_file().await?;
         self.ar.finish().await?;
         debug!("all tarred up!");

pageserver/src/client/mgmt_api.rs

@@ -5,20 +5,22 @@ use utils::id::{TenantId, TimelineId};
 
 pub struct Client {
     mgmt_api_endpoint: String,
+    authorization_header: Option<String>,
     client: hyper::Client<HttpConnector, hyper::Body>,
 }
 
 impl Client {
-    pub fn new(mgmt_api_endpoint: String) -> Self {
+    pub fn new(mgmt_api_endpoint: String, jwt: Option<&str>) -> Self {
         Self {
             mgmt_api_endpoint,
+            authorization_header: jwt.map(|jwt| format!("Bearer {jwt}")),
             client: hyper::client::Client::new(),
         }
     }
 
     pub async fn list_tenants(&self) -> anyhow::Result<Vec<pageserver_api::models::TenantInfo>> {
         let uri = Uri::try_from(format!("{}/v1/tenant", self.mgmt_api_endpoint))?;
-        let resp = self.client.get(uri).await?;
+        let resp = self.get(uri).await?;
         if !resp.status().is_success() {
             anyhow::bail!("status error");
         }
@@ -34,7 +36,22 @@ impl Client {
             "{}/v1/tenant/{tenant_id}/timeline",
             self.mgmt_api_endpoint
         ))?;
-        let resp = self.client.get(uri).await?;
+        let resp = self.get(uri).await?;
+        if !resp.status().is_success() {
+            anyhow::bail!("status error");
+        }
+        let body = hyper::body::to_bytes(resp).await?;
+        Ok(serde_json::from_slice(&body)?)
+    }
+
+    pub async fn timeline_info(
+        &self, tenant_id: TenantId, timeline_id: TimelineId,
+    ) -> anyhow::Result<pageserver_api::models::TimelineInfo> {
+        let uri = Uri::try_from(format!(
+            "{}/v1/tenant/{tenant_id}/timeline/{timeline_id}",
+            self.mgmt_api_endpoint
+        ))?;
+        let resp = self.get(uri).await?;
         if !resp.status().is_success() {
             anyhow::bail!("status error");
         }
@@ -51,11 +68,22 @@ impl Client {
             "{}/v1/tenant/{tenant_id}/timeline/{timeline_id}/keyspace?check_serialization_roundtrip=true",
             self.mgmt_api_endpoint
         ))?;
-        let resp = self.client.get(uri).await?;
+        let resp = self.get(uri).await?;
         if !resp.status().is_success() {
             anyhow::bail!("status error");
         }
         let body = hyper::body::to_bytes(resp).await?;
         Ok(serde_json::from_slice(&body).context("deserialize")?)
     }
+
+    async fn get(&self, uri: Uri) -> hyper::Result<hyper::Response<hyper::Body>> {
+        let req = hyper::Request::builder().uri(uri).method("GET");
+        let req = if let Some(value) = &self.authorization_header {
+            req.header("Authorization", value)
+        } else {
+            req
+        };
+        let req = req.body(hyper::Body::default());
+        self.client.request(req.unwrap()).await
+    }
 }