LCOV - code coverage report
Current view: top level - pageserver/src/tenant/timeline/walreceiver - connection_manager.rs (source / functions) Coverage Total Hit
Test: 07bee600374ccd486c69370d0972d9035964fe68.info Lines: 58.7 % 1193 700
Test Date: 2025-02-20 13:11:02 Functions: 61.8 % 76 47

            Line data    Source code
       1              : //! WAL receiver logic that ensures the pageserver gets connectected to safekeeper,
       2              : //! that contains the latest WAL to stream and this connection does not go stale.
       3              : //!
       4              : //! To achieve that, a storage broker is used: safekepers propagate their timelines' state in it,
       5              : //! the manager subscribes for changes and accumulates those to query the one with the biggest Lsn for connection.
       6              : //! Current connection state is tracked too, to ensure it's not getting stale.
       7              : //!
       8              : //! After every connection or storage broker update fetched, the state gets updated correspondingly and rechecked for the new conneciton leader,
       9              : //! then a (re)connection happens, if necessary.
      10              : //! Only WAL streaming task expects to be finished, other loops (storage broker, connection management) never exit unless cancelled explicitly via the dedicated channel.
      11              : 
      12              : use std::{collections::HashMap, num::NonZeroU64, ops::ControlFlow, sync::Arc, time::Duration};
      13              : 
      14              : use super::{TaskStateUpdate, WalReceiverConf};
      15              : use crate::context::{DownloadBehavior, RequestContext};
      16              : use crate::metrics::{
      17              :     WALRECEIVER_ACTIVE_MANAGERS, WALRECEIVER_BROKER_UPDATES, WALRECEIVER_CANDIDATES_ADDED,
      18              :     WALRECEIVER_CANDIDATES_REMOVED, WALRECEIVER_SWITCHES,
      19              : };
      20              : use crate::task_mgr::TaskKind;
      21              : use crate::tenant::{debug_assert_current_span_has_tenant_and_timeline_id, Timeline};
      22              : use anyhow::Context;
      23              : use chrono::{NaiveDateTime, Utc};
      24              : use pageserver_api::models::TimelineState;
      25              : 
      26              : use storage_broker::proto::TenantTimelineId as ProtoTenantTimelineId;
      27              : use storage_broker::proto::{
      28              :     FilterTenantTimelineId, MessageType, SafekeeperDiscoveryRequest, SafekeeperDiscoveryResponse,
      29              :     SubscribeByFilterRequest, TypeSubscription, TypedMessage,
      30              : };
      31              : use storage_broker::{BrokerClientChannel, Code, Streaming};
      32              : use tokio_util::sync::CancellationToken;
      33              : use tracing::*;
      34              : 
      35              : use postgres_connection::PgConnectionConfig;
      36              : use utils::backoff::{
      37              :     exponential_backoff, DEFAULT_BASE_BACKOFF_SECONDS, DEFAULT_MAX_BACKOFF_SECONDS,
      38              : };
      39              : use utils::postgres_client::{
      40              :     wal_stream_connection_config, ConnectionConfigArgs, PostgresClientProtocol,
      41              : };
      42              : use utils::{
      43              :     id::{NodeId, TenantTimelineId},
      44              :     lsn::Lsn,
      45              : };
      46              : 
      47              : use super::{
      48              :     walreceiver_connection::WalConnectionStatus, walreceiver_connection::WalReceiverError,
      49              :     TaskEvent, TaskHandle,
      50              : };
      51              : 
      52              : pub(crate) struct Cancelled;
      53              : 
      54              : /// Attempts to subscribe for timeline updates, pushed by safekeepers into the broker.
      55              : /// Based on the updates, desides whether to start, keep or stop a WAL receiver task.
      56              : /// If storage broker subscription is cancelled, exits.
      57              : ///
      58              : /// # Cancel-Safety
      59              : ///
      60              : /// Not cancellation-safe. Use `cancel` token to request cancellation.
      61            0 : pub(super) async fn connection_manager_loop_step(
      62            0 :     broker_client: &mut BrokerClientChannel,
      63            0 :     connection_manager_state: &mut ConnectionManagerState,
      64            0 :     ctx: &RequestContext,
      65            0 :     cancel: &CancellationToken,
      66            0 :     manager_status: &std::sync::RwLock<Option<ConnectionManagerStatus>>,
      67            0 : ) -> Result<(), Cancelled> {
      68            0 :     match tokio::select! {
      69            0 :         _ = cancel.cancelled() => { return Err(Cancelled); },
      70            0 :         st = connection_manager_state.timeline.wait_to_become_active(ctx) => { st }
      71              :     } {
      72            0 :         Ok(()) => {}
      73            0 :         Err(new_state) => {
      74            0 :             debug!(
      75              :                 ?new_state,
      76            0 :                 "state changed, stopping wal connection manager loop"
      77              :             );
      78            0 :             return Err(Cancelled);
      79              :         }
      80              :     }
      81              : 
      82            0 :     WALRECEIVER_ACTIVE_MANAGERS.inc();
      83            0 :     scopeguard::defer! {
      84            0 :         WALRECEIVER_ACTIVE_MANAGERS.dec();
      85            0 :     }
      86            0 : 
      87            0 :     let id = TenantTimelineId {
      88            0 :         tenant_id: connection_manager_state.timeline.tenant_shard_id.tenant_id,
      89            0 :         timeline_id: connection_manager_state.timeline.timeline_id,
      90            0 :     };
      91            0 : 
      92            0 :     let mut timeline_state_updates = connection_manager_state
      93            0 :         .timeline
      94            0 :         .subscribe_for_state_updates();
      95            0 : 
      96            0 :     let mut wait_lsn_status = connection_manager_state
      97            0 :         .timeline
      98            0 :         .subscribe_for_wait_lsn_updates();
      99            0 : 
     100            0 :     // TODO: create a separate config option for discovery request interval
     101            0 :     let discovery_request_interval = connection_manager_state.conf.lagging_wal_timeout;
     102            0 :     let mut last_discovery_ts: Option<std::time::Instant> = None;
     103              : 
     104              :     // Subscribe to the broker updates. Stream shares underlying TCP connection
     105              :     // with other streams on this client (other connection managers). When
     106              :     // object goes out of scope, stream finishes in drop() automatically.
     107            0 :     let mut broker_subscription = subscribe_for_timeline_updates(broker_client, id, cancel).await?;
     108            0 :     debug!("Subscribed for broker timeline updates");
     109              : 
     110              :     loop {
     111            0 :         let time_until_next_retry = connection_manager_state.time_until_next_retry();
     112            0 :         let any_activity = connection_manager_state.wal_connection.is_some()
     113            0 :             || !connection_manager_state.wal_stream_candidates.is_empty();
     114              : 
     115              :         // These things are happening concurrently:
     116              :         //
     117              :         //  - cancellation request
     118              :         //  - keep receiving WAL on the current connection
     119              :         //      - if the shared state says we need to change connection, disconnect and return
     120              :         //      - this runs in a separate task and we receive updates via a watch channel
     121              :         //  - change connection if the rules decide so, or if the current connection dies
     122              :         //  - receive updates from broker
     123              :         //      - this might change the current desired connection
     124              :         //  - timeline state changes to something that does not allow walreceiver to run concurrently
     125              :         //  - if there's no connection and no candidates, try to send a discovery request
     126              : 
     127              :         // NB: make sure each of the select expressions are cancellation-safe
     128              :         // (no need for arms to be cancellation-safe).
     129            0 :         tokio::select! {
     130            0 :             _ = cancel.cancelled() => { return Err(Cancelled); }
     131            0 :             Some(wal_connection_update) = async {
     132            0 :                 match connection_manager_state.wal_connection.as_mut() {
     133            0 :                     Some(wal_connection) => Some(wal_connection.connection_task.next_task_event().await),
     134            0 :                     None => None,
     135              :                 }
     136            0 :             } => {
     137            0 :                 let wal_connection = connection_manager_state.wal_connection.as_mut()
     138            0 :                     .expect("Should have a connection, as checked by the corresponding select! guard");
     139            0 :                 match wal_connection_update {
     140            0 :                     TaskEvent::Update(TaskStateUpdate::Started) => {},
     141            0 :                     TaskEvent::Update(TaskStateUpdate::Progress(new_status)) => {
     142            0 :                         if new_status.has_processed_wal {
     143            0 :                             // We have advanced last_record_lsn by processing the WAL received
     144            0 :                             // from this safekeeper. This is good enough to clean unsuccessful
     145            0 :                             // retries history and allow reconnecting to this safekeeper without
     146            0 :                             // sleeping for a long time.
     147            0 :                             connection_manager_state.wal_connection_retries.remove(&wal_connection.sk_id);
     148            0 :                         }
     149            0 :                         wal_connection.status = new_status;
     150              :                     }
     151            0 :                     TaskEvent::End(walreceiver_task_result) => {
     152            0 :                         match walreceiver_task_result {
     153            0 :                             Ok(()) => debug!("WAL receiving task finished"),
     154            0 :                             Err(e) => error!("wal receiver task finished with an error: {e:?}"),
     155              :                         }
     156            0 :                         connection_manager_state.drop_old_connection(false).await;
     157              :                     },
     158              :                 }
     159              :             },
     160              : 
     161              :             // Got a new update from the broker
     162            0 :             broker_update = broker_subscription.message() /* TODO: review cancellation-safety */ => {
     163            0 :                 match broker_update {
     164            0 :                     Ok(Some(broker_update)) => connection_manager_state.register_timeline_update(broker_update),
     165            0 :                     Err(status) => {
     166            0 :                         match status.code() {
     167            0 :                             Code::Unknown if status.message().contains("stream closed because of a broken pipe") || status.message().contains("connection reset") || status.message().contains("error reading a body from connection") => {
     168            0 :                                 // tonic's error handling doesn't provide a clear code for disconnections: we get
     169            0 :                                 // "h2 protocol error: error reading a body from connection: stream closed because of a broken pipe"
     170            0 :                                 // => https://github.com/neondatabase/neon/issues/9562
     171            0 :                                 info!("broker disconnected: {status}");
     172              :                             },
     173              :                             _ => {
     174            0 :                                 warn!("broker subscription failed: {status}");
     175              :                             }
     176              :                         }
     177            0 :                         return Ok(());
     178              :                     }
     179              :                     Ok(None) => {
     180            0 :                         error!("broker subscription stream ended"); // can't happen
     181            0 :                         return Ok(());
     182              :                     }
     183              :                 }
     184              :             },
     185              : 
     186            0 :             new_event = async {
     187              :                 // Reminder: this match arm needs to be cancellation-safe.
     188              :                 loop {
     189            0 :                     if connection_manager_state.timeline.current_state() == TimelineState::Loading {
     190            0 :                         warn!("wal connection manager should only be launched after timeline has become active");
     191            0 :                     }
     192            0 :                     match timeline_state_updates.changed().await {
     193              :                         Ok(()) => {
     194            0 :                             let new_state = connection_manager_state.timeline.current_state();
     195            0 :                             match new_state {
     196              :                                 // we're already active as walreceiver, no need to reactivate
     197            0 :                                 TimelineState::Active => continue,
     198              :                                 TimelineState::Broken { .. } | TimelineState::Stopping => {
     199            0 :                                     debug!("timeline entered terminal state {new_state:?}, stopping wal connection manager loop");
     200            0 :                                     return ControlFlow::Break(());
     201              :                                 }
     202              :                                 TimelineState::Loading => {
     203            0 :                                     warn!("timeline transitioned back to Loading state, that should not happen");
     204            0 :                                     return ControlFlow::Continue(());
     205              :                                 }
     206              :                             }
     207              :                         }
     208            0 :                         Err(_sender_dropped_error) => return ControlFlow::Break(()),
     209              :                     }
     210              :                 }
     211            0 :             } => match new_event {
     212              :                 ControlFlow::Continue(()) => {
     213            0 :                     return Ok(());
     214              :                 }
     215              :                 ControlFlow::Break(()) => {
     216            0 :                     debug!("Timeline is no longer active, stopping wal connection manager loop");
     217            0 :                     return Err(Cancelled);
     218              :                 }
     219              :             },
     220              : 
     221            0 :             Some(()) = async {
     222            0 :                 match time_until_next_retry {
     223            0 :                     Some(sleep_time) => {
     224            0 :                         tokio::time::sleep(sleep_time).await;
     225            0 :                         Some(())
     226              :                     },
     227              :                     None => {
     228            0 :                         debug!("No candidates to retry, waiting indefinitely for the broker events");
     229            0 :                         None
     230              :                     }
     231              :                 }
     232            0 :             } => debug!("Waking up for the next retry after waiting for {time_until_next_retry:?}"),
     233              : 
     234            0 :             Some(()) = async {
     235            0 :                 // Reminder: this match arm needs to be cancellation-safe.
     236            0 :                 // Calculating time needed to wait until sending the next discovery request.
     237            0 :                 // Current implementation is conservative and sends discovery requests only when there are no candidates.
     238            0 : 
     239            0 :                 if any_activity {
     240              :                     // No need to send discovery requests if there is an active connection or candidates.
     241            0 :                     return None;
     242            0 :                 }
     243              : 
     244              :                 // Waiting for an active wait_lsn request.
     245            0 :                 while wait_lsn_status.borrow().is_none() {
     246            0 :                     if wait_lsn_status.changed().await.is_err() {
     247              :                         // wait_lsn_status channel was closed, exiting
     248            0 :                         warn!("wait_lsn_status channel was closed in connection_manager_loop_step");
     249            0 :                         return None;
     250            0 :                     }
     251              :                 }
     252              : 
     253              :                 // All preconditions met, preparing to send a discovery request.
     254            0 :                 let now = std::time::Instant::now();
     255            0 :                 let next_discovery_ts = last_discovery_ts
     256            0 :                     .map(|ts| ts + discovery_request_interval)
     257            0 :                     .unwrap_or_else(|| now);
     258            0 : 
     259            0 :                 if next_discovery_ts > now {
     260              :                     // Prevent sending discovery requests too frequently.
     261            0 :                     tokio::time::sleep(next_discovery_ts - now).await;
     262            0 :                 }
     263              : 
     264            0 :                 let tenant_timeline_id = Some(ProtoTenantTimelineId {
     265            0 :                     tenant_id: id.tenant_id.as_ref().to_owned(),
     266            0 :                     timeline_id: id.timeline_id.as_ref().to_owned(),
     267            0 :                 });
     268            0 :                 let request = SafekeeperDiscoveryRequest { tenant_timeline_id };
     269            0 :                 let msg = TypedMessage {
     270            0 :                     r#type: MessageType::SafekeeperDiscoveryRequest as i32,
     271            0 :                     safekeeper_timeline_info: None,
     272            0 :                     safekeeper_discovery_request: Some(request),
     273            0 :                     safekeeper_discovery_response: None,
     274            0 :                     };
     275            0 : 
     276            0 :                 last_discovery_ts = Some(std::time::Instant::now());
     277            0 :                 info!("No active connection and no candidates, sending discovery request to the broker");
     278              : 
     279              :                 // Cancellation safety: we want to send a message to the broker, but publish_one()
     280              :                 // function can get cancelled by the other select! arm. This is absolutely fine, because
     281              :                 // we just want to receive broker updates and discovery is not important if we already
     282              :                 // receive updates.
     283              :                 //
     284              :                 // It is possible that `last_discovery_ts` will be updated, but the message will not be sent.
     285              :                 // This is totally fine because of the reason above.
     286              : 
     287              :                 // This is a fire-and-forget request, we don't care about the response
     288            0 :                 let _ = broker_client.publish_one(msg).await;
     289            0 :                 debug!("Discovery request sent to the broker");
     290            0 :                 None
     291            0 :             } => {}
     292              :         }
     293              : 
     294            0 :         if let Some(new_candidate) = connection_manager_state.next_connection_candidate() {
     295            0 :             info!("Switching to new connection candidate: {new_candidate:?}");
     296            0 :             connection_manager_state
     297            0 :                 .change_connection(new_candidate, ctx)
     298            0 :                 .await
     299            0 :         }
     300            0 :         *manager_status.write().unwrap() = Some(connection_manager_state.manager_status());
     301              :     }
     302            0 : }
     303              : 
     304              : /// Endlessly try to subscribe for broker updates for a given timeline.
     305            0 : async fn subscribe_for_timeline_updates(
     306            0 :     broker_client: &mut BrokerClientChannel,
     307            0 :     id: TenantTimelineId,
     308            0 :     cancel: &CancellationToken,
     309            0 : ) -> Result<Streaming<TypedMessage>, Cancelled> {
     310            0 :     let mut attempt = 0;
     311              :     loop {
     312            0 :         exponential_backoff(
     313            0 :             attempt,
     314            0 :             DEFAULT_BASE_BACKOFF_SECONDS,
     315            0 :             DEFAULT_MAX_BACKOFF_SECONDS,
     316            0 :             cancel,
     317            0 :         )
     318            0 :         .await;
     319            0 :         attempt += 1;
     320            0 : 
     321            0 :         // subscribe to the specific timeline
     322            0 :         let request = SubscribeByFilterRequest {
     323            0 :             types: vec![
     324            0 :                 TypeSubscription {
     325            0 :                     r#type: MessageType::SafekeeperTimelineInfo as i32,
     326            0 :                 },
     327            0 :                 TypeSubscription {
     328            0 :                     r#type: MessageType::SafekeeperDiscoveryResponse as i32,
     329            0 :                 },
     330            0 :             ],
     331            0 :             tenant_timeline_id: Some(FilterTenantTimelineId {
     332            0 :                 enabled: true,
     333            0 :                 tenant_timeline_id: Some(ProtoTenantTimelineId {
     334            0 :                     tenant_id: id.tenant_id.as_ref().to_owned(),
     335            0 :                     timeline_id: id.timeline_id.as_ref().to_owned(),
     336            0 :                 }),
     337            0 :             }),
     338            0 :         };
     339            0 : 
     340            0 :         match {
     341            0 :             tokio::select! {
     342            0 :                 r = broker_client.subscribe_by_filter(request) => { r }
     343            0 :                 _ = cancel.cancelled() => { return Err(Cancelled); }
     344              :             }
     345              :         } {
     346            0 :             Ok(resp) => {
     347            0 :                 return Ok(resp.into_inner());
     348              :             }
     349            0 :             Err(e) => {
     350            0 :                 // Safekeeper nodes can stop pushing timeline updates to the broker, when no new writes happen and
     351            0 :                 // entire WAL is streamed. Keep this noticeable with logging, but do not warn/error.
     352            0 :                 info!("Attempt #{attempt}, failed to subscribe for timeline {id} updates in broker: {e:#}");
     353            0 :                 continue;
     354              :             }
     355              :         }
     356              :     }
     357            0 : }
     358              : 
     359              : const WALCONNECTION_RETRY_MIN_BACKOFF_SECONDS: f64 = 0.1;
     360              : const WALCONNECTION_RETRY_MAX_BACKOFF_SECONDS: f64 = 15.0;
     361              : const WALCONNECTION_RETRY_BACKOFF_MULTIPLIER: f64 = 1.5;
     362              : 
     363              : /// All data that's needed to run endless broker loop and keep the WAL streaming connection alive, if possible.
     364              : pub(super) struct ConnectionManagerState {
     365              :     id: TenantTimelineId,
     366              :     /// Use pageserver data about the timeline to filter out some of the safekeepers.
     367              :     timeline: Arc<Timeline>,
     368              :     /// Child token of [`super::WalReceiver::cancel`], inherited to all tasks we spawn.
     369              :     cancel: CancellationToken,
     370              :     conf: WalReceiverConf,
     371              :     /// Current connection to safekeeper for WAL streaming.
     372              :     wal_connection: Option<WalConnection>,
     373              :     /// Info about retries and unsuccessful attempts to connect to safekeepers.
     374              :     wal_connection_retries: HashMap<NodeId, RetryInfo>,
     375              :     /// Data about all timelines, available for connection, fetched from storage broker, grouped by their corresponding safekeeper node id.
     376              :     wal_stream_candidates: HashMap<NodeId, BrokerSkTimeline>,
     377              : }
     378              : 
     379              : /// An information about connection manager's current connection and connection candidates.
     380              : #[derive(Debug, Clone)]
     381              : pub struct ConnectionManagerStatus {
     382              :     existing_connection: Option<WalConnectionStatus>,
     383              :     wal_stream_candidates: HashMap<NodeId, BrokerSkTimeline>,
     384              : }
     385              : 
     386              : impl ConnectionManagerStatus {
     387              :     /// Generates a string, describing current connection status in a form, suitable for logging.
     388            0 :     pub fn to_human_readable_string(&self) -> String {
     389            0 :         let mut resulting_string = String::new();
     390            0 :         match &self.existing_connection {
     391            0 :             Some(connection) => {
     392            0 :                 if connection.has_processed_wal {
     393            0 :                     resulting_string.push_str(&format!(
     394            0 :                         " (update {}): streaming WAL from node {}, ",
     395            0 :                         connection.latest_wal_update.format("%Y-%m-%d %H:%M:%S"),
     396            0 :                         connection.node,
     397            0 :                     ));
     398            0 : 
     399            0 :                     match (connection.streaming_lsn, connection.commit_lsn) {
     400            0 :                         (None, None) => resulting_string.push_str("no streaming data"),
     401            0 :                         (None, Some(commit_lsn)) => {
     402            0 :                             resulting_string.push_str(&format!("commit Lsn: {commit_lsn}"))
     403              :                         }
     404            0 :                         (Some(streaming_lsn), None) => {
     405            0 :                             resulting_string.push_str(&format!("streaming Lsn: {streaming_lsn}"))
     406              :                         }
     407            0 :                         (Some(streaming_lsn), Some(commit_lsn)) => resulting_string.push_str(
     408            0 :                             &format!("commit|streaming Lsn: {commit_lsn}|{streaming_lsn}"),
     409            0 :                         ),
     410              :                     }
     411            0 :                 } else if connection.is_connected {
     412            0 :                     resulting_string.push_str(&format!(
     413            0 :                         " (update {}): connecting to node {}",
     414            0 :                         connection
     415            0 :                             .latest_connection_update
     416            0 :                             .format("%Y-%m-%d %H:%M:%S"),
     417            0 :                         connection.node,
     418            0 :                     ));
     419            0 :                 } else {
     420            0 :                     resulting_string.push_str(&format!(
     421            0 :                         " (update {}): initializing node {} connection",
     422            0 :                         connection
     423            0 :                             .latest_connection_update
     424            0 :                             .format("%Y-%m-%d %H:%M:%S"),
     425            0 :                         connection.node,
     426            0 :                     ));
     427            0 :                 }
     428              :             }
     429            0 :             None => resulting_string.push_str(": disconnected"),
     430              :         }
     431              : 
     432            0 :         resulting_string.push_str(", safekeeper candidates (id|update_time|commit_lsn): [");
     433            0 :         let mut candidates = self.wal_stream_candidates.iter().peekable();
     434            0 :         while let Some((node_id, candidate_info)) = candidates.next() {
     435            0 :             resulting_string.push_str(&format!(
     436            0 :                 "({}|{}|{})",
     437            0 :                 node_id,
     438            0 :                 candidate_info.latest_update.format("%H:%M:%S"),
     439            0 :                 Lsn(candidate_info.timeline.commit_lsn)
     440            0 :             ));
     441            0 :             if candidates.peek().is_some() {
     442            0 :                 resulting_string.push_str(", ");
     443            0 :             }
     444              :         }
     445            0 :         resulting_string.push(']');
     446            0 : 
     447            0 :         resulting_string
     448            0 :     }
     449              : }
     450              : 
     451              : /// Current connection data.
     452              : #[derive(Debug)]
     453              : struct WalConnection {
     454              :     /// Time when the connection was initiated.
     455              :     started_at: NaiveDateTime,
     456              :     /// Current safekeeper pageserver is connected to for WAL streaming.
     457              :     sk_id: NodeId,
     458              :     /// Availability zone of the safekeeper.
     459              :     availability_zone: Option<String>,
     460              :     /// Status of the connection.
     461              :     status: WalConnectionStatus,
     462              :     /// WAL streaming task handle.
     463              :     connection_task: TaskHandle<WalConnectionStatus>,
     464              :     /// Have we discovered that other safekeeper has more recent WAL than we do?
     465              :     discovered_new_wal: Option<NewCommittedWAL>,
     466              : }
     467              : 
     468              : /// Notion of a new committed WAL, which exists on other safekeeper.
     469              : #[derive(Debug, Clone, Copy)]
     470              : struct NewCommittedWAL {
     471              :     /// LSN of the new committed WAL.
     472              :     lsn: Lsn,
     473              :     /// When we discovered that the new committed WAL exists on other safekeeper.
     474              :     discovered_at: NaiveDateTime,
     475              : }
     476              : 
     477              : #[derive(Debug, Clone, Copy)]
     478              : struct RetryInfo {
     479              :     next_retry_at: Option<NaiveDateTime>,
     480              :     retry_duration_seconds: f64,
     481              : }
     482              : 
     483              : /// Data about the timeline to connect to, received from the broker.
     484              : #[derive(Debug, Clone)]
     485              : struct BrokerSkTimeline {
     486              :     timeline: SafekeeperDiscoveryResponse,
     487              :     /// Time at which the data was fetched from the broker last time, to track the stale data.
     488              :     latest_update: NaiveDateTime,
     489              : }
     490              : 
     491              : impl ConnectionManagerState {
     492            0 :     pub(super) fn new(
     493            0 :         timeline: Arc<Timeline>,
     494            0 :         conf: WalReceiverConf,
     495            0 :         cancel: CancellationToken,
     496            0 :     ) -> Self {
     497            0 :         let id = TenantTimelineId {
     498            0 :             tenant_id: timeline.tenant_shard_id.tenant_id,
     499            0 :             timeline_id: timeline.timeline_id,
     500            0 :         };
     501            0 :         Self {
     502            0 :             id,
     503            0 :             timeline,
     504            0 :             cancel,
     505            0 :             conf,
     506            0 :             wal_connection: None,
     507            0 :             wal_stream_candidates: HashMap::new(),
     508            0 :             wal_connection_retries: HashMap::new(),
     509            0 :         }
     510            0 :     }
     511              : 
     512           20 :     fn spawn<Fut>(
     513           20 :         &self,
     514           20 :         task: impl FnOnce(
     515           20 :                 tokio::sync::watch::Sender<TaskStateUpdate<WalConnectionStatus>>,
     516           20 :                 CancellationToken,
     517           20 :             ) -> Fut
     518           20 :             + Send
     519           20 :             + 'static,
     520           20 :     ) -> TaskHandle<WalConnectionStatus>
     521           20 :     where
     522           20 :         Fut: std::future::Future<Output = anyhow::Result<()>> + Send,
     523           20 :     {
     524           20 :         // TODO: get rid of TaskHandle
     525           20 :         super::TaskHandle::spawn(&self.cancel, task)
     526           20 :     }
     527              : 
     528              :     /// Shuts down the current connection (if any) and immediately starts another one with the given connection string.
     529            0 :     async fn change_connection(&mut self, new_sk: NewWalConnectionCandidate, ctx: &RequestContext) {
     530            0 :         WALRECEIVER_SWITCHES
     531            0 :             .with_label_values(&[new_sk.reason.name()])
     532            0 :             .inc();
     533            0 : 
     534            0 :         self.drop_old_connection(true).await;
     535              : 
     536            0 :         let node_id = new_sk.safekeeper_id;
     537            0 :         let connect_timeout = self.conf.wal_connect_timeout;
     538            0 :         let ingest_batch_size = self.conf.ingest_batch_size;
     539            0 :         let protocol = self.conf.protocol;
     540            0 :         let timeline = Arc::clone(&self.timeline);
     541            0 :         let ctx = ctx.detached_child(
     542            0 :             TaskKind::WalReceiverConnectionHandler,
     543            0 :             DownloadBehavior::Download,
     544            0 :         );
     545              : 
     546            0 :         let span = info_span!("connection", %node_id);
     547            0 :         let connection_handle = self.spawn(move |events_sender, cancellation| {
     548            0 :             async move {
     549            0 :                 debug_assert_current_span_has_tenant_and_timeline_id();
     550              : 
     551            0 :                 let res = super::walreceiver_connection::handle_walreceiver_connection(
     552            0 :                     timeline,
     553            0 :                     protocol,
     554            0 :                     new_sk.wal_source_connconf,
     555            0 :                     events_sender,
     556            0 :                     cancellation.clone(),
     557            0 :                     connect_timeout,
     558            0 :                     ctx,
     559            0 :                     node_id,
     560            0 :                     ingest_batch_size,
     561            0 :                 )
     562            0 :                 .await;
     563              : 
     564            0 :                 match res {
     565            0 :                     Ok(()) => Ok(()),
     566            0 :                     Err(e) => {
     567            0 :                         match e {
     568            0 :                             WalReceiverError::SuccessfulCompletion(msg) => {
     569            0 :                                 info!("walreceiver connection handling ended with success: {msg}");
     570            0 :                                 Ok(())
     571              :                             }
     572            0 :                             WalReceiverError::ExpectedSafekeeperError(e) => {
     573            0 :                                 info!("walreceiver connection handling ended: {e}");
     574            0 :                                 Ok(())
     575              :                             }
     576              :                             WalReceiverError::ClosedGate => {
     577            0 :                                 info!(
     578            0 :                                     "walreceiver connection handling ended because of closed gate"
     579              :                                 );
     580            0 :                                 Ok(())
     581              :                             }
     582            0 :                             WalReceiverError::Other(e) => {
     583            0 :                                 // give out an error to have task_mgr give it a really verbose logging
     584            0 :                                 if cancellation.is_cancelled() {
     585              :                                     // Ideally we would learn about this via some path other than Other, but
     586              :                                     // that requires refactoring all the intermediate layers of ingest code
     587              :                                     // that only emit anyhow::Error
     588            0 :                                     Ok(())
     589              :                                 } else {
     590            0 :                                     Err(e).context("walreceiver connection handling failure")
     591              :                                 }
     592              :                             }
     593              :                         }
     594              :                     }
     595              :                 }
     596            0 :             }
     597            0 :             .instrument(span)
     598            0 :         });
     599            0 : 
     600            0 :         let now = Utc::now().naive_utc();
     601            0 :         self.wal_connection = Some(WalConnection {
     602            0 :             started_at: now,
     603            0 :             sk_id: new_sk.safekeeper_id,
     604            0 :             availability_zone: new_sk.availability_zone,
     605            0 :             status: WalConnectionStatus {
     606            0 :                 is_connected: false,
     607            0 :                 has_processed_wal: false,
     608            0 :                 latest_connection_update: now,
     609            0 :                 latest_wal_update: now,
     610            0 :                 streaming_lsn: None,
     611            0 :                 commit_lsn: None,
     612            0 :                 node: node_id,
     613            0 :             },
     614            0 :             connection_task: connection_handle,
     615            0 :             discovered_new_wal: None,
     616            0 :         });
     617            0 :     }
     618              : 
     619              :     /// Drops the current connection (if any) and updates retry timeout for the next
     620              :     /// connection attempt to the same safekeeper.
     621              :     ///
     622              :     /// # Cancel-Safety
     623              :     ///
     624              :     /// Not cancellation-safe.
     625            0 :     async fn drop_old_connection(&mut self, needs_shutdown: bool) {
     626            0 :         let wal_connection = match self.wal_connection.take() {
     627            0 :             Some(wal_connection) => wal_connection,
     628            0 :             None => return,
     629              :         };
     630              : 
     631            0 :         if needs_shutdown {
     632            0 :             wal_connection
     633            0 :                 .connection_task
     634            0 :                 .shutdown()
     635            0 :                 // This here is why this function isn't cancellation-safe.
     636            0 :                 // If we got cancelled here, then self.wal_connection is already None and we lose track of the task.
     637            0 :                 // Even if our caller diligently calls Self::shutdown(), it will find a self.wal_connection=None
     638            0 :                 // and thus be ineffective.
     639            0 :                 .await;
     640            0 :         }
     641              : 
     642            0 :         let retry = self
     643            0 :             .wal_connection_retries
     644            0 :             .entry(wal_connection.sk_id)
     645            0 :             .or_insert(RetryInfo {
     646            0 :                 next_retry_at: None,
     647            0 :                 retry_duration_seconds: WALCONNECTION_RETRY_MIN_BACKOFF_SECONDS,
     648            0 :             });
     649            0 : 
     650            0 :         let now = Utc::now().naive_utc();
     651            0 : 
     652            0 :         // Schedule the next retry attempt. We want to have exponential backoff for connection attempts,
     653            0 :         // and we add backoff to the time when we started the connection attempt. If the connection
     654            0 :         // was active for a long time, then next_retry_at will be in the past.
     655            0 :         retry.next_retry_at =
     656            0 :             wal_connection
     657            0 :                 .started_at
     658            0 :                 .checked_add_signed(chrono::Duration::milliseconds(
     659            0 :                     (retry.retry_duration_seconds * 1000.0) as i64,
     660            0 :                 ));
     661              : 
     662            0 :         if let Some(next) = &retry.next_retry_at {
     663            0 :             if next > &now {
     664            0 :                 info!(
     665            0 :                     "Next connection retry to {:?} is at {}",
     666              :                     wal_connection.sk_id, next
     667              :                 );
     668            0 :             }
     669            0 :         }
     670              : 
     671            0 :         let next_retry_duration =
     672            0 :             retry.retry_duration_seconds * WALCONNECTION_RETRY_BACKOFF_MULTIPLIER;
     673            0 :         // Clamp the next retry duration to the maximum allowed.
     674            0 :         let next_retry_duration = next_retry_duration.min(WALCONNECTION_RETRY_MAX_BACKOFF_SECONDS);
     675            0 :         // Clamp the next retry duration to the minimum allowed.
     676            0 :         let next_retry_duration = next_retry_duration.max(WALCONNECTION_RETRY_MIN_BACKOFF_SECONDS);
     677            0 : 
     678            0 :         retry.retry_duration_seconds = next_retry_duration;
     679            0 :     }
     680              : 
     681              :     /// Returns time needed to wait to have a new candidate for WAL streaming.
     682            0 :     fn time_until_next_retry(&self) -> Option<Duration> {
     683            0 :         let now = Utc::now().naive_utc();
     684              : 
     685            0 :         let next_retry_at = self
     686            0 :             .wal_connection_retries
     687            0 :             .values()
     688            0 :             .filter_map(|retry| retry.next_retry_at)
     689            0 :             .filter(|next_retry_at| next_retry_at > &now)
     690            0 :             .min()?;
     691              : 
     692            0 :         (next_retry_at - now).to_std().ok()
     693            0 :     }
     694              : 
     695              :     /// Adds another broker timeline into the state, if its more recent than the one already added there for the same key.
     696            0 :     fn register_timeline_update(&mut self, typed_msg: TypedMessage) {
     697            0 :         let mut is_discovery = false;
     698            0 :         let timeline_update = match typed_msg.r#type() {
     699              :             MessageType::SafekeeperTimelineInfo => {
     700            0 :                 let info = match typed_msg.safekeeper_timeline_info {
     701            0 :                     Some(info) => info,
     702              :                     None => {
     703            0 :                         warn!("bad proto message from broker: no safekeeper_timeline_info");
     704            0 :                         return;
     705              :                     }
     706              :                 };
     707            0 :                 SafekeeperDiscoveryResponse {
     708            0 :                     safekeeper_id: info.safekeeper_id,
     709            0 :                     tenant_timeline_id: info.tenant_timeline_id,
     710            0 :                     commit_lsn: info.commit_lsn,
     711            0 :                     safekeeper_connstr: info.safekeeper_connstr,
     712            0 :                     availability_zone: info.availability_zone,
     713            0 :                     standby_horizon: info.standby_horizon,
     714            0 :                 }
     715              :             }
     716              :             MessageType::SafekeeperDiscoveryResponse => {
     717            0 :                 is_discovery = true;
     718            0 :                 match typed_msg.safekeeper_discovery_response {
     719            0 :                     Some(response) => response,
     720              :                     None => {
     721            0 :                         warn!("bad proto message from broker: no safekeeper_discovery_response");
     722            0 :                         return;
     723              :                     }
     724              :                 }
     725              :             }
     726              :             _ => {
     727              :                 // unexpected message
     728            0 :                 return;
     729              :             }
     730              :         };
     731              : 
     732            0 :         WALRECEIVER_BROKER_UPDATES.inc();
     733            0 : 
     734            0 :         trace!(
     735            0 :             "safekeeper info update: standby_horizon(cutoff)={}",
     736              :             timeline_update.standby_horizon
     737              :         );
     738            0 :         if timeline_update.standby_horizon != 0 {
     739            0 :             // ignore reports from safekeepers not connected to replicas
     740            0 :             self.timeline
     741            0 :                 .standby_horizon
     742            0 :                 .store(Lsn(timeline_update.standby_horizon));
     743            0 :             self.timeline
     744            0 :                 .metrics
     745            0 :                 .standby_horizon_gauge
     746            0 :                 .set(timeline_update.standby_horizon as i64);
     747            0 :         }
     748              : 
     749            0 :         let new_safekeeper_id = NodeId(timeline_update.safekeeper_id);
     750            0 :         let old_entry = self.wal_stream_candidates.insert(
     751            0 :             new_safekeeper_id,
     752            0 :             BrokerSkTimeline {
     753            0 :                 timeline: timeline_update,
     754            0 :                 latest_update: Utc::now().naive_utc(),
     755            0 :             },
     756            0 :         );
     757            0 : 
     758            0 :         if old_entry.is_none() {
     759            0 :             info!(
     760              :                 ?is_discovery,
     761              :                 %new_safekeeper_id,
     762            0 :                 "New SK node was added",
     763              :             );
     764            0 :             WALRECEIVER_CANDIDATES_ADDED.inc();
     765            0 :         }
     766            0 :     }
     767              : 
     768              :     /// Cleans up stale broker records and checks the rest for the new connection candidate.
     769              :     /// Returns a new candidate, if the current state is absent or somewhat lagging, `None` otherwise.
     770              :     /// The current rules for approving new candidates:
     771              :     /// * pick a candidate different from the connected safekeeper with biggest `commit_lsn` and lowest failed connection attemps
     772              :     /// * if there's no such entry, no new candidate found, abort
     773              :     /// * otherwise check if the candidate is much better than the current one
     774              :     ///
     775              :     /// To understand exact rules for determining if the candidate is better than the current one, refer to this function's implementation.
     776              :     /// General rules are following:
     777              :     /// * if connected safekeeper is not present, pick the candidate
     778              :     /// * if we haven't received any updates for some time, pick the candidate
     779              :     /// * if the candidate commit_lsn is much higher than the current one, pick the candidate
     780              :     /// * if the candidate commit_lsn is same, but candidate is located in the same AZ as the pageserver, pick the candidate
     781              :     /// * if connected safekeeper stopped sending us new WAL which is available on other safekeeper, pick the candidate
     782              :     ///
     783              :     /// This way we ensure to keep up with the most up-to-date safekeeper and don't try to jump from one safekeeper to another too frequently.
     784              :     /// Both thresholds are configured per tenant.
     785           36 :     fn next_connection_candidate(&mut self) -> Option<NewWalConnectionCandidate> {
     786           36 :         self.cleanup_old_candidates();
     787           36 : 
     788           36 :         match &self.wal_connection {
     789           20 :             Some(existing_wal_connection) => {
     790           20 :                 let connected_sk_node = existing_wal_connection.sk_id;
     791              : 
     792           20 :                 let (new_sk_id, new_safekeeper_broker_data, new_wal_source_connconf) =
     793           20 :                     self.select_connection_candidate(Some(connected_sk_node))?;
     794           20 :                 let new_availability_zone = new_safekeeper_broker_data.availability_zone.clone();
     795           20 : 
     796           20 :                 let now = Utc::now().naive_utc();
     797           20 :                 if let Ok(latest_interaciton) =
     798           20 :                     (now - existing_wal_connection.status.latest_connection_update).to_std()
     799              :                 {
     800              :                     // Drop connection if we haven't received keepalive message for a while.
     801           20 :                     if latest_interaciton > self.conf.wal_connect_timeout {
     802            4 :                         return Some(NewWalConnectionCandidate {
     803            4 :                             safekeeper_id: new_sk_id,
     804            4 :                             wal_source_connconf: new_wal_source_connconf,
     805            4 :                             availability_zone: new_availability_zone,
     806            4 :                             reason: ReconnectReason::NoKeepAlives {
     807            4 :                                 last_keep_alive: Some(
     808            4 :                                     existing_wal_connection.status.latest_connection_update,
     809            4 :                                 ),
     810            4 :                                 check_time: now,
     811            4 :                                 threshold: self.conf.wal_connect_timeout,
     812            4 :                             },
     813            4 :                         });
     814           16 :                     }
     815            0 :                 }
     816              : 
     817           16 :                 if !existing_wal_connection.status.is_connected {
     818              :                     // We haven't connected yet and we shouldn't switch until connection timeout (condition above).
     819            0 :                     return None;
     820           16 :                 }
     821              : 
     822           16 :                 if let Some(current_commit_lsn) = existing_wal_connection.status.commit_lsn {
     823           16 :                     let new_commit_lsn = Lsn(new_safekeeper_broker_data.commit_lsn);
     824           16 :                     // Check if the new candidate has much more WAL than the current one.
     825           16 :                     match new_commit_lsn.0.checked_sub(current_commit_lsn.0) {
     826           16 :                         Some(new_sk_lsn_advantage) => {
     827           16 :                             if new_sk_lsn_advantage >= self.conf.max_lsn_wal_lag.get() {
     828            4 :                                 return Some(NewWalConnectionCandidate {
     829            4 :                                     safekeeper_id: new_sk_id,
     830            4 :                                     wal_source_connconf: new_wal_source_connconf,
     831            4 :                                     availability_zone: new_availability_zone,
     832            4 :                                     reason: ReconnectReason::LaggingWal {
     833            4 :                                         current_commit_lsn,
     834            4 :                                         new_commit_lsn,
     835            4 :                                         threshold: self.conf.max_lsn_wal_lag,
     836            4 :                                     },
     837            4 :                                 });
     838           12 :                             }
     839           12 :                             // If we have a candidate with the same commit_lsn as the current one, which is in the same AZ as pageserver,
     840           12 :                             // and the current one is not, switch to the new one.
     841           12 :                             if self.conf.availability_zone.is_some()
     842            4 :                                 && existing_wal_connection.availability_zone
     843            4 :                                     != self.conf.availability_zone
     844            4 :                                 && self.conf.availability_zone == new_availability_zone
     845              :                             {
     846            4 :                                 return Some(NewWalConnectionCandidate {
     847            4 :                                     safekeeper_id: new_sk_id,
     848            4 :                                     availability_zone: new_availability_zone,
     849            4 :                                     wal_source_connconf: new_wal_source_connconf,
     850            4 :                                     reason: ReconnectReason::SwitchAvailabilityZone,
     851            4 :                                 });
     852            8 :                             }
     853              :                         }
     854            0 :                         None => debug!(
     855            0 :                             "Best SK candidate has its commit_lsn behind connected SK's commit_lsn"
     856              :                         ),
     857              :                     }
     858            0 :                 }
     859              : 
     860            8 :                 let current_lsn = match existing_wal_connection.status.streaming_lsn {
     861            8 :                     Some(lsn) => lsn,
     862            0 :                     None => self.timeline.get_last_record_lsn(),
     863              :                 };
     864            8 :                 let current_commit_lsn = existing_wal_connection
     865            8 :                     .status
     866            8 :                     .commit_lsn
     867            8 :                     .unwrap_or(current_lsn);
     868            8 :                 let candidate_commit_lsn = Lsn(new_safekeeper_broker_data.commit_lsn);
     869            8 : 
     870            8 :                 // Keep discovered_new_wal only if connected safekeeper has not caught up yet.
     871            8 :                 let mut discovered_new_wal = existing_wal_connection
     872            8 :                     .discovered_new_wal
     873            8 :                     .filter(|new_wal| new_wal.lsn > current_commit_lsn);
     874            8 : 
     875            8 :                 if discovered_new_wal.is_none() {
     876              :                     // Check if the new candidate has more WAL than the current one.
     877              :                     // If the new candidate has more WAL than the current one, we consider switching to the new candidate.
     878            4 :                     discovered_new_wal = if candidate_commit_lsn > current_commit_lsn {
     879            4 :                         trace!(
     880            0 :                             "New candidate has commit_lsn {}, higher than current_commit_lsn {}",
     881              :                             candidate_commit_lsn,
     882              :                             current_commit_lsn
     883              :                         );
     884            4 :                         Some(NewCommittedWAL {
     885            4 :                             lsn: candidate_commit_lsn,
     886            4 :                             discovered_at: Utc::now().naive_utc(),
     887            4 :                         })
     888              :                     } else {
     889            0 :                         None
     890              :                     };
     891            4 :                 }
     892              : 
     893            8 :                 let waiting_for_new_lsn_since = if current_lsn < current_commit_lsn {
     894              :                     // Connected safekeeper has more WAL, but we haven't received updates for some time.
     895            0 :                     trace!(
     896            0 :                         "Connected safekeeper has more WAL, but we haven't received updates for {:?}. current_lsn: {}, current_commit_lsn: {}",
     897            0 :                         (now - existing_wal_connection.status.latest_wal_update).to_std(),
     898              :                         current_lsn,
     899              :                         current_commit_lsn
     900              :                     );
     901            0 :                     Some(existing_wal_connection.status.latest_wal_update)
     902              :                 } else {
     903            8 :                     discovered_new_wal.as_ref().map(|new_wal| {
     904            8 :                         // We know that new WAL is available on other safekeeper, but connected safekeeper don't have it.
     905            8 :                         new_wal
     906            8 :                             .discovered_at
     907            8 :                             .max(existing_wal_connection.status.latest_wal_update)
     908            8 :                     })
     909              :                 };
     910              : 
     911              :                 // If we haven't received any WAL updates for a while and candidate has more WAL, switch to it.
     912            8 :                 if let Some(waiting_for_new_lsn_since) = waiting_for_new_lsn_since {
     913            8 :                     if let Ok(waiting_for_new_wal) = (now - waiting_for_new_lsn_since).to_std() {
     914            4 :                         if candidate_commit_lsn > current_commit_lsn
     915            4 :                             && waiting_for_new_wal > self.conf.lagging_wal_timeout
     916              :                         {
     917            4 :                             return Some(NewWalConnectionCandidate {
     918            4 :                                 safekeeper_id: new_sk_id,
     919            4 :                                 wal_source_connconf: new_wal_source_connconf,
     920            4 :                                 availability_zone: new_availability_zone,
     921            4 :                                 reason: ReconnectReason::NoWalTimeout {
     922            4 :                                     current_lsn,
     923            4 :                                     current_commit_lsn,
     924            4 :                                     candidate_commit_lsn,
     925            4 :                                     last_wal_interaction: Some(
     926            4 :                                         existing_wal_connection.status.latest_wal_update,
     927            4 :                                     ),
     928            4 :                                     check_time: now,
     929            4 :                                     threshold: self.conf.lagging_wal_timeout,
     930            4 :                                 },
     931            4 :                             });
     932            0 :                         }
     933            4 :                     }
     934            0 :                 }
     935              : 
     936            4 :                 self.wal_connection.as_mut().unwrap().discovered_new_wal = discovered_new_wal;
     937              :             }
     938              :             None => {
     939           12 :                 let (new_sk_id, new_safekeeper_broker_data, new_wal_source_connconf) =
     940           16 :                     self.select_connection_candidate(None)?;
     941           12 :                 return Some(NewWalConnectionCandidate {
     942           12 :                     safekeeper_id: new_sk_id,
     943           12 :                     availability_zone: new_safekeeper_broker_data.availability_zone.clone(),
     944           12 :                     wal_source_connconf: new_wal_source_connconf,
     945           12 :                     reason: ReconnectReason::NoExistingConnection,
     946           12 :                 });
     947              :             }
     948              :         }
     949              : 
     950            4 :         None
     951           36 :     }
     952              : 
     953              :     /// Selects the best possible candidate, based on the data collected from the broker updates about the safekeepers.
     954              :     /// Optionally, omits the given node, to support gracefully switching from a healthy safekeeper to another.
     955              :     ///
     956              :     /// The candidate that is chosen:
     957              :     /// * has no pending retry cooldown
     958              :     /// * has greatest commit_lsn among the ones that are left
     959           36 :     fn select_connection_candidate(
     960           36 :         &self,
     961           36 :         node_to_omit: Option<NodeId>,
     962           36 :     ) -> Option<(NodeId, &SafekeeperDiscoveryResponse, PgConnectionConfig)> {
     963           36 :         self.applicable_connection_candidates()
     964           52 :             .filter(|&(sk_id, _, _)| Some(sk_id) != node_to_omit)
     965           40 :             .max_by_key(|(_, info, _)| info.commit_lsn)
     966           36 :     }
     967              : 
     968              :     /// Returns a list of safekeepers that have valid info and ready for connection.
     969              :     /// Some safekeepers are filtered by the retry cooldown.
     970           36 :     fn applicable_connection_candidates(
     971           36 :         &self,
     972           36 :     ) -> impl Iterator<Item = (NodeId, &SafekeeperDiscoveryResponse, PgConnectionConfig)> {
     973           36 :         let now = Utc::now().naive_utc();
     974           36 : 
     975           36 :         self.wal_stream_candidates
     976           36 :             .iter()
     977           72 :             .filter(|(_, info)| Lsn(info.timeline.commit_lsn) != Lsn::INVALID)
     978           64 :             .filter(move |(sk_id, _)| {
     979           64 :                 let next_retry_at = self
     980           64 :                     .wal_connection_retries
     981           64 :                     .get(sk_id)
     982           64 :                     .and_then(|retry_info| {
     983            4 :                         retry_info.next_retry_at
     984           64 :                     });
     985           64 : 
     986           64 :                 next_retry_at.is_none() || next_retry_at.unwrap() <= now
     987           64 :             }).filter_map(|(sk_id, broker_info)| {
     988           60 :                 let info = &broker_info.timeline;
     989           60 :                 if info.safekeeper_connstr.is_empty() {
     990            8 :                     return None; // no connection string, ignore sk
     991           52 :                 }
     992              : 
     993           52 :                 let (shard_number, shard_count, shard_stripe_size) = match self.conf.protocol {
     994              :                     PostgresClientProtocol::Vanilla => {
     995           52 :                         (None, None, None)
     996              :                     },
     997              :                     PostgresClientProtocol::Interpreted { .. } => {
     998            0 :                         let shard_identity = self.timeline.get_shard_identity();
     999            0 :                         (
    1000            0 :                             Some(shard_identity.number.0),
    1001            0 :                             Some(shard_identity.count.0),
    1002            0 :                             Some(shard_identity.stripe_size.0),
    1003            0 :                         )
    1004              :                     }
    1005              :                 };
    1006              : 
    1007           52 :                 let connection_conf_args = ConnectionConfigArgs {
    1008           52 :                     protocol: self.conf.protocol,
    1009           52 :                     ttid: self.id,
    1010           52 :                     shard_number,
    1011           52 :                     shard_count,
    1012           52 :                     shard_stripe_size,
    1013           52 :                     listen_pg_addr_str: info.safekeeper_connstr.as_ref(),
    1014           52 :                     auth_token: self.conf.auth_token.as_ref().map(|t| t.as_str()),
    1015           52 :                     availability_zone: self.conf.availability_zone.as_deref()
    1016           52 :                 };
    1017           52 : 
    1018           52 :                 match wal_stream_connection_config(connection_conf_args) {
    1019           52 :                     Ok(connstr) => Some((*sk_id, info, connstr)),
    1020            0 :                     Err(e) => {
    1021            0 :                         error!("Failed to create wal receiver connection string from broker data of safekeeper node {}: {e:#}", sk_id);
    1022            0 :                         None
    1023              :                     }
    1024              :                 }
    1025           60 :             })
    1026           36 :     }
    1027              : 
    1028              :     /// Remove candidates which haven't sent broker updates for a while.
    1029           36 :     fn cleanup_old_candidates(&mut self) {
    1030           36 :         let mut node_ids_to_remove = Vec::with_capacity(self.wal_stream_candidates.len());
    1031           36 :         let lagging_wal_timeout = self.conf.lagging_wal_timeout;
    1032           36 : 
    1033           76 :         self.wal_stream_candidates.retain(|node_id, broker_info| {
    1034           76 :             if let Ok(time_since_latest_broker_update) =
    1035           76 :                 (Utc::now().naive_utc() - broker_info.latest_update).to_std()
    1036              :             {
    1037           76 :                 let should_retain = time_since_latest_broker_update < lagging_wal_timeout;
    1038           76 :                 if !should_retain {
    1039            4 :                     node_ids_to_remove.push(*node_id);
    1040           72 :                 }
    1041           76 :                 should_retain
    1042              :             } else {
    1043            0 :                 true
    1044              :             }
    1045           76 :         });
    1046           36 : 
    1047           36 :         if !node_ids_to_remove.is_empty() {
    1048            8 :             for node_id in node_ids_to_remove {
    1049            4 :                 info!("Safekeeper node {node_id} did not send events for over {lagging_wal_timeout:?}, not retrying the connections");
    1050            4 :                 self.wal_connection_retries.remove(&node_id);
    1051            4 :                 WALRECEIVER_CANDIDATES_REMOVED.inc();
    1052              :             }
    1053           32 :         }
    1054           36 :     }
    1055              : 
    1056              :     /// # Cancel-Safety
    1057              :     ///
    1058              :     /// Not cancellation-safe.
    1059            0 :     pub(super) async fn shutdown(mut self) {
    1060            0 :         if let Some(wal_connection) = self.wal_connection.take() {
    1061            0 :             wal_connection.connection_task.shutdown().await;
    1062            0 :         }
    1063            0 :     }
    1064              : 
    1065            0 :     fn manager_status(&self) -> ConnectionManagerStatus {
    1066            0 :         ConnectionManagerStatus {
    1067            0 :             existing_connection: self.wal_connection.as_ref().map(|conn| conn.status),
    1068            0 :             wal_stream_candidates: self.wal_stream_candidates.clone(),
    1069            0 :         }
    1070            0 :     }
    1071              : }
    1072              : 
    1073              : #[derive(Debug)]
    1074              : struct NewWalConnectionCandidate {
    1075              :     safekeeper_id: NodeId,
    1076              :     wal_source_connconf: PgConnectionConfig,
    1077              :     availability_zone: Option<String>,
    1078              :     reason: ReconnectReason,
    1079              : }
    1080              : 
    1081              : /// Stores the reason why WAL connection was switched, for furter debugging purposes.
    1082              : #[derive(Debug, PartialEq, Eq)]
    1083              : enum ReconnectReason {
    1084              :     NoExistingConnection,
    1085              :     LaggingWal {
    1086              :         current_commit_lsn: Lsn,
    1087              :         new_commit_lsn: Lsn,
    1088              :         threshold: NonZeroU64,
    1089              :     },
    1090              :     SwitchAvailabilityZone,
    1091              :     NoWalTimeout {
    1092              :         current_lsn: Lsn,
    1093              :         current_commit_lsn: Lsn,
    1094              :         candidate_commit_lsn: Lsn,
    1095              :         last_wal_interaction: Option<NaiveDateTime>,
    1096              :         check_time: NaiveDateTime,
    1097              :         threshold: Duration,
    1098              :     },
    1099              :     NoKeepAlives {
    1100              :         last_keep_alive: Option<NaiveDateTime>,
    1101              :         check_time: NaiveDateTime,
    1102              :         threshold: Duration,
    1103              :     },
    1104              : }
    1105              : 
    1106              : impl ReconnectReason {
    1107            0 :     fn name(&self) -> &str {
    1108            0 :         match self {
    1109            0 :             ReconnectReason::NoExistingConnection => "NoExistingConnection",
    1110            0 :             ReconnectReason::LaggingWal { .. } => "LaggingWal",
    1111            0 :             ReconnectReason::SwitchAvailabilityZone => "SwitchAvailabilityZone",
    1112            0 :             ReconnectReason::NoWalTimeout { .. } => "NoWalTimeout",
    1113            0 :             ReconnectReason::NoKeepAlives { .. } => "NoKeepAlives",
    1114              :         }
    1115            0 :     }
    1116              : }
    1117              : 
    1118              : #[cfg(test)]
    1119              : mod tests {
    1120              :     use super::*;
    1121              :     use crate::tenant::harness::{TenantHarness, TIMELINE_ID};
    1122              :     use pageserver_api::config::defaults::DEFAULT_WAL_RECEIVER_PROTOCOL;
    1123              :     use url::Host;
    1124              : 
    1125           76 :     fn dummy_broker_sk_timeline(
    1126           76 :         commit_lsn: u64,
    1127           76 :         safekeeper_connstr: &str,
    1128           76 :         latest_update: NaiveDateTime,
    1129           76 :     ) -> BrokerSkTimeline {
    1130           76 :         BrokerSkTimeline {
    1131           76 :             timeline: SafekeeperDiscoveryResponse {
    1132           76 :                 safekeeper_id: 0,
    1133           76 :                 tenant_timeline_id: None,
    1134           76 :                 commit_lsn,
    1135           76 :                 safekeeper_connstr: safekeeper_connstr.to_owned(),
    1136           76 :                 availability_zone: None,
    1137           76 :                 standby_horizon: 0,
    1138           76 :             },
    1139           76 :             latest_update,
    1140           76 :         }
    1141           76 :     }
    1142              : 
    1143              :     #[tokio::test]
    1144            4 :     async fn no_connection_no_candidate() -> anyhow::Result<()> {
    1145            4 :         let harness = TenantHarness::create("no_connection_no_candidate").await?;
    1146            4 :         let mut state = dummy_state(&harness).await;
    1147            4 :         let now = Utc::now().naive_utc();
    1148            4 : 
    1149            4 :         let lagging_wal_timeout = chrono::Duration::from_std(state.conf.lagging_wal_timeout)?;
    1150            4 :         let delay_over_threshold = now - lagging_wal_timeout - lagging_wal_timeout;
    1151            4 : 
    1152            4 :         state.wal_connection = None;
    1153            4 :         state.wal_stream_candidates = HashMap::from([
    1154            4 :             (NodeId(0), dummy_broker_sk_timeline(1, "", now)),
    1155            4 :             (NodeId(1), dummy_broker_sk_timeline(0, "no_commit_lsn", now)),
    1156            4 :             (NodeId(2), dummy_broker_sk_timeline(0, "no_commit_lsn", now)),
    1157            4 :             (
    1158            4 :                 NodeId(3),
    1159            4 :                 dummy_broker_sk_timeline(
    1160            4 :                     1 + state.conf.max_lsn_wal_lag.get(),
    1161            4 :                     "delay_over_threshold",
    1162            4 :                     delay_over_threshold,
    1163            4 :                 ),
    1164            4 :             ),
    1165            4 :         ]);
    1166            4 : 
    1167            4 :         let no_candidate = state.next_connection_candidate();
    1168            4 :         assert!(
    1169            4 :             no_candidate.is_none(),
    1170            4 :             "Expected no candidate selected out of non full data options, but got {no_candidate:?}"
    1171            4 :         );
    1172            4 : 
    1173            4 :         Ok(())
    1174            4 :     }
    1175              : 
    1176              :     #[tokio::test]
    1177            4 :     async fn connection_no_candidate() -> anyhow::Result<()> {
    1178            4 :         let harness = TenantHarness::create("connection_no_candidate").await?;
    1179            4 :         let mut state = dummy_state(&harness).await;
    1180            4 :         let now = Utc::now().naive_utc();
    1181            4 : 
    1182            4 :         let connected_sk_id = NodeId(0);
    1183            4 :         let current_lsn = 100_000;
    1184            4 : 
    1185            4 :         let connection_status = WalConnectionStatus {
    1186            4 :             is_connected: true,
    1187            4 :             has_processed_wal: true,
    1188            4 :             latest_connection_update: now,
    1189            4 :             latest_wal_update: now,
    1190            4 :             commit_lsn: Some(Lsn(current_lsn)),
    1191            4 :             streaming_lsn: Some(Lsn(current_lsn)),
    1192            4 :             node: NodeId(1),
    1193            4 :         };
    1194            4 : 
    1195            4 :         state.conf.max_lsn_wal_lag = NonZeroU64::new(100).unwrap();
    1196            4 :         state.wal_connection = Some(WalConnection {
    1197            4 :             started_at: now,
    1198            4 :             sk_id: connected_sk_id,
    1199            4 :             availability_zone: None,
    1200            4 :             status: connection_status,
    1201            4 :             connection_task: state.spawn(move |sender, _| async move {
    1202            4 :                 sender
    1203            4 :                     .send(TaskStateUpdate::Progress(connection_status))
    1204            4 :                     .ok();
    1205            4 :                 Ok(())
    1206            4 :             }),
    1207            4 :             discovered_new_wal: None,
    1208            4 :         });
    1209            4 :         state.wal_stream_candidates = HashMap::from([
    1210            4 :             (
    1211            4 :                 connected_sk_id,
    1212            4 :                 dummy_broker_sk_timeline(
    1213            4 :                     current_lsn + state.conf.max_lsn_wal_lag.get() * 2,
    1214            4 :                     DUMMY_SAFEKEEPER_HOST,
    1215            4 :                     now,
    1216            4 :                 ),
    1217            4 :             ),
    1218            4 :             (
    1219            4 :                 NodeId(1),
    1220            4 :                 dummy_broker_sk_timeline(current_lsn, "not_advanced_lsn", now),
    1221            4 :             ),
    1222            4 :             (
    1223            4 :                 NodeId(2),
    1224            4 :                 dummy_broker_sk_timeline(
    1225            4 :                     current_lsn + state.conf.max_lsn_wal_lag.get() / 2,
    1226            4 :                     "not_enough_advanced_lsn",
    1227            4 :                     now,
    1228            4 :                 ),
    1229            4 :             ),
    1230            4 :         ]);
    1231            4 : 
    1232            4 :         let no_candidate = state.next_connection_candidate();
    1233            4 :         assert!(
    1234            4 :             no_candidate.is_none(),
    1235            4 :             "Expected no candidate selected out of valid options since candidate Lsn data is ignored and others' was not advanced enough, but got {no_candidate:?}"
    1236            4 :         );
    1237            4 : 
    1238            4 :         Ok(())
    1239            4 :     }
    1240              : 
    1241              :     #[tokio::test]
    1242            4 :     async fn no_connection_candidate() -> anyhow::Result<()> {
    1243            4 :         let harness = TenantHarness::create("no_connection_candidate").await?;
    1244            4 :         let mut state = dummy_state(&harness).await;
    1245            4 :         let now = Utc::now().naive_utc();
    1246            4 : 
    1247            4 :         state.wal_connection = None;
    1248            4 :         state.wal_stream_candidates = HashMap::from([(
    1249            4 :             NodeId(0),
    1250            4 :             dummy_broker_sk_timeline(
    1251            4 :                 1 + state.conf.max_lsn_wal_lag.get(),
    1252            4 :                 DUMMY_SAFEKEEPER_HOST,
    1253            4 :                 now,
    1254            4 :             ),
    1255            4 :         )]);
    1256            4 : 
    1257            4 :         let only_candidate = state
    1258            4 :             .next_connection_candidate()
    1259            4 :             .expect("Expected one candidate selected out of the only data option, but got none");
    1260            4 :         assert_eq!(only_candidate.safekeeper_id, NodeId(0));
    1261            4 :         assert_eq!(
    1262            4 :             only_candidate.reason,
    1263            4 :             ReconnectReason::NoExistingConnection,
    1264            4 :             "Should select new safekeeper due to missing connection, even if there's also a lag in the wal over the threshold"
    1265            4 :         );
    1266            4 :         assert_eq!(
    1267            4 :             only_candidate.wal_source_connconf.host(),
    1268            4 :             &Host::Domain(DUMMY_SAFEKEEPER_HOST.to_owned())
    1269            4 :         );
    1270            4 : 
    1271            4 :         let selected_lsn = 100_000;
    1272            4 :         state.wal_stream_candidates = HashMap::from([
    1273            4 :             (
    1274            4 :                 NodeId(0),
    1275            4 :                 dummy_broker_sk_timeline(selected_lsn - 100, "smaller_commit_lsn", now),
    1276            4 :             ),
    1277            4 :             (
    1278            4 :                 NodeId(1),
    1279            4 :                 dummy_broker_sk_timeline(selected_lsn, DUMMY_SAFEKEEPER_HOST, now),
    1280            4 :             ),
    1281            4 :             (
    1282            4 :                 NodeId(2),
    1283            4 :                 dummy_broker_sk_timeline(selected_lsn + 100, "", now),
    1284            4 :             ),
    1285            4 :         ]);
    1286            4 :         let biggest_wal_candidate = state.next_connection_candidate().expect(
    1287            4 :             "Expected one candidate selected out of multiple valid data options, but got none",
    1288            4 :         );
    1289            4 : 
    1290            4 :         assert_eq!(biggest_wal_candidate.safekeeper_id, NodeId(1));
    1291            4 :         assert_eq!(
    1292            4 :             biggest_wal_candidate.reason,
    1293            4 :             ReconnectReason::NoExistingConnection,
    1294            4 :             "Should select new safekeeper due to missing connection, even if there's also a lag in the wal over the threshold"
    1295            4 :         );
    1296            4 :         assert_eq!(
    1297            4 :             biggest_wal_candidate.wal_source_connconf.host(),
    1298            4 :             &Host::Domain(DUMMY_SAFEKEEPER_HOST.to_owned())
    1299            4 :         );
    1300            4 : 
    1301            4 :         Ok(())
    1302            4 :     }
    1303              : 
    1304              :     #[tokio::test]
    1305            4 :     async fn candidate_with_many_connection_failures() -> anyhow::Result<()> {
    1306            4 :         let harness = TenantHarness::create("candidate_with_many_connection_failures").await?;
    1307            4 :         let mut state = dummy_state(&harness).await;
    1308            4 :         let now = Utc::now().naive_utc();
    1309            4 : 
    1310            4 :         let current_lsn = Lsn(100_000).align();
    1311            4 :         let bigger_lsn = Lsn(current_lsn.0 + 100).align();
    1312            4 : 
    1313            4 :         state.wal_connection = None;
    1314            4 :         state.wal_stream_candidates = HashMap::from([
    1315            4 :             (
    1316            4 :                 NodeId(0),
    1317            4 :                 dummy_broker_sk_timeline(bigger_lsn.0, DUMMY_SAFEKEEPER_HOST, now),
    1318            4 :             ),
    1319            4 :             (
    1320            4 :                 NodeId(1),
    1321            4 :                 dummy_broker_sk_timeline(current_lsn.0, DUMMY_SAFEKEEPER_HOST, now),
    1322            4 :             ),
    1323            4 :         ]);
    1324            4 :         state.wal_connection_retries = HashMap::from([(
    1325            4 :             NodeId(0),
    1326            4 :             RetryInfo {
    1327            4 :                 next_retry_at: now.checked_add_signed(chrono::Duration::hours(1)),
    1328            4 :                 retry_duration_seconds: WALCONNECTION_RETRY_MAX_BACKOFF_SECONDS,
    1329            4 :             },
    1330            4 :         )]);
    1331            4 : 
    1332            4 :         let candidate_with_less_errors = state
    1333            4 :             .next_connection_candidate()
    1334            4 :             .expect("Expected one candidate selected, but got none");
    1335            4 :         assert_eq!(
    1336            4 :             candidate_with_less_errors.safekeeper_id,
    1337            4 :             NodeId(1),
    1338            4 :             "Should select the node with no pending retry cooldown"
    1339            4 :         );
    1340            4 : 
    1341            4 :         Ok(())
    1342            4 :     }
    1343              : 
    1344              :     #[tokio::test]
    1345            4 :     async fn lsn_wal_over_threshold_current_candidate() -> anyhow::Result<()> {
    1346            4 :         let harness = TenantHarness::create("lsn_wal_over_threshcurrent_candidate").await?;
    1347            4 :         let mut state = dummy_state(&harness).await;
    1348            4 :         let current_lsn = Lsn(100_000).align();
    1349            4 :         let now = Utc::now().naive_utc();
    1350            4 : 
    1351            4 :         let connected_sk_id = NodeId(0);
    1352            4 :         let new_lsn = Lsn(current_lsn.0 + state.conf.max_lsn_wal_lag.get() + 1);
    1353            4 : 
    1354            4 :         let connection_status = WalConnectionStatus {
    1355            4 :             is_connected: true,
    1356            4 :             has_processed_wal: true,
    1357            4 :             latest_connection_update: now,
    1358            4 :             latest_wal_update: now,
    1359            4 :             commit_lsn: Some(current_lsn),
    1360            4 :             streaming_lsn: Some(current_lsn),
    1361            4 :             node: connected_sk_id,
    1362            4 :         };
    1363            4 : 
    1364            4 :         state.wal_connection = Some(WalConnection {
    1365            4 :             started_at: now,
    1366            4 :             sk_id: connected_sk_id,
    1367            4 :             availability_zone: None,
    1368            4 :             status: connection_status,
    1369            4 :             connection_task: state.spawn(move |sender, _| async move {
    1370            4 :                 sender
    1371            4 :                     .send(TaskStateUpdate::Progress(connection_status))
    1372            4 :                     .ok();
    1373            4 :                 Ok(())
    1374            4 :             }),
    1375            4 :             discovered_new_wal: None,
    1376            4 :         });
    1377            4 :         state.wal_stream_candidates = HashMap::from([
    1378            4 :             (
    1379            4 :                 connected_sk_id,
    1380            4 :                 dummy_broker_sk_timeline(current_lsn.0, DUMMY_SAFEKEEPER_HOST, now),
    1381            4 :             ),
    1382            4 :             (
    1383            4 :                 NodeId(1),
    1384            4 :                 dummy_broker_sk_timeline(new_lsn.0, "advanced_by_lsn_safekeeper", now),
    1385            4 :             ),
    1386            4 :         ]);
    1387            4 : 
    1388            4 :         let over_threshcurrent_candidate = state.next_connection_candidate().expect(
    1389            4 :             "Expected one candidate selected out of multiple valid data options, but got none",
    1390            4 :         );
    1391            4 : 
    1392            4 :         assert_eq!(over_threshcurrent_candidate.safekeeper_id, NodeId(1));
    1393            4 :         assert_eq!(
    1394            4 :             over_threshcurrent_candidate.reason,
    1395            4 :             ReconnectReason::LaggingWal {
    1396            4 :                 current_commit_lsn: current_lsn,
    1397            4 :                 new_commit_lsn: new_lsn,
    1398            4 :                 threshold: state.conf.max_lsn_wal_lag
    1399            4 :             },
    1400            4 :             "Should select bigger WAL safekeeper if it starts to lag enough"
    1401            4 :         );
    1402            4 :         assert_eq!(
    1403            4 :             over_threshcurrent_candidate.wal_source_connconf.host(),
    1404            4 :             &Host::Domain("advanced_by_lsn_safekeeper".to_owned())
    1405            4 :         );
    1406            4 : 
    1407            4 :         Ok(())
    1408            4 :     }
    1409              : 
    1410              :     #[tokio::test]
    1411            4 :     async fn timeout_connection_threshold_current_candidate() -> anyhow::Result<()> {
    1412            4 :         let harness =
    1413            4 :             TenantHarness::create("timeout_connection_threshold_current_candidate").await?;
    1414            4 :         let mut state = dummy_state(&harness).await;
    1415            4 :         let current_lsn = Lsn(100_000).align();
    1416            4 :         let now = Utc::now().naive_utc();
    1417            4 : 
    1418            4 :         let wal_connect_timeout = chrono::Duration::from_std(state.conf.wal_connect_timeout)?;
    1419            4 :         let time_over_threshold =
    1420            4 :             Utc::now().naive_utc() - wal_connect_timeout - wal_connect_timeout;
    1421            4 : 
    1422            4 :         let connection_status = WalConnectionStatus {
    1423            4 :             is_connected: true,
    1424            4 :             has_processed_wal: true,
    1425            4 :             latest_connection_update: time_over_threshold,
    1426            4 :             latest_wal_update: time_over_threshold,
    1427            4 :             commit_lsn: Some(current_lsn),
    1428            4 :             streaming_lsn: Some(current_lsn),
    1429            4 :             node: NodeId(1),
    1430            4 :         };
    1431            4 : 
    1432            4 :         state.wal_connection = Some(WalConnection {
    1433            4 :             started_at: now,
    1434            4 :             sk_id: NodeId(1),
    1435            4 :             availability_zone: None,
    1436            4 :             status: connection_status,
    1437            4 :             connection_task: state.spawn(move |sender, _| async move {
    1438            4 :                 sender
    1439            4 :                     .send(TaskStateUpdate::Progress(connection_status))
    1440            4 :                     .ok();
    1441            4 :                 Ok(())
    1442            4 :             }),
    1443            4 :             discovered_new_wal: None,
    1444            4 :         });
    1445            4 :         state.wal_stream_candidates = HashMap::from([(
    1446            4 :             NodeId(0),
    1447            4 :             dummy_broker_sk_timeline(current_lsn.0, DUMMY_SAFEKEEPER_HOST, now),
    1448            4 :         )]);
    1449            4 : 
    1450            4 :         let over_threshcurrent_candidate = state.next_connection_candidate().expect(
    1451            4 :             "Expected one candidate selected out of multiple valid data options, but got none",
    1452            4 :         );
    1453            4 : 
    1454            4 :         assert_eq!(over_threshcurrent_candidate.safekeeper_id, NodeId(0));
    1455            4 :         match over_threshcurrent_candidate.reason {
    1456            4 :             ReconnectReason::NoKeepAlives {
    1457            4 :                 last_keep_alive,
    1458            4 :                 threshold,
    1459            4 :                 ..
    1460            4 :             } => {
    1461            4 :                 assert_eq!(last_keep_alive, Some(time_over_threshold));
    1462            4 :                 assert_eq!(threshold, state.conf.lagging_wal_timeout);
    1463            4 :             }
    1464            4 :             unexpected => panic!("Unexpected reason: {unexpected:?}"),
    1465            4 :         }
    1466            4 :         assert_eq!(
    1467            4 :             over_threshcurrent_candidate.wal_source_connconf.host(),
    1468            4 :             &Host::Domain(DUMMY_SAFEKEEPER_HOST.to_owned())
    1469            4 :         );
    1470            4 : 
    1471            4 :         Ok(())
    1472            4 :     }
    1473              : 
    1474              :     #[tokio::test]
    1475            4 :     async fn timeout_wal_over_threshold_current_candidate() -> anyhow::Result<()> {
    1476            4 :         let harness = TenantHarness::create("timeout_wal_over_threshold_current_candidate").await?;
    1477            4 :         let mut state = dummy_state(&harness).await;
    1478            4 :         let current_lsn = Lsn(100_000).align();
    1479            4 :         let new_lsn = Lsn(100_100).align();
    1480            4 :         let now = Utc::now().naive_utc();
    1481            4 : 
    1482            4 :         let lagging_wal_timeout = chrono::Duration::from_std(state.conf.lagging_wal_timeout)?;
    1483            4 :         let time_over_threshold =
    1484            4 :             Utc::now().naive_utc() - lagging_wal_timeout - lagging_wal_timeout;
    1485            4 : 
    1486            4 :         let connection_status = WalConnectionStatus {
    1487            4 :             is_connected: true,
    1488            4 :             has_processed_wal: true,
    1489            4 :             latest_connection_update: now,
    1490            4 :             latest_wal_update: time_over_threshold,
    1491            4 :             commit_lsn: Some(current_lsn),
    1492            4 :             streaming_lsn: Some(current_lsn),
    1493            4 :             node: NodeId(1),
    1494            4 :         };
    1495            4 : 
    1496            4 :         state.wal_connection = Some(WalConnection {
    1497            4 :             started_at: now,
    1498            4 :             sk_id: NodeId(1),
    1499            4 :             availability_zone: None,
    1500            4 :             status: connection_status,
    1501            4 :             connection_task: state.spawn(move |_, _| async move { Ok(()) }),
    1502            4 :             discovered_new_wal: Some(NewCommittedWAL {
    1503            4 :                 discovered_at: time_over_threshold,
    1504            4 :                 lsn: new_lsn,
    1505            4 :             }),
    1506            4 :         });
    1507            4 :         state.wal_stream_candidates = HashMap::from([(
    1508            4 :             NodeId(0),
    1509            4 :             dummy_broker_sk_timeline(new_lsn.0, DUMMY_SAFEKEEPER_HOST, now),
    1510            4 :         )]);
    1511            4 : 
    1512            4 :         let over_threshcurrent_candidate = state.next_connection_candidate().expect(
    1513            4 :             "Expected one candidate selected out of multiple valid data options, but got none",
    1514            4 :         );
    1515            4 : 
    1516            4 :         assert_eq!(over_threshcurrent_candidate.safekeeper_id, NodeId(0));
    1517            4 :         match over_threshcurrent_candidate.reason {
    1518            4 :             ReconnectReason::NoWalTimeout {
    1519            4 :                 current_lsn,
    1520            4 :                 current_commit_lsn,
    1521            4 :                 candidate_commit_lsn,
    1522            4 :                 last_wal_interaction,
    1523            4 :                 threshold,
    1524            4 :                 ..
    1525            4 :             } => {
    1526            4 :                 assert_eq!(current_lsn, current_lsn);
    1527            4 :                 assert_eq!(current_commit_lsn, current_lsn);
    1528            4 :                 assert_eq!(candidate_commit_lsn, new_lsn);
    1529            4 :                 assert_eq!(last_wal_interaction, Some(time_over_threshold));
    1530            4 :                 assert_eq!(threshold, state.conf.lagging_wal_timeout);
    1531            4 :             }
    1532            4 :             unexpected => panic!("Unexpected reason: {unexpected:?}"),
    1533            4 :         }
    1534            4 :         assert_eq!(
    1535            4 :             over_threshcurrent_candidate.wal_source_connconf.host(),
    1536            4 :             &Host::Domain(DUMMY_SAFEKEEPER_HOST.to_owned())
    1537            4 :         );
    1538            4 : 
    1539            4 :         Ok(())
    1540            4 :     }
    1541              : 
    1542              :     const DUMMY_SAFEKEEPER_HOST: &str = "safekeeper_connstr";
    1543              : 
    1544           32 :     async fn dummy_state(harness: &TenantHarness) -> ConnectionManagerState {
    1545           32 :         let (tenant, ctx) = harness.load().await;
    1546           32 :         let timeline = tenant
    1547           32 :             .create_test_timeline(TIMELINE_ID, Lsn(0x8), crate::DEFAULT_PG_VERSION, &ctx)
    1548           32 :             .await
    1549           32 :             .expect("Failed to create an empty timeline for dummy wal connection manager");
    1550           32 : 
    1551           32 :         ConnectionManagerState {
    1552           32 :             id: TenantTimelineId {
    1553           32 :                 tenant_id: harness.tenant_shard_id.tenant_id,
    1554           32 :                 timeline_id: TIMELINE_ID,
    1555           32 :             },
    1556           32 :             timeline,
    1557           32 :             cancel: CancellationToken::new(),
    1558           32 :             conf: WalReceiverConf {
    1559           32 :                 protocol: DEFAULT_WAL_RECEIVER_PROTOCOL,
    1560           32 :                 wal_connect_timeout: Duration::from_secs(1),
    1561           32 :                 lagging_wal_timeout: Duration::from_secs(1),
    1562           32 :                 max_lsn_wal_lag: NonZeroU64::new(1024 * 1024).unwrap(),
    1563           32 :                 auth_token: None,
    1564           32 :                 availability_zone: None,
    1565           32 :                 ingest_batch_size: 1,
    1566           32 :             },
    1567           32 :             wal_connection: None,
    1568           32 :             wal_stream_candidates: HashMap::new(),
    1569           32 :             wal_connection_retries: HashMap::new(),
    1570           32 :         }
    1571           32 :     }
    1572              : 
    1573              :     #[tokio::test]
    1574            4 :     async fn switch_to_same_availability_zone() -> anyhow::Result<()> {
    1575            4 :         // Pageserver and one of safekeepers will be in the same availability zone
    1576            4 :         // and pageserver should prefer to connect to it.
    1577            4 :         let test_az = Some("test_az".to_owned());
    1578            4 : 
    1579            4 :         let harness = TenantHarness::create("switch_to_same_availability_zone").await?;
    1580            4 :         let mut state = dummy_state(&harness).await;
    1581            4 :         state.conf.availability_zone.clone_from(&test_az);
    1582            4 :         let current_lsn = Lsn(100_000).align();
    1583            4 :         let now = Utc::now().naive_utc();
    1584            4 : 
    1585            4 :         let connected_sk_id = NodeId(0);
    1586            4 : 
    1587            4 :         let connection_status = WalConnectionStatus {
    1588            4 :             is_connected: true,
    1589            4 :             has_processed_wal: true,
    1590            4 :             latest_connection_update: now,
    1591            4 :             latest_wal_update: now,
    1592            4 :             commit_lsn: Some(current_lsn),
    1593            4 :             streaming_lsn: Some(current_lsn),
    1594            4 :             node: connected_sk_id,
    1595            4 :         };
    1596            4 : 
    1597            4 :         state.wal_connection = Some(WalConnection {
    1598            4 :             started_at: now,
    1599            4 :             sk_id: connected_sk_id,
    1600            4 :             availability_zone: None,
    1601            4 :             status: connection_status,
    1602            4 :             connection_task: state.spawn(move |sender, _| async move {
    1603            4 :                 sender
    1604            4 :                     .send(TaskStateUpdate::Progress(connection_status))
    1605            4 :                     .ok();
    1606            4 :                 Ok(())
    1607            4 :             }),
    1608            4 :             discovered_new_wal: None,
    1609            4 :         });
    1610            4 : 
    1611            4 :         // We have another safekeeper with the same commit_lsn, and it have the same availability zone as
    1612            4 :         // the current pageserver.
    1613            4 :         let mut same_az_sk = dummy_broker_sk_timeline(current_lsn.0, "same_az", now);
    1614            4 :         same_az_sk.timeline.availability_zone.clone_from(&test_az);
    1615            4 : 
    1616            4 :         state.wal_stream_candidates = HashMap::from([
    1617            4 :             (
    1618            4 :                 connected_sk_id,
    1619            4 :                 dummy_broker_sk_timeline(current_lsn.0, DUMMY_SAFEKEEPER_HOST, now),
    1620            4 :             ),
    1621            4 :             (NodeId(1), same_az_sk),
    1622            4 :         ]);
    1623            4 : 
    1624            4 :         // We expect that pageserver will switch to the safekeeper in the same availability zone,
    1625            4 :         // even if it has the same commit_lsn.
    1626            4 :         let next_candidate = state.next_connection_candidate().expect(
    1627            4 :             "Expected one candidate selected out of multiple valid data options, but got none",
    1628            4 :         );
    1629            4 : 
    1630            4 :         assert_eq!(next_candidate.safekeeper_id, NodeId(1));
    1631            4 :         assert_eq!(
    1632            4 :             next_candidate.reason,
    1633            4 :             ReconnectReason::SwitchAvailabilityZone,
    1634            4 :             "Should switch to the safekeeper in the same availability zone, if it has the same commit_lsn"
    1635            4 :         );
    1636            4 :         assert_eq!(
    1637            4 :             next_candidate.wal_source_connconf.host(),
    1638            4 :             &Host::Domain("same_az".to_owned())
    1639            4 :         );
    1640            4 : 
    1641            4 :         Ok(())
    1642            4 :     }
    1643              : }
        

Generated by: LCOV version 2.1-beta