test artifact pruning in the integration test

Author: iliana

nexus/db-queries/src/db/datastore/target_release.rs

@@ -208,6 +208,7 @@ impl DataStore {
 }
 
 /// Represents information fetched about recent target releases
+#[derive(Debug)]
 pub struct RecentTargetReleases {
     /// distinct target releases found
     pub releases: BTreeSet<TufRepoUuid>,

nexus/tests/integration_tests/updates.rs

@@ -426,7 +426,7 @@ async fn test_repo_upload() -> Result<()> {
 
     // Upload a new repository with a different system version but no other
     // changes. This should be accepted.
-    {
+    let initial_installinator_doc = {
         let tweaks = &[ManifestTweak::SystemVersion("2.0.0".parse().unwrap())];
         let response = trust_root
             .assemble_repo(&logctx.log, tweaks)
@@ -507,7 +507,9 @@ async fn test_repo_upload() -> Result<()> {
             description, get_description,
             "initial description matches fetched description"
         );
-    }
+
+        installinator_doc_1
+    };
     // The installinator document changed, so the generation number is bumped to
     // 3.
     assert_eq!(
@@ -523,6 +525,75 @@ async fn test_repo_upload() -> Result<()> {
     assert_eq!(status.last_run_counters.put_ok, 3);
     assert_eq!(status.last_run_counters.copy_ok, 1);
     assert_eq!(status.local_repos, 1);
+    // Run the replication background task again; the local repos should be
+    // dropped.
+    let status =
+        run_tuf_artifact_replication_step(&cptestctx.lockstep_client).await;
+    eprintln!("{status:?}");
+    assert_eq!(status.last_run_counters.put_config_ok, 4);
+    assert_eq!(status.last_run_counters.list_ok, 4);
+    assert_eq!(status.last_run_counters.sum(), 8);
+    assert_eq!(status.local_repos, 0);
+
+    // Verify the initial installinator document is present on all sled-agents.
+    let installinator_doc_hash = initial_installinator_doc.hash.to_string();
+    for sled_agent in &cptestctx.sled_agents {
+        for dir in sled_agent.sled_agent().artifact_store().storage_paths() {
+            let path = dir.join(&installinator_doc_hash);
+            assert!(path.exists(), "{path} does not exist");
+        }
+    }
+    // Collect watchers for all of the sled-agent artifact delete reconcilers.
+    let mut delete_watchers = cptestctx
+        .sled_agents
+        .iter()
+        .map(|sled_agent| {
+            sled_agent.sled_agent().artifact_store().create_delete_watcher()
+        })
+        .collect::<Vec<_>>();
+    // Manually prune the first repo.
+    let initial_repo = datastore
+        .tuf_repo_get_by_version(
+            &opctx,
+            "1.0.0".parse::<Version>().unwrap().into(),
+        )
+        .await?;
+    let recent_releases =
+        datastore.target_release_fetch_recent_distinct(&opctx, 3).await?;
+    datastore
+        .tuf_repo_mark_pruned(
+            &opctx,
+            status.generation,
+            &recent_releases,
+            initial_repo.repo.id(),
+        )
+        .await
+        .unwrap();
+    // Marking a repository as pruned bumps the generation number.
+    assert_eq!(
+        datastore.tuf_get_generation(&opctx).await.unwrap(),
+        4u32.into()
+    );
+    // Run the replication background task; we should see new configs be put.
+    let status =
+        run_tuf_artifact_replication_step(&cptestctx.lockstep_client).await;
+    eprintln!("{status:?}");
+    assert_eq!(status.last_run_counters.put_config_ok, 4);
+    assert_eq!(status.last_run_counters.list_ok, 4);
+    assert_eq!(status.last_run_counters.sum(), 8);
+    assert_eq!(status.generation, 4u32.into());
+    // Wait for the delete reconciler to finish on all sled agents.
+    futures::future::join_all(
+        delete_watchers.iter_mut().map(|watcher| watcher.changed()),
+    )
+    .await;
+    // Verify the installinator document from the initial repo is deleted.
+    for sled_agent in &cptestctx.sled_agents {
+        for dir in sled_agent.sled_agent().artifact_store().storage_paths() {
+            let path = dir.join(&installinator_doc_hash);
+            assert!(!path.exists(), "{path} was not deleted");
+        }
+    }
 
     cptestctx.teardown().await;
     Ok(())

sled-agent/src/artifact_store.rs

@@ -81,10 +81,6 @@ pub struct ArtifactStore<T: DatasetsManager> {
     ledger_tx: mpsc::Sender<LedgerManagerRequest>,
     config: watch::Receiver<Option<ArtifactConfig>>,
     pub(crate) storage: T,
-
-    /// Used for synchronization in unit tests.
-    #[cfg(test)]
-    delete_done: watch::Receiver<Generation>,
 }
 
 impl<T: DatasetsManager> ArtifactStore<T> {
@@ -136,14 +132,10 @@ impl<T: DatasetsManager> ArtifactStore<T> {
             config_tx,
         ));
 
-        #[cfg(test)]
-        let (done_signal, delete_done) = watch::channel(0u32.into());
         tokio::task::spawn(delete_reconciler(
             log.clone(),
             storage.clone(),
             config.clone(),
-            #[cfg(test)]
-            done_signal,
         ));
 
         ArtifactStore {
@@ -155,9 +147,6 @@ impl<T: DatasetsManager> ArtifactStore<T> {
             ledger_tx,
             config,
             storage,
-
-            #[cfg(test)]
-            delete_done,
         }
     }
 }
@@ -506,7 +495,6 @@ async fn delete_reconciler<T: DatasetsManager>(
     log: Logger,
     storage: T,
     mut receiver: watch::Receiver<Option<ArtifactConfig>>,
-    #[cfg(test)] done_signal: watch::Sender<Generation>,
 ) {
     while let Ok(()) = receiver.changed().await {
         let generation = match receiver.borrow_and_update().as_ref() {
@@ -580,12 +568,7 @@ async fn delete_reconciler<T: DatasetsManager>(
                 }
             }
         }
-        #[cfg(test)]
-        done_signal.send_if_modified(|old| {
-            let modified = *old != generation;
-            *old = generation;
-            modified
-        });
+        storage.signal_delete_done(generation);
     }
     warn!(log, "Delete reconciler sender dropped");
 }
@@ -602,6 +585,8 @@ pub trait DatasetsManager: Clone + Send + Sync + 'static {
     async fn copy_permit(&self) -> Option<OwnedSemaphorePermit> {
         None
     }
+
+    fn signal_delete_done(&self, _generation: Generation) {}
 }
 
 impl DatasetsManager for InternalDisksReceiver {
@@ -908,16 +893,20 @@ mod test {
     use camino_tempfile::Utf8TempDir;
     use futures::stream::{self, StreamExt};
     use hex_literal::hex;
+    use omicron_common::api::external::Generation;
     use omicron_test_utils::dev::test_setup_log;
     use sled_agent_api::ArtifactConfig;
     use tokio::io::AsyncReadExt;
     use tokio::sync::oneshot;
+    use tokio::sync::watch;
     use tufaceous_artifact::ArtifactHash;
 
     use super::{ArtifactStore, DatasetsManager, Error};
 
     #[derive(Clone)]
     struct TestBackend {
+        delete_done_tx: watch::Sender<Generation>,
+        delete_done_rx: watch::Receiver<Generation>,
         datasets: Vec<Utf8PathBuf>,
         _tempdir: Arc<Utf8TempDir>,
     }
@@ -934,7 +923,13 @@ mod test {
                 datasets.push(dataset)
             }
 
-            TestBackend { datasets, _tempdir: tempdir }
+            let (delete_done_tx, delete_done_rx) = watch::channel(0u32.into());
+            TestBackend {
+                delete_done_tx,
+                delete_done_rx,
+                datasets,
+                _tempdir: tempdir,
+            }
         }
     }
 
@@ -944,6 +939,14 @@ mod test {
         ) -> impl Iterator<Item = camino::Utf8PathBuf> + '_ {
             self.datasets.iter().cloned()
         }
+
+        fn signal_delete_done(&self, generation: Generation) {
+            self.delete_done_tx.send_if_modified(|old| {
+                let modified = *old != generation;
+                *old = generation;
+                modified
+            });
+        }
     }
 
     const TEST_ARTIFACT: Bytes = Bytes::from_static(b"I'm an artifact!\n");
@@ -1121,7 +1124,7 @@ mod test {
         }
 
         // clear `delete_done` so we can synchronize with the delete reconciler
-        store.delete_done.mark_unchanged();
+        store.storage.delete_done_rx.mark_unchanged();
         // put a new config that says we don't want the artifact anymore.
         config.generation = config.generation.next();
         config.artifacts.remove(&TEST_HASH);
@@ -1130,7 +1133,7 @@ mod test {
         // has actually occurred yet
         assert!(store.list().await.unwrap().list.is_empty());
         // wait for deletion to actually complete
-        store.delete_done.changed().await.unwrap();
+        store.storage.delete_done_rx.changed().await.unwrap();
         // get fails, because it has been deleted
         assert!(matches!(
             store.get(TEST_HASH).await,

sled-agent/src/sim/artifact_store.rs

@@ -7,13 +7,15 @@
 
 use std::sync::Arc;
 
+use camino::Utf8Path;
 use camino_tempfile::Utf8TempDir;
 use dropshot::{
     Body, ConfigDropshot, FreeformBody, HttpError, HttpResponseOk, HttpServer,
     Path, RequestContext, ServerBuilder,
 };
+use omicron_common::api::external::Generation;
 use repo_depot_api::*;
-use tokio::sync::{OwnedSemaphorePermit, Semaphore};
+use tokio::sync::{OwnedSemaphorePermit, Semaphore, watch};
 
 use crate::artifact_store::{ArtifactStore, DatasetsManager};
 
@@ -28,10 +30,16 @@ pub struct SimArtifactStorage {
     // Semaphore to keep track of how many copy requests are in flight, and to
     // be able to await on their completion. Used in integration tests.
     copy_semaphore: Arc<Semaphore>,
+
+    // Watch channel to be able to await on the delete reconciler completing in
+    // integration tests.
+    delete_done_tx: watch::Sender<Generation>,
+    delete_done_rx: watch::Receiver<Generation>,
 }
 
 impl SimArtifactStorage {
     pub(super) fn new() -> SimArtifactStorage {
+        let (delete_done_tx, delete_done_rx) = watch::channel(0u32.into());
         SimArtifactStorage {
             dirs: Arc::new([
                 camino_tempfile::tempdir().unwrap(),
@@ -40,6 +48,8 @@ impl SimArtifactStorage {
             copy_semaphore: Arc::new(
                 const { Semaphore::const_new(MAX_PERMITS as usize) },
             ),
+            delete_done_tx,
+            delete_done_rx,
         }
     }
 }
@@ -54,6 +64,14 @@ impl DatasetsManager for SimArtifactStorage {
     async fn copy_permit(&self) -> Option<OwnedSemaphorePermit> {
         Some(self.copy_semaphore.clone().acquire_owned().await.unwrap())
     }
+
+    fn signal_delete_done(&self, generation: Generation) {
+        self.delete_done_tx.send_if_modified(|old| {
+            let modified = *old != generation;
+            *old = generation;
+            modified
+        });
+    }
 }
 
 impl ArtifactStore<SimArtifactStorage> {
@@ -73,6 +91,10 @@ impl ArtifactStore<SimArtifactStorage> {
         .unwrap()
     }
 
+    pub fn storage_paths(&self) -> impl Iterator<Item = &Utf8Path> {
+        self.storage.dirs.iter().map(|p| p.path())
+    }
+
     pub async fn wait_for_copy_tasks(&self) {
         // Acquire a permit for MAX_PERMITS, which requires that all copy tasks
         // have dropped their permits. Then immediately drop it.
@@ -83,6 +105,12 @@ impl ArtifactStore<SimArtifactStorage> {
             .await
             .unwrap();
     }
+
+    pub fn create_delete_watcher(&self) -> watch::Receiver<Generation> {
+        let mut watcher = self.storage.delete_done_rx.clone();
+        watcher.mark_unchanged();
+        watcher
+    }
 }
 
 /// Implementation of the Repo Depot API backed by an