Refactoring: load balancer

ankitects · dae · Mar 26, 2025 · Mar 12, 2025 · Mar 12, 2025 · Mar 14, 2025
commit 939fd4ab1aaa87cd5b2a9caabd248eb5f2ff0374
diff --git a/proto/anki/collection.proto b/proto/anki/collection.proto
@@ -19,6 +19,7 @@ service CollectionService {
   rpc MergeUndoEntries(generic.UInt32) returns (OpChanges);
   rpc LatestProgress(generic.Empty) returns (Progress);
   rpc SetWantsAbort(generic.Empty) returns (generic.Empty);
+  rpc ClearStudyQueues(generic.Empty) returns (generic.Empty);
 }
 
 // Implicitly includes any of the above methods that are not listed in the

diff --git a/pylib/anki/collection.py b/pylib/anki/collection.py
@@ -986,6 +986,7 @@ def _get_enable_load_balancer(self) -> bool:
         return self.get_config_bool(Config.Bool.LOAD_BALANCER_ENABLED)
 
     def _set_enable_load_balancer(self, value: bool) -> None:
+        self._backend.clear_study_queues()
         self.set_config_bool(Config.Bool.LOAD_BALANCER_ENABLED, value)
 
     load_balancer_enabled = property(

diff --git a/rslib/src/collection/service.rs b/rslib/src/collection/service.rs
@@ -49,4 +49,9 @@ impl crate::services::CollectionService for Collection {
         self.state.progress.lock().unwrap().want_abort = true;
         Ok(())
     }
+
+    fn clear_study_queues(&mut self) -> error::Result<()> {
+        self.clear_study_queues();
+        Ok(())
+    }
 }
diff --git a/rslib/src/scheduler/answering/mod.rs b/rslib/src/scheduler/answering/mod.rs
@@ -83,7 +83,7 @@ impl CardStateUpdater {
     /// state handling code from the rest of the Anki codebase.
     pub(crate) fn state_context<'a>(
         &'a self,
-        load_balancer: Option<LoadBalancerContext<'a>>,
+        load_balancer_ctx: Option<LoadBalancerContext<'a>>,
     ) -> StateContext<'a> {
         StateContext {
             fuzz_factor: get_fuzz_factor(self.fuzz_seed),
@@ -96,8 +96,8 @@ impl CardStateUpdater {
             interval_multiplier: self.config.inner.interval_multiplier,
             maximum_review_interval: self.config.inner.maximum_review_interval,
             leech_threshold: self.config.inner.leech_threshold,
-            load_balancer: load_balancer
-                .map(|load_balancer| load_balancer.set_fuzz_seed(self.fuzz_seed)),
+            load_balancer_ctx: load_balancer_ctx
+                .map(|load_balancer_ctx| load_balancer_ctx.set_fuzz_seed(self.fuzz_seed)),
             relearn_steps: self.relearn_steps(),
             lapse_multiplier: self.config.inner.lapse_multiplier,
             minimum_lapse_interval: self.config.inner.minimum_lapse_interval,
@@ -240,22 +240,23 @@ impl Collection {
         let ctx = self.card_state_updater(card)?;
         let current = ctx.current_card_state();
 
-        let load_balancer = self
+        let load_balancer_ctx = self
             .get_config_bool(BoolKey::LoadBalancerEnabled)
             .then(|| {
                 let deckconfig_id = deck.config_id();
 
                 self.state.card_queues.as_ref().and_then(|card_queues| {
-                    Some(
-                        card_queues
-                            .load_balancer
-                            .review_context(note_id, deckconfig_id?),
-                    )
+                    match &card_queues.load_balancer {
+                        None => None,
+                        Some(load_balancer) => {
+                            Some(load_balancer.review_context(note_id, deckconfig_id?))
+                        }
+                    }
                 })
             })
             .flatten();
 
-        let state_ctx = ctx.state_context(load_balancer);
+        let state_ctx = ctx.state_context(load_balancer_ctx);
         Ok(current.next_states(&state_ctx))
     }
 
@@ -349,16 +350,13 @@ impl Collection {
             self.add_leech_tag(card.note_id)?;
         }
 
-        if card.queue == CardQueue::Review {
+        if self.get_config_bool(BoolKey::LoadBalancerEnabled) && card.queue == CardQueue::Review {
             let deck = self.get_deck(card.deck_id)?;
             if let Some(card_queues) = self.state.card_queues.as_mut() {
                 if let Some(deckconfig_id) = deck.and_then(|deck| deck.config_id()) {
-                    card_queues.load_balancer.add_card(
-                        card.id,
-                        card.note_id,
-                        deckconfig_id,
-                        card.interval,
-                    )
+                    if let Some(load_balancer) = &mut card_queues.load_balancer {
+                        load_balancer.add_card(card.id, card.note_id, deckconfig_id, card.interval)
+                    }
                 }
             }
         }

diff --git a/rslib/src/scheduler/queue/builder/mod.rs b/rslib/src/scheduler/queue/builder/mod.rs
@@ -107,7 +107,7 @@ pub(super) struct QueueBuilder {
     pub(super) learning: Vec<DueCard>,
     pub(super) day_learning: Vec<DueCard>,
     limits: LimitTreeMap,
-    load_balancer: LoadBalancer,
+    load_balancer: Option<LoadBalancer>,
     context: Context,
 }
 
@@ -146,16 +146,21 @@ impl QueueBuilder {
         let sort_options = sort_options(&root_deck, &config_map);
         let deck_map = col.storage.get_decks_map()?;
 
-        let did_to_dcid = deck_map
-            .values()
-            .filter_map(|deck| Some((deck.id, deck.config_id()?)))
-            .collect::<HashMap<_, _>>();
-        let load_balancer = LoadBalancer::new(
-            timing.days_elapsed,
-            did_to_dcid,
-            col.timing_today()?.next_day_at,
-            &col.storage,
-        )?;
+        let load_balancer = col
+            .get_config_bool(BoolKey::LoadBalancerEnabled)
+            .then(|| {
+                let did_to_dcid = deck_map
+                    .values()
+                    .filter_map(|deck| Some((deck.id, deck.config_id()?)))
+                    .collect::<HashMap<_, _>>();
+                LoadBalancer::new(
+                    timing.days_elapsed,
+                    did_to_dcid,
+                    col.timing_today()?.next_day_at,
+                    &col.storage,
+                )
+            })
+            .transpose()?;
 
         Ok(QueueBuilder {
             new: Vec::new(),

diff --git a/rslib/src/scheduler/queue/mod.rs b/rslib/src/scheduler/queue/mod.rs
@@ -38,7 +38,7 @@ pub(crate) struct CardQueues {
     /// counts are zero. Ensures we don't show a newly-due learning card after a
     /// user returns from editing a review card.
     current_learning_cutoff: TimestampSecs,
-    pub(crate) load_balancer: LoadBalancer,
+    pub(crate) load_balancer: Option<LoadBalancer>,
 }
 
 #[derive(Debug, Copy, Clone)]

diff --git a/rslib/src/scheduler/queue/undo.rs b/rslib/src/scheduler/queue/undo.rs
@@ -39,13 +39,17 @@ impl Collection {
                     queues.push_undo_entry(update.entry);
                 }
 
-                if let Some(card_queues) = self.state.card_queues.as_mut() {
-                    match &update.entry {
-                        QueueEntry::IntradayLearning(entry) => {
-                            card_queues.load_balancer.remove_card(entry.id);
-                        }
-                        QueueEntry::Main(entry) => {
-                            card_queues.load_balancer.remove_card(entry.id);
+                if self.get_config_bool(BoolKey::LoadBalancerEnabled) {
+                    if let Some(card_queues) = self.state.card_queues.as_mut() {
+                        if let Some(load_balancer) = &mut card_queues.load_balancer {
+                            match &update.entry {
+                                QueueEntry::IntradayLearning(entry) => {
+                                    load_balancer.remove_card(entry.id);
+                                }
+                                QueueEntry::Main(entry) => {
+                                    load_balancer.remove_card(entry.id);
+                                }
+                            }
                         }
                     }
                 }

diff --git a/rslib/src/scheduler/states/fuzz.rs b/rslib/src/scheduler/states/fuzz.rs
@@ -34,9 +34,11 @@ static FUZZ_RANGES: [FuzzRange; 3] = [
 impl StateContext<'_> {
     /// Apply fuzz, respecting the passed bounds.
     pub(crate) fn with_review_fuzz(&self, interval: f32, minimum: u32, maximum: u32) -> u32 {
-        self.load_balancer
+        self.load_balancer_ctx
             .as_ref()
-            .and_then(|load_balancer| load_balancer.find_interval(interval, minimum, maximum))
+            .and_then(|load_balancer_ctx| {
+                load_balancer_ctx.find_interval(interval, minimum, maximum)
+            })
             .unwrap_or_else(|| with_review_fuzz(self.fuzz_factor, interval, minimum, maximum))
     }
 }

diff --git a/rslib/src/scheduler/states/mod.rs b/rslib/src/scheduler/states/mod.rs
@@ -102,7 +102,7 @@ pub(crate) struct StateContext<'a> {
     pub interval_multiplier: f32,
     pub maximum_review_interval: u32,
     pub leech_threshold: u32,
-    pub load_balancer: Option<LoadBalancerContext<'a>>,
+    pub load_balancer_ctx: Option<LoadBalancerContext<'a>>,
 
     // relearning
     pub relearn_steps: LearningSteps<'a>,
@@ -137,7 +137,7 @@ impl StateContext<'_> {
             interval_multiplier: 1.0,
             maximum_review_interval: 36500,
             leech_threshold: 8,
-            load_balancer: None,
+            load_balancer_ctx: None,
             relearn_steps: LearningSteps::new(&[10.0]),
             lapse_multiplier: 0.0,
             minimum_lapse_interval: 1,