LCOV - code coverage report
Current view: top level - storage_controller/src - heartbeater.rs (source / functions) Coverage Total Hit
Test: 5445d246133daeceb0507e6cc0797ab7c1c70cb8.info Lines: 0.0 % 272 0
Test Date: 2025-03-12 18:05:02 Functions: 0.0 % 41 0

            Line data    Source code
       1              : use std::collections::HashMap;
       2              : use std::fmt::Debug;
       3              : use std::future::Future;
       4              : use std::sync::Arc;
       5              : use std::time::{Duration, Instant};
       6              : 
       7              : use futures::StreamExt;
       8              : use futures::stream::FuturesUnordered;
       9              : use pageserver_api::controller_api::{NodeAvailability, SkSchedulingPolicy};
      10              : use pageserver_api::models::PageserverUtilization;
      11              : use reqwest::Certificate;
      12              : use safekeeper_api::models::SafekeeperUtilization;
      13              : use safekeeper_client::mgmt_api;
      14              : use thiserror::Error;
      15              : use tokio_util::sync::CancellationToken;
      16              : use utils::id::NodeId;
      17              : use utils::logging::SecretString;
      18              : 
      19              : use crate::node::Node;
      20              : use crate::safekeeper::Safekeeper;
      21              : 
      22              : struct HeartbeaterTask<Server, State> {
      23              :     receiver: tokio::sync::mpsc::UnboundedReceiver<HeartbeatRequest<Server, State>>,
      24              :     cancel: CancellationToken,
      25              : 
      26              :     state: HashMap<NodeId, State>,
      27              : 
      28              :     max_offline_interval: Duration,
      29              :     max_warming_up_interval: Duration,
      30              :     jwt_token: Option<String>,
      31              :     ssl_ca_cert: Option<Certificate>,
      32              : }
      33              : 
      34              : #[derive(Debug, Clone)]
      35              : pub(crate) enum PageserverState {
      36              :     Available {
      37              :         last_seen_at: Instant,
      38              :         utilization: PageserverUtilization,
      39              :     },
      40              :     WarmingUp {
      41              :         started_at: Instant,
      42              :     },
      43              :     Offline,
      44              : }
      45              : 
      46              : #[derive(Debug, Clone)]
      47              : pub(crate) enum SafekeeperState {
      48              :     Available {
      49              :         last_seen_at: Instant,
      50              :         utilization: SafekeeperUtilization,
      51              :     },
      52              :     Offline,
      53              : }
      54              : 
      55              : #[derive(Debug)]
      56              : pub(crate) struct AvailablityDeltas<State>(pub Vec<(NodeId, State)>);
      57              : 
      58              : #[derive(Debug, Error)]
      59              : pub(crate) enum HeartbeaterError {
      60              :     #[error("Cancelled")]
      61              :     Cancel,
      62              : }
      63              : 
      64              : struct HeartbeatRequest<Server, State> {
      65              :     servers: Arc<HashMap<NodeId, Server>>,
      66              :     reply: tokio::sync::oneshot::Sender<Result<AvailablityDeltas<State>, HeartbeaterError>>,
      67              : }
      68              : 
      69              : pub(crate) struct Heartbeater<Server, State> {
      70              :     sender: tokio::sync::mpsc::UnboundedSender<HeartbeatRequest<Server, State>>,
      71              : }
      72              : 
      73              : #[allow(private_bounds)]
      74              : impl<Server: Send + Sync + 'static, State: Debug + Send + 'static> Heartbeater<Server, State>
      75              : where
      76              :     HeartbeaterTask<Server, State>: HeartBeat<Server, State>,
      77              : {
      78            0 :     pub(crate) fn new(
      79            0 :         jwt_token: Option<String>,
      80            0 :         ssl_ca_cert: Option<Certificate>,
      81            0 :         max_offline_interval: Duration,
      82            0 :         max_warming_up_interval: Duration,
      83            0 :         cancel: CancellationToken,
      84            0 :     ) -> Self {
      85            0 :         let (sender, receiver) =
      86            0 :             tokio::sync::mpsc::unbounded_channel::<HeartbeatRequest<Server, State>>();
      87            0 :         let mut heartbeater = HeartbeaterTask::new(
      88            0 :             receiver,
      89            0 :             jwt_token,
      90            0 :             ssl_ca_cert,
      91            0 :             max_offline_interval,
      92            0 :             max_warming_up_interval,
      93            0 :             cancel,
      94            0 :         );
      95            0 :         tokio::task::spawn(async move { heartbeater.run().await });
      96            0 : 
      97            0 :         Self { sender }
      98            0 :     }
      99              : 
     100            0 :     pub(crate) async fn heartbeat(
     101            0 :         &self,
     102            0 :         servers: Arc<HashMap<NodeId, Server>>,
     103            0 :     ) -> Result<AvailablityDeltas<State>, HeartbeaterError> {
     104            0 :         let (sender, receiver) = tokio::sync::oneshot::channel();
     105            0 :         self.sender
     106            0 :             .send(HeartbeatRequest {
     107            0 :                 servers,
     108            0 :                 reply: sender,
     109            0 :             })
     110            0 :             .map_err(|_| HeartbeaterError::Cancel)?;
     111              : 
     112            0 :         receiver
     113            0 :             .await
     114            0 :             .map_err(|_| HeartbeaterError::Cancel)
     115            0 :             .and_then(|x| x)
     116            0 :     }
     117              : }
     118              : 
     119              : impl<Server, State: Debug> HeartbeaterTask<Server, State>
     120              : where
     121              :     HeartbeaterTask<Server, State>: HeartBeat<Server, State>,
     122              : {
     123            0 :     fn new(
     124            0 :         receiver: tokio::sync::mpsc::UnboundedReceiver<HeartbeatRequest<Server, State>>,
     125            0 :         jwt_token: Option<String>,
     126            0 :         ssl_ca_cert: Option<Certificate>,
     127            0 :         max_offline_interval: Duration,
     128            0 :         max_warming_up_interval: Duration,
     129            0 :         cancel: CancellationToken,
     130            0 :     ) -> Self {
     131            0 :         Self {
     132            0 :             receiver,
     133            0 :             cancel,
     134            0 :             state: HashMap::new(),
     135            0 :             max_offline_interval,
     136            0 :             max_warming_up_interval,
     137            0 :             jwt_token,
     138            0 :             ssl_ca_cert,
     139            0 :         }
     140            0 :     }
     141            0 :     async fn run(&mut self) {
     142              :         loop {
     143            0 :             tokio::select! {
     144            0 :                 request = self.receiver.recv() => {
     145            0 :                     match request {
     146            0 :                         Some(req) => {
     147            0 :                             if req.reply.is_closed() {
     148              :                                 // Prevent a possibly infinite buildup of the receiver channel, if requests arrive faster than we can handle them
     149            0 :                                 continue;
     150            0 :                             }
     151            0 :                             let res = self.heartbeat(req.servers).await;
     152              :                             // Ignore the return value in order to not panic if the heartbeat function's future was cancelled
     153            0 :                             _ = req.reply.send(res);
     154              :                         },
     155            0 :                         None => { return; }
     156              :                     }
     157              :                 },
     158            0 :                 _ = self.cancel.cancelled() => return
     159              :             }
     160              :         }
     161            0 :     }
     162              : }
     163              : 
     164              : pub(crate) trait HeartBeat<Server, State> {
     165              :     fn heartbeat(
     166              :         &mut self,
     167              :         pageservers: Arc<HashMap<NodeId, Server>>,
     168              :     ) -> impl Future<Output = Result<AvailablityDeltas<State>, HeartbeaterError>> + Send;
     169              : }
     170              : 
     171              : impl HeartBeat<Node, PageserverState> for HeartbeaterTask<Node, PageserverState> {
     172            0 :     async fn heartbeat(
     173            0 :         &mut self,
     174            0 :         pageservers: Arc<HashMap<NodeId, Node>>,
     175            0 :     ) -> Result<AvailablityDeltas<PageserverState>, HeartbeaterError> {
     176            0 :         let mut new_state = HashMap::new();
     177            0 : 
     178            0 :         let mut heartbeat_futs = FuturesUnordered::new();
     179            0 :         for (node_id, node) in &*pageservers {
     180            0 :             heartbeat_futs.push({
     181            0 :                 let ssl_ca_cert = self.ssl_ca_cert.clone();
     182            0 :                 let jwt_token = self.jwt_token.clone();
     183            0 :                 let cancel = self.cancel.clone();
     184            0 : 
     185            0 :                 // Clone the node and mark it as available such that the request
     186            0 :                 // goes through to the pageserver even when the node is marked offline.
     187            0 :                 // This doesn't impact the availability observed by [`crate::service::Service`].
     188            0 :                 let mut node_clone = node.clone();
     189            0 :                 node_clone
     190            0 :                     .set_availability(NodeAvailability::Active(PageserverUtilization::full()));
     191            0 : 
     192            0 :                 async move {
     193            0 :                     let response = node_clone
     194            0 :                         .with_client_retries(
     195            0 :                             |client| async move { client.get_utilization().await },
     196            0 :                             &jwt_token,
     197            0 :                             &ssl_ca_cert,
     198            0 :                             3,
     199            0 :                             3,
     200            0 :                             Duration::from_secs(1),
     201            0 :                             &cancel,
     202            0 :                         )
     203            0 :                         .await;
     204              : 
     205            0 :                     let response = match response {
     206            0 :                         Some(r) => r,
     207              :                         None => {
     208              :                             // This indicates cancellation of the request.
     209              :                             // We ignore the node in this case.
     210            0 :                             return None;
     211              :                         }
     212              :                     };
     213              : 
     214            0 :                     let status = if let Ok(utilization) = response {
     215            0 :                         PageserverState::Available {
     216            0 :                             last_seen_at: Instant::now(),
     217            0 :                             utilization,
     218            0 :                         }
     219            0 :                     } else if let NodeAvailability::WarmingUp(last_seen_at) =
     220            0 :                         node.get_availability()
     221              :                     {
     222            0 :                         PageserverState::WarmingUp {
     223            0 :                             started_at: *last_seen_at,
     224            0 :                         }
     225              :                     } else {
     226            0 :                         PageserverState::Offline
     227              :                     };
     228              : 
     229            0 :                     Some((*node_id, status))
     230            0 :                 }
     231            0 :             });
     232            0 :         }
     233              : 
     234              :         loop {
     235            0 :             let maybe_status = tokio::select! {
     236            0 :                 next = heartbeat_futs.next() => {
     237            0 :                     match next {
     238            0 :                         Some(result) => result,
     239            0 :                         None => { break; }
     240              :                     }
     241              :                 },
     242            0 :                 _ = self.cancel.cancelled() => { return Err(HeartbeaterError::Cancel); }
     243              :             };
     244              : 
     245            0 :             if let Some((node_id, status)) = maybe_status {
     246            0 :                 new_state.insert(node_id, status);
     247            0 :             }
     248              :         }
     249              : 
     250            0 :         let mut warming_up = 0;
     251            0 :         let mut offline = 0;
     252            0 :         for state in new_state.values() {
     253            0 :             match state {
     254            0 :                 PageserverState::WarmingUp { .. } => {
     255            0 :                     warming_up += 1;
     256            0 :                 }
     257            0 :                 PageserverState::Offline { .. } => offline += 1,
     258            0 :                 PageserverState::Available { .. } => {}
     259              :             }
     260              :         }
     261              : 
     262            0 :         tracing::info!(
     263            0 :             "Heartbeat round complete for {} nodes, {} warming-up, {} offline",
     264            0 :             new_state.len(),
     265              :             warming_up,
     266              :             offline
     267              :         );
     268              : 
     269            0 :         let mut deltas = Vec::new();
     270            0 :         let now = Instant::now();
     271            0 :         for (node_id, ps_state) in new_state.iter_mut() {
     272              :             use std::collections::hash_map::Entry::*;
     273            0 :             let entry = self.state.entry(*node_id);
     274            0 : 
     275            0 :             let mut needs_update = false;
     276            0 :             match entry {
     277            0 :                 Occupied(ref occ) => match (occ.get(), &ps_state) {
     278            0 :                     (PageserverState::Offline, PageserverState::Offline) => {}
     279            0 :                     (PageserverState::Available { last_seen_at, .. }, PageserverState::Offline) => {
     280            0 :                         if now - *last_seen_at >= self.max_offline_interval {
     281            0 :                             deltas.push((*node_id, ps_state.clone()));
     282            0 :                             needs_update = true;
     283            0 :                         }
     284              :                     }
     285            0 :                     (_, PageserverState::WarmingUp { started_at }) => {
     286            0 :                         if now - *started_at >= self.max_warming_up_interval {
     287            0 :                             *ps_state = PageserverState::Offline;
     288            0 :                         }
     289              : 
     290            0 :                         deltas.push((*node_id, ps_state.clone()));
     291            0 :                         needs_update = true;
     292              :                     }
     293            0 :                     _ => {
     294            0 :                         deltas.push((*node_id, ps_state.clone()));
     295            0 :                         needs_update = true;
     296            0 :                     }
     297              :                 },
     298            0 :                 Vacant(_) => {
     299            0 :                     // This is a new node. Don't generate a delta for it.
     300            0 :                     deltas.push((*node_id, ps_state.clone()));
     301            0 :                 }
     302              :             }
     303              : 
     304            0 :             match entry {
     305            0 :                 Occupied(mut occ) if needs_update => {
     306            0 :                     (*occ.get_mut()) = ps_state.clone();
     307            0 :                 }
     308            0 :                 Vacant(vac) => {
     309            0 :                     vac.insert(ps_state.clone());
     310            0 :                 }
     311            0 :                 _ => {}
     312              :             }
     313              :         }
     314              : 
     315            0 :         Ok(AvailablityDeltas(deltas))
     316            0 :     }
     317              : }
     318              : 
     319              : impl HeartBeat<Safekeeper, SafekeeperState> for HeartbeaterTask<Safekeeper, SafekeeperState> {
     320            0 :     async fn heartbeat(
     321            0 :         &mut self,
     322            0 :         safekeepers: Arc<HashMap<NodeId, Safekeeper>>,
     323            0 :     ) -> Result<AvailablityDeltas<SafekeeperState>, HeartbeaterError> {
     324            0 :         let mut new_state = HashMap::new();
     325            0 : 
     326            0 :         let mut heartbeat_futs = FuturesUnordered::new();
     327            0 :         for (node_id, sk) in &*safekeepers {
     328            0 :             if sk.scheduling_policy() == SkSchedulingPolicy::Decomissioned {
     329            0 :                 continue;
     330            0 :             }
     331            0 :             heartbeat_futs.push({
     332            0 :                 let jwt_token = self
     333            0 :                     .jwt_token
     334            0 :                     .as_ref()
     335            0 :                     .map(|t| SecretString::from(t.to_owned()));
     336            0 :                 let ssl_ca_cert = self.ssl_ca_cert.clone();
     337            0 :                 let cancel = self.cancel.clone();
     338            0 : 
     339            0 :                 async move {
     340            0 :                     let response = sk
     341            0 :                         .with_client_retries(
     342            0 :                             |client| async move { client.get_utilization().await },
     343            0 :                             &jwt_token,
     344            0 :                             &ssl_ca_cert,
     345            0 :                             3,
     346            0 :                             3,
     347            0 :                             Duration::from_secs(1),
     348            0 :                             &cancel,
     349            0 :                         )
     350            0 :                         .await;
     351              : 
     352            0 :                     let status = match response {
     353            0 :                         Ok(utilization) => SafekeeperState::Available {
     354            0 :                             last_seen_at: Instant::now(),
     355            0 :                             utilization,
     356            0 :                         },
     357              :                         Err(mgmt_api::Error::Cancelled) => {
     358              :                             // This indicates cancellation of the request.
     359              :                             // We ignore the node in this case.
     360            0 :                             return None;
     361              :                         }
     362            0 :                         Err(e) => {
     363            0 :                             tracing::info!(
     364            0 :                                 "Marking safekeeper {} at as offline: {e}",
     365            0 :                                 sk.base_url()
     366              :                             );
     367            0 :                             SafekeeperState::Offline
     368              :                         }
     369              :                     };
     370              : 
     371            0 :                     Some((*node_id, status))
     372            0 :                 }
     373            0 :             });
     374            0 :         }
     375              : 
     376              :         loop {
     377            0 :             let maybe_status = tokio::select! {
     378            0 :                 next = heartbeat_futs.next() => {
     379            0 :                     match next {
     380            0 :                         Some(result) => result,
     381            0 :                         None => { break; }
     382              :                     }
     383              :                 },
     384            0 :                 _ = self.cancel.cancelled() => { return Err(HeartbeaterError::Cancel); }
     385              :             };
     386              : 
     387            0 :             if let Some((node_id, status)) = maybe_status {
     388            0 :                 new_state.insert(node_id, status);
     389            0 :             }
     390              :         }
     391              : 
     392            0 :         let mut offline = 0;
     393            0 :         for state in new_state.values() {
     394            0 :             match state {
     395            0 :                 SafekeeperState::Offline { .. } => offline += 1,
     396            0 :                 SafekeeperState::Available { .. } => {}
     397              :             }
     398              :         }
     399              : 
     400            0 :         tracing::info!(
     401            0 :             "Heartbeat round complete for {} safekeepers, {} offline",
     402            0 :             new_state.len(),
     403              :             offline
     404              :         );
     405              : 
     406            0 :         let mut deltas = Vec::new();
     407            0 :         let now = Instant::now();
     408            0 :         for (node_id, sk_state) in new_state.iter_mut() {
     409              :             use std::collections::hash_map::Entry::*;
     410            0 :             let entry = self.state.entry(*node_id);
     411            0 : 
     412            0 :             let mut needs_update = false;
     413            0 :             match entry {
     414            0 :                 Occupied(ref occ) => match (occ.get(), &sk_state) {
     415            0 :                     (SafekeeperState::Offline, SafekeeperState::Offline) => {}
     416            0 :                     (SafekeeperState::Available { last_seen_at, .. }, SafekeeperState::Offline) => {
     417            0 :                         if now - *last_seen_at >= self.max_offline_interval {
     418            0 :                             deltas.push((*node_id, sk_state.clone()));
     419            0 :                             needs_update = true;
     420            0 :                         }
     421              :                     }
     422            0 :                     _ => {
     423            0 :                         deltas.push((*node_id, sk_state.clone()));
     424            0 :                         needs_update = true;
     425            0 :                     }
     426              :                 },
     427            0 :                 Vacant(_) => {
     428            0 :                     // This is a new node. Don't generate a delta for it.
     429            0 :                     deltas.push((*node_id, sk_state.clone()));
     430            0 :                 }
     431              :             }
     432              : 
     433            0 :             match entry {
     434            0 :                 Occupied(mut occ) if needs_update => {
     435            0 :                     (*occ.get_mut()) = sk_state.clone();
     436            0 :                 }
     437            0 :                 Vacant(vac) => {
     438            0 :                     vac.insert(sk_state.clone());
     439            0 :                 }
     440            0 :                 _ => {}
     441              :             }
     442              :         }
     443              : 
     444            0 :         Ok(AvailablityDeltas(deltas))
     445            0 :     }
     446              : }
        

Generated by: LCOV version 2.1-beta