From 2c23ed4873e1326291fb682cc19f399205821717 Mon Sep 17 00:00:00 2001 From: Joonas Koivunen Date: Fri, 15 Dec 2023 21:45:29 +0000 Subject: [PATCH] scan_metadata: rewrite main loop with progress (compile fail) --- s3_scrubber/src/scan_metadata.rs | 216 ++++++++++++++++++++++++++----- 1 file changed, 186 insertions(+), 30 deletions(-) diff --git a/s3_scrubber/src/scan_metadata.rs b/s3_scrubber/src/scan_metadata.rs index fd5ebf8f16..d55b2620bf 100644 --- a/s3_scrubber/src/scan_metadata.rs +++ b/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 { + 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, ) -> anyhow::Result { 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()) }