remake: download wakes on end instead of ticker

2 files changed, 15 insertions, 6 deletions

diff --git a/src/remake/content.rs b/src/remake/content.rs
index d036b87..467d083 100644
--- a/src/remake/content.rs
+++ b/src/remake/content.rs
@@ -11,7 +11,6 @@
 
 use std::collections::HashSet;
 use std::sync::Arc;
-use std::time::Duration;
 
 use tracing::{info, warn};
 
@@ -59,9 +58,7 @@ pub(super) async fn ensure<H: Host>(
                 .ok_or_else(|| format!("couldn't download content {}", reference.blake3));
         }
         // Someone else is downloading it.
-        while host.remake().content_fetch_running(&hash) {
-            tokio::time::sleep(Duration::from_millis(500)).await;
-        }
+        host.remake().content_fetch_ended(&hash).await;
     }
 }
 
diff --git a/src/remake/local.rs b/src/remake/local.rs
index ce133ec..c80a5a5 100644
--- a/src/remake/local.rs
+++ b/src/remake/local.rs
@@ -25,6 +25,8 @@ pub(crate) struct RemakeState {
     started_at: UnixMillis,
     // A std mutex: every access is short and never spans an await.
     inner: Mutex<Inner>,
+    /// A download of content ended (`end_content_fetch`).
+    content_fetch_ended: tokio::sync::Notify,
 }
 
 impl RemakeState {
@@ -33,6 +35,7 @@ impl RemakeState {
         Self {
             started_at,
             inner: Mutex::default(),
+            content_fetch_ended: tokio::sync::Notify::new(),
         }
     }
 }
@@ -339,10 +342,19 @@ impl RemakeState {
 
     pub(crate) fn end_content_fetch(&self, hash: &[u8; 32]) {
         self.lock().content_fetches.remove(hash);
+        self.content_fetch_ended.notify_waiters();
     }
 
-    pub(crate) fn content_fetch_running(&self, hash: &[u8; 32]) -> bool {
-        self.lock().content_fetches.contains(hash)
+    /// Waits until no download of `hash` runs.
+    pub(crate) async fn content_fetch_ended(&self, hash: &[u8; 32]) {
+        loop {
+            // Before looking, so that an end in between isn't missed.
+            let ended = self.content_fetch_ended.notified();
+            if !self.lock().content_fetches.contains(hash) {
+                return;
+            }
+            ended.await;
+        }
     }
 
     /// Drops content that no remake in `referenced` needs, once its grace has passed.