System Initiative

use std::{ collections::HashSet, future::IntoFuture as _, }; use dal::{ ChangeSetStatus, ServicesContext, TransactionsError, WorkspaceSnapshotAddress, }; use si_data_pg::{ PgError, PgPoolError, }; use si_layer_cache::LayerDbError; use strum::IntoEnumIterator; use telemetry::prelude::*; use thiserror::Error; use tokio_util::{ sync::CancellationToken, task::TaskTracker, }; use crate::{ Config, init, }; #[remain::sorted] #[derive(Debug, Error)] pub enum GarbageCollectorError { #[error("error while initializing: {0}")] Init(#[from] init::InitError), #[error("Layer DB error: {0}")] LayerDb(#[from] LayerDbError), #[error("Pg error: {0}")] Pg(#[from] PgError), #[error("Pg Pool error: {0}")] PgPool(#[from] PgPoolError), #[error("Transactions error: {0}")] Transactions(#[from] TransactionsError), #[error("Unable to query workspace snapshots")] UnableToQuerySnapshots, } type Result<T> = std::result::Result<T, GarbageCollectorError>; pub struct SnapshotGarbageCollector { services_context: ServicesContext, } impl SnapshotGarbageCollector { #[instrument(name = "sdf.snapshot_garbage_collector.new", level = "info", skip_all)] pub async fn new( config: Config, task_tracker: &TaskTracker, task_token: CancellationToken, ) -> Result<Self> { let (services_context, layer_db_graceful_shutdown) = init::services_context_from_config(&config, task_token).await?; task_tracker.spawn(layer_db_graceful_shutdown.into_future()); Ok(Self { services_context }) } #[instrument( name = "sdf.snapshot_garbage_collector.garbage_collect_snapshots", level = "info", skip_all )] pub async fn garbage_collect_snapshots(self) -> Result<()> { let span = current_span_for_instrument_at!("info"); let dal_context = self.services_context.clone().into_builder(true); let ctx = dal_context .build_default(None) .await .map_err(|err| span.record_err(err))?; let ctx = &ctx; let mut open_change_set_snapshot_ids = HashSet::new(); let mut all_snapshot_ids = HashSet::new(); // Gather the WorkspaceSnapshotAddress of all open change sets. let open_statuses: Vec<String> = ChangeSetStatus::iter() .filter_map(|status| { if status.is_active() { Some(status.to_string()) } else { None } }) .collect(); info!( "Change set status(es) to consider 'active': {:?}", &open_statuses ); let change_set_snapshot_rows = ctx.txns().await?.pg().query( "SELECT workspace_snapshot_address AS snapshot_id FROM change_set_pointers WHERE status = ANY($1::text[]) GROUP BY workspace_snapshot_address", &[&open_statuses], ).await?; for row in change_set_snapshot_rows { let snapshot_id: WorkspaceSnapshotAddress = row.try_get("snapshot_id")?; open_change_set_snapshot_ids.insert(snapshot_id); } info!( "Found {} distinct snapshot address(es) for open change sets.", open_change_set_snapshot_ids.len() ); // Gather the WorkspaceSnapshotAddress of all existing snapshots that are // at least an hour old. // // By only considering the snapshots that are at least an hour old, we avoid // race conditions where a change set is created, or modified between when we // queried the change_set_pointers table and the workspace_snapshots table. // We can't rely on transactional integrity to avoid race conditions as the // tables are in completely separate databases. let snapshot_id_rows = ctx .layer_db() .workspace_snapshot() .cache .pg() .query( "SELECT key AS snapshot_id FROM workspace_snapshots WHERE created_at < NOW() - '1 hour'::interval GROUP BY key", &[], ) .await? .ok_or_else(|| GarbageCollectorError::UnableToQuerySnapshots)?; for row in snapshot_id_rows { let snapshot_id: WorkspaceSnapshotAddress = row.try_get("snapshot_id")?; all_snapshot_ids.insert(snapshot_id); } info!( "Found {} distinct snapshot address(es) older than cutoff.", all_snapshot_ids.len() ); // Any WorkspaceSnapshotAddress not in both open_change_set_snapshot_ids // and all_snapshot_ids is for a closed/applied/abandoned change set and // can be deleted, as we do not allow re-opening change sets. let snapshot_ids_to_delete: HashSet<_> = all_snapshot_ids .difference(&open_change_set_snapshot_ids) .collect(); info!( "Found {} snapshot address(es) only used by inactive change sets.", snapshot_ids_to_delete.len() ); let mut counter = 0; for key in snapshot_ids_to_delete.iter().take(10_000) { ctx.layer_db() .workspace_snapshot() .cache .pg() .delete(&key.to_string()) .await?; counter += 1; if counter % 100 == 0 { info!("Deleted {} snapshot addresses.", counter); } } info!("Deleted {} snapshot address(es).", counter); ctx.commit().await?; span.record_ok(); Ok(()) } }