Move Leader lock to separate file

Author: m-sz

crates/autopilot/src/leader_lock_tracker.rs

@@ -0,0 +1,129 @@
+use {database::leader_pg_lock::LeaderLock, observe::metrics};
+
+/// Tracks the autopilot leader lock status.
+/// Leader lock status is tracked based on calls to try_acquire()
+#[expect(clippy::large_enum_variant)]
+pub enum LeaderLockTracker {
+    /// Leader lock mechanism is disabled.
+    /// Only one autopilot instance should be running at all times.
+    Disabled,
+    /// Leader lock mechanism is enabled.
+    /// This allows for multiple autopilots to run simultaneously with only one
+    /// driving the auctions forward.
+    /// The autopilot followers (not holding the lock) will keep their caches
+    /// warm. It facilitates zero downtime deployments.
+    Enabled {
+        /// Are we currently the leader since last call to try_acquire()
+        is_leader: bool,
+        /// Were we previously the leader sinec the last call to try_acquire()
+        was_leader: bool,
+        leader_lock: LeaderLock,
+    },
+}
+
+impl LeaderLockTracker {
+    pub fn new(leader_lock: Option<LeaderLock>) -> Self {
+        match leader_lock {
+            Some(leader_lock) => Self::Enabled {
+                is_leader: false,
+                was_leader: false,
+                leader_lock,
+            },
+            None => Self::Disabled,
+        }
+    }
+
+    /// Tries to acquire the leader lock if it's enabled
+    /// If not, does nothing
+    /// Should be called at the beginning of every run loop iteration
+    pub async fn try_acquire(&mut self) {
+        let Self::Enabled {
+            is_leader,
+            was_leader,
+            leader_lock,
+        } = self
+        else {
+            return;
+        };
+
+        *was_leader = *is_leader;
+
+        *is_leader = leader_lock.try_acquire().await.unwrap_or_else(|err| {
+            tracing::error!(?err, "failed to become leader");
+            Metrics::leader_lock_error();
+            false
+        });
+
+        if self.just_stepped_up() {
+            tracing::info!("Stepped up as a leader");
+            Metrics::leader_step_up();
+        }
+    }
+
+    /// Releases the leader lock if it was held
+    /// Should be called after breaking out of run loop (for example: due to
+    /// shutdown)
+    pub async fn release(self) {
+        if let Self::Enabled {
+            mut leader_lock,
+            is_leader: true,
+            ..
+        } = self
+        {
+            tracing::info!("Shutdown received, stepping down as the leader");
+            leader_lock.release().await;
+            Metrics::leader_step_down();
+        }
+    }
+
+    /// Returns true if the last try_acquire call resulted in acquiring the
+    /// leader lock If the feature is disabled, always returns false
+    pub fn just_stepped_up(&self) -> bool {
+        matches!(
+            self,
+            Self::Enabled {
+                is_leader: true,
+                was_leader: false,
+                ..
+            }
+        )
+    }
+
+    /// Returns true if the leader lock is being held
+    /// If the feature is disabled, always returns true
+    pub fn is_leader(&self) -> bool {
+        match self {
+            Self::Enabled { is_leader, .. } => *is_leader,
+            _ => true,
+        }
+    }
+}
+#[derive(prometheus_metric_storage::MetricStorage)]
+#[metric(subsystem = "leader_lock_tracker")]
+struct Metrics {
+    /// Tracks the current leader status
+    /// 1 - is currently autopilot leader
+    /// 0 - is currently autopilot follower
+    is_leader: prometheus::IntGauge,
+
+    /// Trackes the count of errors acquiring leader lock (should never happen)
+    leader_lock_error: prometheus::IntCounter,
+}
+
+impl Metrics {
+    fn get() -> &'static Self {
+        Metrics::instance(metrics::get_storage_registry()).unwrap()
+    }
+
+    fn leader_step_up() {
+        Self::get().is_leader.set(1)
+    }
+
+    fn leader_step_down() {
+        Self::get().is_leader.set(0)
+    }
+
+    fn leader_lock_error() {
+        Self::get().leader_lock_error.inc()
+    }
+}

crates/autopilot/src/lib.rs

@@ -4,6 +4,7 @@ pub mod database;
 pub mod domain;
 pub mod event_updater;
 pub mod infra;
+mod leader_lock_tracker;
 mod maintenance;
 pub mod periodic_db_cleanup;
 pub mod run;

crates/autopilot/src/run_loop.rs

@@ -20,14 +20,15 @@ use {
             self,
             solvers::dto::{settle, solve},
         },
+        leader_lock_tracker::LeaderLockTracker,
         maintenance::Maintenance,
         run::Liveness,
         shutdown_controller::ShutdownController,
         solvable_orders::SolvableOrdersCache,
     },
     ::observe::metrics,
     anyhow::{Context, Result},
-    database::{leader_pg_lock::LeaderLock, order_events::OrderEventLabel},
+    database::order_events::OrderEventLabel,
     ethrpc::block_stream::BlockInfo,
     futures::{FutureExt, TryFutureExt},
     itertools::Itertools,
@@ -90,95 +91,6 @@ pub struct RunLoop {
     winner_selection: winner_selection::Arbitrator,
 }
 
-/// Tracks the autopilot leader lock status.
-#[expect(clippy::large_enum_variant)]
-enum LeaderLockTracker {
-    Disabled,
-    Enabled {
-        is_leader: bool,
-        was_leader: bool,
-        leader_lock: LeaderLock,
-    },
-}
-
-impl LeaderLockTracker {
-    pub fn new(leader_lock: Option<LeaderLock>) -> Self {
-        match leader_lock {
-            Some(leader_lock) => Self::Enabled {
-                is_leader: false,
-                was_leader: false,
-                leader_lock,
-            },
-            None => Self::Disabled,
-        }
-    }
-
-    /// Tries to acquire the leader lock if it's enabled
-    /// If not, does nothing
-    /// Should be called at the beginning of every run loop iteration
-    pub async fn try_acquire(&mut self) {
-        let Self::Enabled {
-            is_leader,
-            was_leader,
-            leader_lock,
-        } = self
-        else {
-            return;
-        };
-
-        *was_leader = *is_leader;
-
-        *is_leader = leader_lock.try_acquire().await.unwrap_or_else(|err| {
-            tracing::error!(?err, "failed to become leader");
-            Metrics::leader_lock_error();
-            false
-        });
-
-        if self.just_stepped_up() {
-            tracing::info!("Stepped up as a leader");
-            Metrics::leader_step_up();
-        }
-    }
-
-    /// Releases the leader lock if it was held
-    /// Should be called after breaking out of run loop (for example: due to
-    /// shutdown)
-    pub async fn release(self) {
-        if let Self::Enabled {
-            mut leader_lock,
-            is_leader: true,
-            ..
-        } = self
-        {
-            tracing::info!("Shutdown received, stepping down as the leader");
-            leader_lock.release().await;
-            Metrics::leader_step_down();
-        }
-    }
-
-    /// Returns true if the last try_acquire call resulted in acquiring the
-    /// leader lock If the feature is disabled, always returns false
-    pub fn just_stepped_up(&self) -> bool {
-        match self {
-            Self::Enabled {
-                is_leader: true,
-                was_leader: false,
-                ..
-            } => true,
-            _ => false,
-        }
-    }
-
-    /// Returns true if the leader lock is being held
-    /// If the feature is disabled, always returns true
-    pub fn is_leader(&self) -> bool {
-        match self {
-            Self::Enabled { is_leader, .. } => *is_leader,
-            _ => true,
-        }
-    }
-}
-
 impl RunLoop {
     #[expect(clippy::too_many_arguments)]
     pub fn new(
@@ -1078,14 +990,6 @@ struct Metrics {
     /// function is started.
     #[metric(buckets(0, 0.25, 0.5, 0.75, 1, 1.5, 2, 2.5, 3, 4, 5, 6))]
     current_block_delay: prometheus::Histogram,
-
-    /// Tracks the current leader status
-    /// 1 - is currently autopilot leader
-    /// 0 - is currently autopilot follower
-    is_leader: prometheus::IntGauge,
-
-    /// Trackes the count of errors acquiring leader lock (should never happen)
-    leader_lock_error: prometheus::IntCounter,
 }
 
 impl Metrics {
@@ -1186,18 +1090,6 @@ impl Metrics {
             .current_block_delay
             .observe(init_block_timestamp.elapsed().as_secs_f64())
     }
-
-    fn leader_step_up() {
-        Self::get().is_leader.set(1)
-    }
-
-    fn leader_step_down() {
-        Self::get().is_leader.set(0)
-    }
-
-    fn leader_lock_error() {
-        Self::get().leader_lock_error.inc()
-    }
 }
 
 pub mod observe {

crates/observe/src/metrics.rs

@@ -91,9 +91,9 @@ pub fn serve_metrics(
     startup: Arc<Option<AtomicBool>>,
 ) -> JoinHandle<()> {
     let filter = handle_metrics()
-        .or(handle_liveness(liveness))
-        .or(handle_readiness(readiness))
-        .or(handle_startup(startup));
+        .or(handle_liveness_probe(liveness))
+        .or(handle_readiness_probe(readiness))
+        .or(handle_startup_probe(startup));
     tracing::info!(%address, "serving metrics");
     task::spawn(warp::serve(filter).bind(address))
 }
@@ -104,7 +104,7 @@ pub fn handle_metrics() -> impl Filter<Extract = (impl Reply,), Error = Rejectio
     warp::path("metrics").map(move || encode(registry))
 }
 
-fn handle_liveness(
+fn handle_liveness_probe(
     liveness_checker: Arc<dyn LivenessChecking>,
 ) -> impl Filter<Extract = (impl Reply,), Error = Rejection> + Clone {
     warp::path("liveness").and_then(move || {
@@ -120,7 +120,7 @@ fn handle_liveness(
     })
 }
 
-fn handle_readiness(
+fn handle_readiness_probe(
     readiness: Arc<Option<AtomicBool>>,
 ) -> impl Filter<Extract = (impl Reply,), Error = Rejection> + Clone {
     warp::path("ready").and_then(move || {
@@ -143,7 +143,7 @@ fn handle_readiness(
     })
 }
 
-fn handle_startup(
+fn handle_startup_probe(
     startup: Arc<Option<AtomicBool>>,
 ) -> impl Filter<Extract = (impl Reply,), Error = Rejection> + Clone {
     warp::path("startup").and_then(move || {