Fix check_download_progress [scheduler]

It was jailing workers unfairly after a new assignment has been applied
subsquid · Dec 23, 2024 · 1d00914 · 1d00914
1 parent ecadac8
commit 1d00914
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 7 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/crates/network-scheduler/Cargo.toml b/crates/network-scheduler/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "network-scheduler"
-version = "2.0.0"
+version = "2.0.1"
 edition = "2021"
 
 [dependencies]

diff --git a/crates/network-scheduler/src/worker_state.rs b/crates/network-scheduler/src/worker_state.rs
@@ -177,10 +177,7 @@ impl WorkerState {
 
     /// Check if the worker is making progress with downloading missing chunks.
     /// Returns true iff the worker is fully synced or making progress.
-    pub fn check_download_progress(
-        &mut self,
-        _units: &DashMap<UnitId, SchedulingUnit>,
-    ) -> bool {
+    pub fn check_download_progress(&mut self, _units: &DashMap<UnitId, SchedulingUnit>) -> bool {
         assert!(!self.jailed);
         let Some(last_assignment) = self.last_assignment.as_ref() else {
             return true; // worker doesn't have any assignment
@@ -215,7 +212,12 @@ impl WorkerState {
         }
     }
 
-    pub fn reset_download_progress(&mut self, _units: &DashMap<UnitId, SchedulingUnit>) {
+    pub fn reset_download_progress(&mut self, units: &DashMap<UnitId, SchedulingUnit>) {
+        self.num_missing_chunks = self
+            .assigned_units
+            .iter()
+            .map(|unit| units.get(unit).map_or(0, |u| u.num_chunks()))
+            .sum::<usize>() as u32;
         self.last_assignment = Some(SystemTime::now());
     }