fix test_repo_upload flake (#8066)

Closes #8057.

Author: iliana

nexus/tests/integration_tests/updates.rs

@@ -179,6 +179,12 @@ async fn test_repo_upload() -> Result<()> {
     // The local repo is not deleted until the next task run.
     assert_eq!(status.local_repos, 1);
 
+    // Wait for all the copy requests to complete.
+    futures::future::join_all(cptestctx.sled_agents.iter().map(|sled_agent| {
+        sled_agent.sled_agent().artifact_store().wait_for_copy_tasks()
+    }))
+    .await;
+
     // Run the replication background task again; the local repos should be
     // dropped.
     let status =

sled-agent/src/artifact_store.rs

@@ -45,7 +45,7 @@ use sled_storage::manager::StorageHandle;
 use slog::{Logger, error, info};
 use slog_error_chain::{InlineErrorChain, SlogInlineError};
 use tokio::fs::File;
-use tokio::sync::{mpsc, oneshot, watch};
+use tokio::sync::{OwnedSemaphorePermit, mpsc, oneshot, watch};
 use tokio::task::JoinSet;
 use tufaceous_artifact::ArtifactHash;
 
@@ -77,12 +77,12 @@ const TEMP_SUBDIR: &str = "tmp";
 /// - for PUT, we try to write to all datasets, logging errors as we go; if we
 ///   successfully write the artifact to at least one, we return OK.
 /// - for GET, we look in each dataset until we find it.
-pub(crate) struct ArtifactStore<T: DatasetsManager> {
+pub struct ArtifactStore<T: DatasetsManager> {
     log: Logger,
     reqwest_client: reqwest::Client,
     ledger_tx: mpsc::Sender<LedgerManagerRequest>,
     config: watch::Receiver<Option<ArtifactConfig>>,
-    storage: T,
+    pub(crate) storage: T,
 
     /// Used for synchronization in unit tests.
     #[cfg(test)]
@@ -400,6 +400,7 @@ impl<T: DatasetsManager> ArtifactStore<T> {
     ) -> Result<(), Error> {
         // Check that there's no conflict before we send the upstream request.
         let writer = self.writer(sha256, generation).await?;
+        let permit = self.storage.copy_permit().await;
 
         let client = repo_depot_client::Client::new_with_client(
             depot_base_url,
@@ -421,6 +422,7 @@ impl<T: DatasetsManager> ArtifactStore<T> {
         let log = self.log.clone();
         let base_url = depot_base_url.to_owned();
         tokio::task::spawn(async move {
+            let _permit = permit;
             let stream = response.into_inner().into_inner().map_err(|err| {
                 Error::DepotCopy {
                     sha256,
@@ -601,10 +603,15 @@ async fn delete_reconciler<T: DatasetsManager>(
 /// Abstracts over what kind of sled agent we are; each of the real sled agent,
 /// simulated sled agent, and this module's unit tests have different ways of
 /// keeping track of the datasets on the system.
-pub(crate) trait DatasetsManager: Clone + Send + Sync + 'static {
+pub trait DatasetsManager: Clone + Send + Sync + 'static {
     fn artifact_storage_paths(
         &self,
     ) -> impl Future<Output = impl Iterator<Item = Utf8PathBuf> + Send + '_> + Send;
+
+    #[expect(async_fn_in_trait)]
+    async fn copy_permit(&self) -> Option<OwnedSemaphorePermit> {
+        None
+    }
 }
 
 impl DatasetsManager for StorageHandle {

sled-agent/src/sim/artifact_store.rs

@@ -5,20 +5,29 @@
 //! Implementation of `crate::artifact_store::StorageBackend` for our simulated
 //! storage.
 
+use std::sync::Arc;
+
 use camino_tempfile::Utf8TempDir;
 use dropshot::{
     Body, ConfigDropshot, FreeformBody, HttpError, HttpResponseOk, HttpServer,
     Path, RequestContext, ServerBuilder,
 };
 use repo_depot_api::*;
-use std::sync::Arc;
+use tokio::sync::{OwnedSemaphorePermit, Semaphore};
 
 use crate::artifact_store::{ArtifactStore, DatasetsManager};
 
+// Semaphore mostly uses usize but in `acquire_many` it unfortunately uses u32.
+const MAX_PERMITS: u32 = u32::MAX >> 3;
+
 #[derive(Clone)]
-pub(super) struct SimArtifactStorage {
+pub struct SimArtifactStorage {
     // We simulate the two M.2s with two separate temporary directories.
     dirs: Arc<[Utf8TempDir; 2]>,
+
+    // 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>,
 }
 
 impl SimArtifactStorage {
@@ -28,6 +37,9 @@ impl SimArtifactStorage {
                 camino_tempfile::tempdir().unwrap(),
                 camino_tempfile::tempdir().unwrap(),
             ]),
+            copy_semaphore: Arc::new(
+                const { Semaphore::const_new(MAX_PERMITS as usize) },
+            ),
         }
     }
 }
@@ -38,6 +50,10 @@ impl DatasetsManager for SimArtifactStorage {
     ) -> impl Iterator<Item = camino::Utf8PathBuf> + '_ {
         self.dirs.iter().map(|tempdir| tempdir.path().to_owned())
     }
+
+    async fn copy_permit(&self) -> Option<OwnedSemaphorePermit> {
+        Some(self.copy_semaphore.clone().acquire_owned().await.unwrap())
+    }
 }
 
 impl ArtifactStore<SimArtifactStorage> {
@@ -56,6 +72,17 @@ impl ArtifactStore<SimArtifactStorage> {
         .start()
         .unwrap()
     }
+
+    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.
+        let _permit = self
+            .storage
+            .copy_semaphore
+            .acquire_many(MAX_PERMITS)
+            .await
+            .unwrap();
+    }
 }
 
 /// Implementation of the Repo Depot API backed by an

sled-agent/src/sim/sled_agent.rs

@@ -504,7 +504,7 @@ impl SledAgent {
         &self.updates
     }
 
-    pub(super) fn artifact_store(&self) -> &ArtifactStore<SimArtifactStorage> {
+    pub fn artifact_store(&self) -> &ArtifactStore<SimArtifactStorage> {
         self.repo_depot.app_private()
     }