scan_metadata: rewrite main loop with progress (compile fail)

s3_scrubber/src/scan_metadata.rs

@@ -1,16 +1,21 @@
 use std::collections::{HashMap, HashSet};
+use std::sync::atomic::{AtomicUsize, Ordering};
+use std::sync::Arc;
 
 use crate::checks::{
     branch_cleanup_and_check_errors, list_timeline_blobs, BlobDataParseResult, S3TimelineBlobData,
     TimelineAnalysis,
 };
-use crate::metadata_stream::{stream_tenant_timelines, stream_tenants};
-use crate::{init_remote, BucketConfig, NodeKind, RootTarget, TenantShardTimelineId};
-use aws_sdk_s3::Client;
-use futures_util::{pin_mut, StreamExt, TryStreamExt};
+use crate::metadata_stream::stream_tenants;
+use crate::{init_remote, BucketConfig, NodeKind, TenantShardTimelineId};
+
+use futures_util::StreamExt;
 use histogram::Histogram;
 use pageserver::tenant::IndexPart;
+
+use pageserver_api::shard::TenantShardId;
 use serde::Serialize;
+use tracing::Instrument;
 
 #[derive(Serialize)]
 pub struct MetadataSummary {
@@ -179,12 +184,19 @@ Timeline layer count: {6}
     }
 }
 
+#[derive(Debug)]
+enum Either<A, B> {
+    Left(A),
+    Right(B),
+}
+
 /// Scan the pageserver metadata in an S3 bucket, reporting errors and statistics.
 pub async fn scan_metadata(
     bucket_config: BucketConfig,
     tenant_ids: Vec<TenantShardId>,
 ) -> anyhow::Result<MetadataSummary> {
     let (s3_client, target) = init_remote(bucket_config, NodeKind::Pageserver)?;
+    let target = Arc::new(target);
 
     let tenants = if tenant_ids.is_empty() {
         futures::future::Either::Left(stream_tenants(&s3_client, &target))
@@ -192,37 +204,181 @@ pub async fn scan_metadata(
         futures::future::Either::Right(futures::stream::iter(tenant_ids.into_iter().map(Ok)))
     };
 
-    // How many tenants to process in parallel.  We need to be mindful of pageservers
-    // accessing the same per tenant prefixes, so use a lower setting than pageservers.
-    const CONCURRENCY: usize = 32;
+    let tenants = tenants.fuse();
 
-    // Generate a stream of TenantTimelineId
-    let timelines = tenants.map_ok(|t| stream_tenant_timelines(&s3_client, &target, t));
-    let timelines = timelines.try_buffer_unordered(CONCURRENCY);
-    let timelines = timelines.try_flatten();
+    let mut tenants = std::pin::pin!(tenants);
 
-    // Generate a stream of S3TimelineBlobData
-    async fn report_on_timeline(
-        s3_client: &Client,
-        target: &RootTarget,
-        ttid: TenantShardTimelineId,
-    ) -> anyhow::Result<(TenantShardTimelineId, S3TimelineBlobData)> {
-        let data = list_timeline_blobs(s3_client, ttid, target).await?;
-        Ok((ttid, data))
-    }
-    let timelines = timelines.map_ok(|ttid| report_on_timeline(&s3_client, &target, ttid));
-    let timelines = timelines.try_buffer_unordered(CONCURRENCY);
+    let mut js = tokio::task::JoinSet::new();
+    let mut consumed_all = false;
 
-    let mut summary = MetadataSummary::new();
-    pin_mut!(timelines);
-    while let Some(i) = timelines.next().await {
-        let (ttid, data) = i?;
-        summary.update_data(&data);
+    let summary = MetadataSummary::new();
 
-        let analysis = branch_cleanup_and_check_errors(&ttid, &target, None, None, Some(data));
+    // have timeline and timeline blob listings fight over the same semaphore
+    let timeline_listings = Arc::new(tokio::sync::Semaphore::new(50));
+    let blob_listings = timeline_listings.clone();
 
-        summary.update_analysis(&ttid, &analysis);
+    let spawned_tenants = AtomicUsize::new(0);
+    let spawned_timelines = Arc::new(AtomicUsize::new(0));
+    let completed_tenants = AtomicUsize::new(0);
+    let completed_timelines = AtomicUsize::new(0);
+
+    let s3_client = s3_client.clone();
+    let target = target.clone();
+
+    let summary = std::sync::Mutex::new(summary);
+
+    let scan_tenants = async {
+        let timeline_listings = timeline_listings;
+        let blob_listings = blob_listings;
+
+        // used to control whether to receive more tenants
+        let mut more_tenants = true;
+
+        loop {
+            let next_start = tokio::select! {
+                next_tenant = tenants.next(), if !consumed_all && more_tenants => {
+                    match next_tenant {
+                        Some(Ok(tenant_id)) => Either::Left(tenant_id),
+                        Some(Err(e)) => {
+                            consumed_all = true;
+                            tracing::error!("tenant streaming failed with: {e:?}");
+                            continue;
+                        }
+                        None => {
+                            consumed_all = true;
+                            continue;
+                        }
+                    }
+                },
+
+                next = js.join_next(), if !js.is_empty() => {
+                    more_tenants = js.len() < 10;
+
+                    match next.unwrap() {
+                        Ok(Either::Left((tenant_id, timelines))) => {
+                            completed_tenants.fetch_add(1, Ordering::Relaxed);
+
+                            Either::Right((tenant_id, timelines))
+                        }
+                        Ok(Either::Right(Some((ttid, data)))) => {
+                            completed_timelines.fetch_add(1, Ordering::Relaxed);
+
+                            let ttid: TenantShardTimelineId = ttid;
+                            {
+                                let _e = tracing::info_span!("analysis", tenant_shard_id=%ttid.tenant_shard_id, timeline_id=%ttid.timeline_id).entered();
+                                let summary = &mut summary.lock().unwrap();
+                                summary.update_data(&data);
+                                let analysis = branch_cleanup_and_check_errors(&ttid, &target, None, None, Some(data));
+                                summary.update_analysis(&ttid, &analysis);
+                            }
+                            continue;
+                        }
+                        Ok(Either::Right(None)) => {
+                            completed_timelines.fetch_add(1, Ordering::Relaxed);
+                            continue;
+                        }
+                        Err(je) if je.is_cancelled() => unreachable!("not used"),
+                        Err(je) if je.is_panic() => {
+                            continue;
+                        },
+                        Err(je) => {
+                            tracing::error!("unknown join error: {je:?}");
+                            continue;
+                        }
+                    }
+                },
+
+                else => break,
+            };
+
+            let s3_client = s3_client.clone();
+            let target = target.clone();
+            let timeline_listings = timeline_listings.clone();
+            let blob_listings = blob_listings.clone();
+
+            match next_start {
+                Either::Left(tenant_shard_id) => {
+                    let span = tracing::info_span!("get_timelines", tenant_id=%tenant_shard_id.tenant_id, shard_id=%tenant_shard_id.shard_slug());
+                    js.spawn(
+                        async move {
+                            let _permit = timeline_listings.acquire().await;
+                            let timelines = crate::metadata_stream::get_tenant_timelines(
+                                &s3_client,
+                                &target,
+                                tenant_shard_id,
+                            )
+                            .await;
+
+                            Either::Left((tenant_shard_id, timelines))
+                        }
+                        .instrument(span),
+                    );
+
+                    more_tenants = js.len() < 1000;
+
+                    spawned_tenants.fetch_add(1, Ordering::Relaxed);
+                }
+                Either::Right((tenant_shard_id, timelines)) => {
+                    for timeline_id in timelines {
+                        let timeline_id = match timeline_id {
+                            Ok(timeline_id) => timeline_id,
+                            Err(e) => {
+                                tracing::error!("failed to fetch a timeline: {e:?}");
+                                continue;
+                            }
+                        };
+                        let s3_client = s3_client.clone();
+                        let target = target.clone();
+                        let blob_listings = blob_listings.clone();
+                        let span = tracing::info_span!("list_timelines_blobs", tenant_id=%tenant_shard_id.tenant_id, shard_id=%tenant_shard_id.shard_slug(), %timeline_id);
+                        js.spawn(
+                            async move {
+                                let _permit = blob_listings.acquire().await;
+                                let ttid = TenantShardTimelineId::new(tenant_shard_id, timeline_id);
+                                match list_timeline_blobs(&s3_client, ttid, &target).await {
+                                    Ok(data) => Either::Right(Some((ttid, data))),
+                                    Err(e) => {
+                                        tracing::error!("listing failed {e:?}");
+                                        Either::Right(None)
+                                    }
+                                }
+                            }
+                            .instrument(span),
+                        );
+
+                        spawned_timelines.fetch_add(1, Ordering::Relaxed);
+                        tokio::task::yield_now().await;
+                    }
+                }
+            }
+        }
+    };
+
+    let started_at = std::time::Instant::now();
+
+    {
+        let mut scan_tenants = std::pin::pin!(scan_tenants);
+
+        loop {
+            let res =
+                tokio::time::timeout(std::time::Duration::from_secs(1), &mut scan_tenants).await;
+
+            let spawned_tenants = spawned_tenants.load(Ordering::Relaxed);
+            let completed_tenants = completed_tenants.load(Ordering::Relaxed);
+            let spawned_timelines = spawned_timelines.load(Ordering::Relaxed);
+            let completed_timelines = completed_timelines.load(Ordering::Relaxed);
+
+            match res {
+                Ok(()) => {
+                    tracing::info!("progress tenants: {completed_tenants} / {spawned_tenants}, timelines: {completed_timelines} / {spawned_timelines} after {:?}", started_at.elapsed());
+                    break;
+                }
+                Err(_timeout) => {
+                    tracing::info!("progress tenants: {completed_tenants} / {spawned_tenants}, timelines: {completed_timelines} / {spawned_timelines}");
+                }
+            }
+        }
     }
 
-    Ok(summary)
+    Ok(summary.into_inner().unwrap())
 }