remake: download wakes on end instead of ticker

1 files changed, 14 insertions, 2 deletions

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.