LCOV - code coverage report
Current view: top level - pageserver/src/tenant/secondary - heatmap_uploader.rs (source / functions) Coverage Total Hit
Test: 5fe7fa8d483b39476409aee736d6d5e32728bfac.info Lines: 0.0 % 277 0
Test Date: 2025-03-12 16:10:49 Functions: 0.0 % 26 0

            Line data    Source code
       1              : use std::collections::HashMap;
       2              : use std::pin::Pin;
       3              : use std::sync::{Arc, Weak};
       4              : use std::time::{Duration, Instant};
       5              : 
       6              : use futures::Future;
       7              : use pageserver_api::shard::TenantShardId;
       8              : use remote_storage::{GenericRemoteStorage, TimeoutOrCancel};
       9              : use tokio_util::sync::CancellationToken;
      10              : use tracing::{Instrument, info_span, instrument};
      11              : use utils::backoff;
      12              : use utils::completion::Barrier;
      13              : use utils::crashsafe::path_with_suffix_extension;
      14              : use utils::yielding_loop::yielding_loop;
      15              : 
      16              : use super::heatmap::HeatMapTenant;
      17              : use super::scheduler::{
      18              :     self, JobGenerator, RunningJob, SchedulingResult, TenantBackgroundJobs, period_jitter,
      19              :     period_warmup,
      20              : };
      21              : use super::{CommandRequest, SecondaryTenantError, UploadCommand};
      22              : use crate::TEMP_FILE_SUFFIX;
      23              : use crate::metrics::SECONDARY_MODE;
      24              : use crate::tenant::Tenant;
      25              : use crate::tenant::config::AttachmentMode;
      26              : use crate::tenant::mgr::{GetTenantError, TenantManager};
      27              : use crate::tenant::remote_timeline_client::remote_heatmap_path;
      28              : use crate::tenant::span::debug_assert_current_span_has_tenant_id;
      29              : use crate::tenant::tasks::{BackgroundLoopKind, warn_when_period_overrun};
      30              : use crate::virtual_file::VirtualFile;
      31              : 
      32            0 : pub(super) async fn heatmap_uploader_task(
      33            0 :     tenant_manager: Arc<TenantManager>,
      34            0 :     remote_storage: GenericRemoteStorage,
      35            0 :     command_queue: tokio::sync::mpsc::Receiver<CommandRequest<UploadCommand>>,
      36            0 :     background_jobs_can_start: Barrier,
      37            0 :     cancel: CancellationToken,
      38            0 : ) {
      39            0 :     let concurrency = tenant_manager.get_conf().heatmap_upload_concurrency;
      40            0 : 
      41            0 :     let generator = HeatmapUploader {
      42            0 :         tenant_manager,
      43            0 :         remote_storage,
      44            0 :         cancel: cancel.clone(),
      45            0 :         tenants: HashMap::new(),
      46            0 :     };
      47            0 :     let mut scheduler = Scheduler::new(generator, concurrency);
      48            0 : 
      49            0 :     scheduler
      50            0 :         .run(command_queue, background_jobs_can_start, cancel)
      51            0 :         .instrument(info_span!("heatmap_upload_scheduler"))
      52            0 :         .await
      53            0 : }
      54              : 
      55              : /// This type is owned by a single task ([`heatmap_uploader_task`]) which runs an event
      56              : /// handling loop and mutates it as needed: there are no locks here, because that event loop
      57              : /// can hold &mut references to this type throughout.
      58              : struct HeatmapUploader {
      59              :     tenant_manager: Arc<TenantManager>,
      60              :     remote_storage: GenericRemoteStorage,
      61              :     cancel: CancellationToken,
      62              : 
      63              :     tenants: HashMap<TenantShardId, UploaderTenantState>,
      64              : }
      65              : 
      66              : struct WriteInProgress {
      67              :     barrier: Barrier,
      68              : }
      69              : 
      70              : impl RunningJob for WriteInProgress {
      71            0 :     fn get_barrier(&self) -> Barrier {
      72            0 :         self.barrier.clone()
      73            0 :     }
      74              : }
      75              : 
      76              : struct UploadPending {
      77              :     tenant: Arc<Tenant>,
      78              :     last_upload: Option<LastUploadState>,
      79              :     target_time: Option<Instant>,
      80              :     period: Option<Duration>,
      81              : }
      82              : 
      83              : impl scheduler::PendingJob for UploadPending {
      84            0 :     fn get_tenant_shard_id(&self) -> &TenantShardId {
      85            0 :         self.tenant.get_tenant_shard_id()
      86            0 :     }
      87              : }
      88              : 
      89              : struct WriteComplete {
      90              :     tenant_shard_id: TenantShardId,
      91              :     completed_at: Instant,
      92              :     uploaded: Option<LastUploadState>,
      93              :     next_upload: Option<Instant>,
      94              : }
      95              : 
      96              : impl scheduler::Completion for WriteComplete {
      97            0 :     fn get_tenant_shard_id(&self) -> &TenantShardId {
      98            0 :         &self.tenant_shard_id
      99            0 :     }
     100              : }
     101              : 
     102              : /// The heatmap uploader keeps a little bit of per-tenant state, mainly to remember
     103              : /// when we last did a write.  We only populate this after doing at least one
     104              : /// write for a tenant -- this avoids holding state for tenants that have
     105              : /// uploads disabled.
     106              : struct UploaderTenantState {
     107              :     // This Weak only exists to enable culling idle instances of this type
     108              :     // when the Tenant has been deallocated.
     109              :     tenant: Weak<Tenant>,
     110              : 
     111              :     /// Digest of the serialized heatmap that we last successfully uploaded
     112              :     last_upload_state: Option<LastUploadState>,
     113              : 
     114              :     /// When the last upload attempt completed (may have been successful or failed)
     115              :     last_upload: Option<Instant>,
     116              : 
     117              :     /// When should we next do an upload?  None means never.
     118              :     next_upload: Option<Instant>,
     119              : }
     120              : 
     121              : type Scheduler = TenantBackgroundJobs<
     122              :     HeatmapUploader,
     123              :     UploadPending,
     124              :     WriteInProgress,
     125              :     WriteComplete,
     126              :     UploadCommand,
     127              : >;
     128              : 
     129              : impl JobGenerator<UploadPending, WriteInProgress, WriteComplete, UploadCommand>
     130              :     for HeatmapUploader
     131              : {
     132            0 :     async fn schedule(&mut self) -> SchedulingResult<UploadPending> {
     133            0 :         // Cull any entries in self.tenants whose Arc<Tenant> is gone
     134            0 :         self.tenants
     135            0 :             .retain(|_k, v| v.tenant.upgrade().is_some() && v.next_upload.is_some());
     136            0 : 
     137            0 :         let now = Instant::now();
     138            0 : 
     139            0 :         let mut result = SchedulingResult {
     140            0 :             jobs: Vec::new(),
     141            0 :             want_interval: None,
     142            0 :         };
     143            0 : 
     144            0 :         let tenants = self.tenant_manager.get_attached_active_tenant_shards();
     145            0 : 
     146            0 :         yielding_loop(1000, &self.cancel, tenants.into_iter(), |tenant| {
     147            0 :             let period = match tenant.get_heatmap_period() {
     148              :                 None => {
     149              :                     // Heatmaps are disabled for this tenant
     150            0 :                     return;
     151              :                 }
     152            0 :                 Some(period) => {
     153            0 :                     // If any tenant has asked for uploads more frequent than our scheduling interval,
     154            0 :                     // reduce it to match so that we can keep up.  This is mainly useful in testing, where
     155            0 :                     // we may set rather short intervals.
     156            0 :                     result.want_interval = match result.want_interval {
     157            0 :                         None => Some(period),
     158            0 :                         Some(existing) => Some(std::cmp::min(period, existing)),
     159              :                     };
     160              : 
     161            0 :                     period
     162            0 :                 }
     163            0 :             };
     164            0 : 
     165            0 :             // Stale attachments do not upload anything: if we are in this state, there is probably some
     166            0 :             // other attachment in mode Single or Multi running on another pageserver, and we don't
     167            0 :             // want to thrash and overwrite their heatmap uploads.
     168            0 :             if tenant.get_attach_mode() == AttachmentMode::Stale {
     169            0 :                 return;
     170            0 :             }
     171            0 : 
     172            0 :             // Create an entry in self.tenants if one doesn't already exist: this will later be updated
     173            0 :             // with the completion time in on_completion.
     174            0 :             let state = self
     175            0 :                 .tenants
     176            0 :                 .entry(*tenant.get_tenant_shard_id())
     177            0 :                 .or_insert_with(|| UploaderTenantState {
     178            0 :                     tenant: Arc::downgrade(&tenant),
     179            0 :                     last_upload: None,
     180            0 :                     next_upload: Some(now.checked_add(period_warmup(period)).unwrap_or(now)),
     181            0 :                     last_upload_state: None,
     182            0 :                 });
     183            0 : 
     184            0 :             // Decline to do the upload if insufficient time has passed
     185            0 :             if state.next_upload.map(|nu| nu > now).unwrap_or(false) {
     186            0 :                 return;
     187            0 :             }
     188            0 : 
     189            0 :             let last_upload = state.last_upload_state.clone();
     190            0 :             result.jobs.push(UploadPending {
     191            0 :                 tenant,
     192            0 :                 last_upload,
     193            0 :                 target_time: state.next_upload,
     194            0 :                 period: Some(period),
     195            0 :             });
     196            0 :         })
     197            0 :         .await
     198            0 :         .ok();
     199            0 : 
     200            0 :         result
     201            0 :     }
     202              : 
     203            0 :     fn spawn(
     204            0 :         &mut self,
     205            0 :         job: UploadPending,
     206            0 :     ) -> (
     207            0 :         WriteInProgress,
     208            0 :         Pin<Box<dyn Future<Output = WriteComplete> + Send>>,
     209            0 :     ) {
     210            0 :         let UploadPending {
     211            0 :             tenant,
     212            0 :             last_upload,
     213            0 :             target_time,
     214            0 :             period,
     215            0 :         } = job;
     216            0 : 
     217            0 :         let remote_storage = self.remote_storage.clone();
     218            0 :         let (completion, barrier) = utils::completion::channel();
     219            0 :         let tenant_shard_id = *tenant.get_tenant_shard_id();
     220            0 :         (WriteInProgress { barrier }, Box::pin(async move {
     221            0 :             // Guard for the barrier in [`WriteInProgress`]
     222            0 :             let _completion = completion;
     223            0 : 
     224            0 :             let started_at = Instant::now();
     225            0 :             let uploaded = match upload_tenant_heatmap(remote_storage, &tenant, last_upload.clone()).await {
     226            0 :                 Ok(UploadHeatmapOutcome::Uploaded(uploaded)) => {
     227            0 :                     let duration = Instant::now().duration_since(started_at);
     228            0 :                     SECONDARY_MODE
     229            0 :                         .upload_heatmap_duration
     230            0 :                         .observe(duration.as_secs_f64());
     231            0 :                     SECONDARY_MODE.upload_heatmap.inc();
     232            0 :                     Some(uploaded)
     233              :                 }
     234            0 :                 Ok(UploadHeatmapOutcome::NoChange | UploadHeatmapOutcome::Skipped) => last_upload,
     235            0 :                 Err(UploadHeatmapError::Upload(e)) => {
     236            0 :                     tracing::warn!(
     237            0 :                         "Failed to upload heatmap for tenant {}: {e:#}",
     238            0 :                         tenant.get_tenant_shard_id(),
     239              :                     );
     240            0 :                     let duration = Instant::now().duration_since(started_at);
     241            0 :                     SECONDARY_MODE
     242            0 :                         .upload_heatmap_duration
     243            0 :                         .observe(duration.as_secs_f64());
     244            0 :                     SECONDARY_MODE.upload_heatmap_errors.inc();
     245            0 :                     last_upload
     246              :                 }
     247              :                 Err(UploadHeatmapError::Cancelled) => {
     248            0 :                     tracing::info!("Cancelled heatmap upload, shutting down");
     249            0 :                     last_upload
     250              :                 }
     251              :             };
     252              : 
     253            0 :             let now = Instant::now();
     254              : 
     255              :             // If the job had a target execution time, we may check our final execution
     256              :             // time against that for observability purposes.
     257            0 :             if let (Some(target_time), Some(period)) = (target_time, period) {
     258            0 :                 // Elapsed time includes any scheduling lag as well as the execution of the job
     259            0 :                 let elapsed = now.duration_since(target_time);
     260            0 : 
     261            0 :                 warn_when_period_overrun(elapsed, period, BackgroundLoopKind::HeatmapUpload);
     262            0 :             }
     263              : 
     264            0 :             let next_upload = tenant
     265            0 :                 .get_heatmap_period()
     266            0 :                 .and_then(|period| now.checked_add(period_jitter(period, 5)));
     267            0 : 
     268            0 :             WriteComplete {
     269            0 :                     tenant_shard_id: *tenant.get_tenant_shard_id(),
     270            0 :                     completed_at: now,
     271            0 :                     uploaded,
     272            0 :                     next_upload,
     273            0 :                 }
     274            0 :         }.instrument(info_span!(parent: None, "heatmap_upload", tenant_id=%tenant_shard_id.tenant_id, shard_id=%tenant_shard_id.shard_slug()))))
     275            0 :     }
     276              : 
     277            0 :     fn on_command(
     278            0 :         &mut self,
     279            0 :         command: UploadCommand,
     280            0 :     ) -> Result<UploadPending, SecondaryTenantError> {
     281            0 :         let tenant_shard_id = command.get_tenant_shard_id();
     282            0 : 
     283            0 :         tracing::info!(
     284            0 :             tenant_id=%tenant_shard_id.tenant_id, shard_id=%tenant_shard_id.shard_slug(),
     285            0 :             "Starting heatmap write on command");
     286            0 :         let tenant = self
     287            0 :             .tenant_manager
     288            0 :             .get_attached_tenant_shard(*tenant_shard_id)?;
     289            0 :         if !tenant.is_active() {
     290            0 :             return Err(GetTenantError::NotActive(*tenant_shard_id).into());
     291            0 :         }
     292            0 : 
     293            0 :         Ok(UploadPending {
     294            0 :             // Ignore our state for last digest: this forces an upload even if nothing has changed
     295            0 :             last_upload: None,
     296            0 :             tenant,
     297            0 :             target_time: None,
     298            0 :             period: None,
     299            0 :         })
     300            0 :     }
     301              : 
     302              :     #[instrument(skip_all, fields(tenant_id=%completion.tenant_shard_id.tenant_id, shard_id=%completion.tenant_shard_id.shard_slug()))]
     303              :     fn on_completion(&mut self, completion: WriteComplete) {
     304              :         tracing::debug!("Heatmap upload completed");
     305              :         let WriteComplete {
     306              :             tenant_shard_id,
     307              :             completed_at,
     308              :             uploaded,
     309              :             next_upload,
     310              :         } = completion;
     311              :         use std::collections::hash_map::Entry;
     312              :         match self.tenants.entry(tenant_shard_id) {
     313              :             Entry::Vacant(_) => {
     314              :                 // Tenant state was dropped, nothing to update.
     315              :             }
     316              :             Entry::Occupied(mut entry) => {
     317              :                 entry.get_mut().last_upload = Some(completed_at);
     318              :                 entry.get_mut().last_upload_state = uploaded;
     319              :                 entry.get_mut().next_upload = next_upload
     320              :             }
     321              :         }
     322              :     }
     323              : }
     324              : 
     325              : enum UploadHeatmapOutcome {
     326              :     /// We successfully wrote to remote storage, with this digest.
     327              :     Uploaded(LastUploadState),
     328              :     /// We did not upload because the heatmap digest was unchanged since the last upload
     329              :     NoChange,
     330              :     /// We skipped the upload for some reason, such as tenant/timeline not ready
     331              :     Skipped,
     332              : }
     333              : 
     334              : #[derive(thiserror::Error, Debug)]
     335              : enum UploadHeatmapError {
     336              :     #[error("Cancelled")]
     337              :     Cancelled,
     338              : 
     339              :     #[error(transparent)]
     340              :     Upload(#[from] anyhow::Error),
     341              : }
     342              : 
     343              : /// Digests describing the heatmap we most recently uploaded successfully.
     344              : ///
     345              : /// md5 is generally a bad hash.  We use it because it's convenient for interop with AWS S3's ETag,
     346              : /// which is also an md5sum.
     347              : #[derive(Clone)]
     348              : struct LastUploadState {
     349              :     // Digest of json-encoded HeatMapTenant
     350              :     uploaded_digest: md5::Digest,
     351              : 
     352              :     // Digest without atimes set.
     353              :     layers_only_digest: md5::Digest,
     354              : }
     355              : 
     356              : /// The inner upload operation.  This will skip if `last_digest` is Some and matches the digest
     357              : /// of the object we would have uploaded.
     358            0 : async fn upload_tenant_heatmap(
     359            0 :     remote_storage: GenericRemoteStorage,
     360            0 :     tenant: &Arc<Tenant>,
     361            0 :     last_upload: Option<LastUploadState>,
     362            0 : ) -> Result<UploadHeatmapOutcome, UploadHeatmapError> {
     363            0 :     debug_assert_current_span_has_tenant_id();
     364            0 : 
     365            0 :     let generation = tenant.get_generation();
     366            0 :     debug_assert!(!generation.is_none());
     367            0 :     if generation.is_none() {
     368              :         // We do not expect this: None generations should only appear in historic layer metadata, not in running Tenants
     369            0 :         tracing::warn!("Skipping heatmap upload for tenant with generation==None");
     370            0 :         return Ok(UploadHeatmapOutcome::Skipped);
     371            0 :     }
     372            0 : 
     373            0 :     let mut heatmap = HeatMapTenant {
     374            0 :         timelines: Vec::new(),
     375            0 :         generation,
     376            0 :         upload_period_ms: tenant.get_heatmap_period().map(|p| p.as_millis()),
     377            0 :     };
     378            0 :     let timelines = tenant.timelines.lock().unwrap().clone();
     379              : 
     380              :     // Ensure that Tenant::shutdown waits for any upload in flight: this is needed because otherwise
     381              :     // when we delete a tenant, we might race with an upload in flight and end up leaving a heatmap behind
     382              :     // in remote storage.
     383            0 :     let Ok(_guard) = tenant.gate.enter() else {
     384            0 :         tracing::info!("Skipping heatmap upload for tenant which is shutting down");
     385            0 :         return Err(UploadHeatmapError::Cancelled);
     386              :     };
     387              : 
     388            0 :     for (timeline_id, timeline) in timelines {
     389            0 :         let heatmap_timeline = timeline.generate_heatmap().await;
     390            0 :         match heatmap_timeline {
     391              :             None => {
     392            0 :                 tracing::debug!(
     393            0 :                     "Skipping heatmap upload because timeline {timeline_id} is not ready"
     394              :                 );
     395            0 :                 return Ok(UploadHeatmapOutcome::Skipped);
     396              :             }
     397            0 :             Some(heatmap_timeline) => {
     398            0 :                 heatmap.timelines.push(heatmap_timeline);
     399            0 :             }
     400              :         }
     401              :     }
     402              : 
     403              :     // Serialize the heatmap
     404            0 :     let bytes = serde_json::to_vec(&heatmap).map_err(|e| anyhow::anyhow!(e))?;
     405              : 
     406              :     // Drop out early if nothing changed since our last upload
     407            0 :     let digest = md5::compute(&bytes);
     408            0 :     if Some(&digest) == last_upload.as_ref().map(|d| &d.uploaded_digest) {
     409            0 :         return Ok(UploadHeatmapOutcome::NoChange);
     410            0 :     }
     411            0 : 
     412            0 :     // Calculate a digest that omits atimes, so that we can distinguish actual changes in
     413            0 :     // layers from changes only in atimes.
     414            0 :     let heatmap_size_bytes = heatmap.get_stats().bytes;
     415            0 :     let layers_only_bytes =
     416            0 :         serde_json::to_vec(&heatmap.strip_atimes()).map_err(|e| anyhow::anyhow!(e))?;
     417            0 :     let layers_only_digest = md5::compute(&layers_only_bytes);
     418            0 :     if heatmap_size_bytes < tenant.get_checkpoint_distance() {
     419              :         // For small tenants, skip upload if only atimes changed. This avoids doing frequent
     420              :         // uploads from long-idle tenants whose atimes are just incremented by periodic
     421              :         // size calculations.
     422            0 :         if Some(&layers_only_digest) == last_upload.as_ref().map(|d| &d.layers_only_digest) {
     423            0 :             return Ok(UploadHeatmapOutcome::NoChange);
     424            0 :         }
     425            0 :     }
     426              : 
     427            0 :     let bytes = bytes::Bytes::from(bytes);
     428            0 :     let size = bytes.len();
     429            0 : 
     430            0 :     let path = remote_heatmap_path(tenant.get_tenant_shard_id());
     431            0 : 
     432            0 :     let cancel = &tenant.cancel;
     433            0 : 
     434            0 :     tracing::debug!("Uploading {size} byte heatmap to {path}");
     435            0 :     if let Err(e) = backoff::retry(
     436            0 :         || async {
     437            0 :             let bytes = futures::stream::once(futures::future::ready(Ok(bytes.clone())));
     438            0 :             remote_storage
     439            0 :                 .upload_storage_object(bytes, size, &path, cancel)
     440            0 :                 .await
     441            0 :         },
     442            0 :         TimeoutOrCancel::caused_by_cancel,
     443            0 :         3,
     444            0 :         u32::MAX,
     445            0 :         "Uploading heatmap",
     446            0 :         cancel,
     447            0 :     )
     448            0 :     .await
     449            0 :     .ok_or_else(|| anyhow::anyhow!("Shutting down"))
     450            0 :     .and_then(|x| x)
     451              :     {
     452            0 :         if cancel.is_cancelled() {
     453            0 :             return Err(UploadHeatmapError::Cancelled);
     454              :         } else {
     455            0 :             return Err(e.into());
     456              :         }
     457            0 :     }
     458            0 : 
     459            0 :     // After a successful upload persist the fresh heatmap to disk.
     460            0 :     // When restarting, the tenant will read the heatmap from disk
     461            0 :     // and additively generate a new heatmap (see [`Timeline::generate_heatmap`]).
     462            0 :     // If the heatmap is stale, the additive generation can lead to keeping previously
     463            0 :     // evicted timelines on the secondarie's disk.
     464            0 :     let tenant_shard_id = tenant.get_tenant_shard_id();
     465            0 :     let heatmap_path = tenant.conf.tenant_heatmap_path(tenant_shard_id);
     466            0 :     let temp_path = path_with_suffix_extension(&heatmap_path, TEMP_FILE_SUFFIX);
     467            0 :     if let Err(err) = VirtualFile::crashsafe_overwrite(heatmap_path, temp_path, bytes).await {
     468            0 :         tracing::warn!("Non fatal IO error writing to disk after heatmap upload: {err}");
     469            0 :     }
     470              : 
     471            0 :     tracing::info!("Successfully uploaded {size} byte heatmap to {path}");
     472              : 
     473            0 :     Ok(UploadHeatmapOutcome::Uploaded(LastUploadState {
     474            0 :         uploaded_digest: digest,
     475            0 :         layers_only_digest,
     476            0 :     }))
     477            0 : }
        

Generated by: LCOV version 2.1-beta