Refactor: Make Load Balancer Optional Throughout Codebase (#3860)

* Refactoring: load balancer

* Update about.py

* Refactoring: load balancer

* Update about.py

* Clean the code

* Remove config check from get_scheduling_states

* Backend method for the load balancer

* Refactor backend method for the load balancer

ftl/core/actions.ftl

@@ -50,6 +50,7 @@ actions-select = Select
 actions-shortcut-key = Shortcut key: { $val }
 actions-suspend-card = Suspend Card
 actions-set-due-date = Set Due Date
+actions-toggle-load-balancer = Toggle Load Balancer
 actions-grade-now = Grade Now
 actions-answer-card = Answer Card
 actions-unbury-unsuspend = Unbury/Unsuspend

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 SetLoadBalancerEnabled(generic.Bool) returns (OpChanges);
 }
 
 // Implicitly includes any of the above methods that are not listed in the

pylib/anki/collection.py

@@ -982,14 +982,14 @@ class Collection(DeprecatedNamesMixin):
         )
         return self.set_config(key, value, undoable=undoable)
 
-    def _get_enable_load_balancer(self) -> bool:
+    def _get_load_balancer_enabled(self) -> bool:
         return self.get_config_bool(Config.Bool.LOAD_BALANCER_ENABLED)
 
-    def _set_enable_load_balancer(self, value: bool) -> None:
+    def _set_load_balancer_enabled(self, value: bool) -> None:
-        self.set_config_bool(Config.Bool.LOAD_BALANCER_ENABLED, value)
+        self._backend.set_load_balancer_enabled(value)
 
     load_balancer_enabled = property(
-        fget=_get_enable_load_balancer, fset=_set_enable_load_balancer
+        fget=_get_load_balancer_enabled, fset=_set_load_balancer_enabled
     )
 
     def _get_enable_fsrs_short_term_with_steps(self) -> bool:

qt/aqt/about.py

@@ -218,6 +218,7 @@ def show(mw: aqt.AnkiQt) -> QDialog:
             "Luc Mcgrady",
             "Brayan Oliveira",
             "Market345",
+            "Yuki",
         )
     )
 

rslib/src/collection/service.rs

@@ -4,6 +4,8 @@ use anki_proto::generic;
 
 use crate::collection::Collection;
 use crate::error;
+use crate::prelude::BoolKey;
+use crate::prelude::Op;
 use crate::progress::progress_to_proto;
 
 impl crate::services::CollectionService for Collection {
@@ -49,4 +51,15 @@ impl crate::services::CollectionService for Collection {
         self.state.progress.lock().unwrap().want_abort = true;
         Ok(())
     }
+
+    fn set_load_balancer_enabled(
+        &mut self,
+        input: generic::Bool,
+    ) -> error::Result<anki_proto::collection::OpChanges> {
+        self.transact(Op::ToggleLoadBalancer, |col| {
+            col.set_config(BoolKey::LoadBalancerEnabled, &input.val)?;
+            Ok(())
+        })
+        .map(Into::into)
+    }
 }

rslib/src/ops.rs

@@ -36,6 +36,7 @@ pub enum Op {
     SetFlag,
     SortCards,
     Suspend,
+    ToggleLoadBalancer,
     UnburyUnsuspend,
     UpdateCard,
     UpdateConfig,
@@ -66,6 +67,7 @@ impl Op {
             Op::RenameDeck => tr.actions_rename_deck(),
             Op::ScheduleAsNew => tr.actions_forget_card(),
             Op::SetDueDate => tr.actions_set_due_date(),
+            Op::ToggleLoadBalancer => tr.actions_toggle_load_balancer(),
             Op::GradeNow => tr.actions_grade_now(),
             Op::Suspend => tr.studying_suspend(),
             Op::UnburyUnsuspend => tr.actions_unbury_unsuspend(),
@@ -170,7 +172,10 @@ impl OpChanges {
             || (c.config
                 && matches!(
                     self.op,
-                    Op::SetCurrentDeck | Op::UpdatePreferences | Op::UpdateDeckConfig
+                    Op::SetCurrentDeck
+                        | Op::UpdatePreferences
+                        | Op::UpdateDeckConfig
+                        | Op::ToggleLoadBalancer
                 ))
             || c.deck_config
     }

rslib/src/scheduler/answering/mod.rs

@@ -84,7 +84,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),
@@ -97,8 +97,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
+            load_balancer_ctx: load_balancer_ctx
-                .map(|load_balancer| load_balancer.set_fuzz_seed(self.fuzz_seed)),
+                .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,
@@ -241,22 +241,16 @@ impl Collection {
         let ctx = self.card_state_updater(card)?;
         let current = ctx.current_card_state();
 
-        let load_balancer = self
+        let load_balancer_ctx = self.state.card_queues.as_ref().and_then(|card_queues| {
-            .get_config_bool(BoolKey::LoadBalancerEnabled)
+            match card_queues.load_balancer.as_ref() {
-            .then(|| {
+                None => None,
-                let deckconfig_id = deck.config_id();
+                Some(load_balancer) => {
+                    Some(load_balancer.review_context(note_id, deck.config_id()?))
+                }
+            }
+        });
 
-                self.state.card_queues.as_ref().and_then(|card_queues| {
+        let state_ctx = ctx.state_context(load_balancer_ctx);
-                    Some(
-                        card_queues
-                            .load_balancer
-                            .review_context(note_id, deckconfig_id?),
-                    )
-                })
-            })
-            .flatten();
-
-        let state_ctx = ctx.state_context(load_balancer);
         Ok(current.next_states(&state_ctx))
     }
 
@@ -354,12 +348,9 @@ impl Collection {
             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(
+                    if let Some(load_balancer) = card_queues.load_balancer.as_mut() {
-                        card.id,
+                        load_balancer.add_card(card.id, card.note_id, deckconfig_id, card.interval)
-                        card.note_id,
+                    }
-                        deckconfig_id,
-                        card.interval,
-                    )
                 }
             }
         }

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
+        let load_balancer = col
-            .values()
+            .get_config_bool(BoolKey::LoadBalancerEnabled)
-            .filter_map(|deck| Some((deck.id, deck.config_id()?)))
+            .then(|| {
-            .collect::<HashMap<_, _>>();
+                let did_to_dcid = deck_map
-        let load_balancer = LoadBalancer::new(
+                    .values()
-            timing.days_elapsed,
+                    .filter_map(|deck| Some((deck.id, deck.config_id()?)))
-            did_to_dcid,
+                    .collect::<HashMap<_, _>>();
-            col.timing_today()?.next_day_at,
+                LoadBalancer::new(
-            &col.storage,
+                    timing.days_elapsed,
-        )?;
+                    did_to_dcid,
+                    col.timing_today()?.next_day_at,
+                    &col.storage,
+                )
+            })
+            .transpose()?;
 
         Ok(QueueBuilder {
             new: Vec::new(),

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)]

rslib/src/scheduler/queue/undo.rs

@@ -40,12 +40,14 @@ impl Collection {
                 }
 
                 if let Some(card_queues) = self.state.card_queues.as_mut() {
-                    match &update.entry {
+                    if let Some(load_balancer) = card_queues.load_balancer.as_mut() {
-                        QueueEntry::IntradayLearning(entry) => {
+                        match &update.entry {
-                            card_queues.load_balancer.remove_card(entry.id);
+                            QueueEntry::IntradayLearning(entry) => {
-                        }
+                                load_balancer.remove_card(entry.id);
-                        QueueEntry::Main(entry) => {
+                            }
-                            card_queues.load_balancer.remove_card(entry.id);
+                            QueueEntry::Main(entry) => {
+                                load_balancer.remove_card(entry.id);
+                            }
                         }
                     }
                 }

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))
     }
 }

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,