LCOV - code coverage report
Current view: top level - safekeeper/tests/walproposer_sim - safekeeper.rs (source / functions) Coverage Total Hit
Test: a1cc1f33dc9899e4da66eb51e44e911a4b3bd648.info Lines: 95.0 % 281 267
Test Date: 2025-07-31 11:35:14 Functions: 94.9 % 39 37

            Line data    Source code
       1              : //! Safekeeper communication endpoint to WAL proposer (compute node).
       2              : //! Gets messages from the network, passes them down to consensus module and
       3              : //! sends replies back.
       4              : 
       5              : use std::collections::HashMap;
       6              : use std::sync::Arc;
       7              : use std::time::Duration;
       8              : 
       9              : use anyhow::{Result, bail};
      10              : use bytes::{Bytes, BytesMut};
      11              : use camino::Utf8PathBuf;
      12              : use desim::executor::{self, PollSome};
      13              : use desim::network::TCP;
      14              : use desim::node_os::NodeOs;
      15              : use desim::proto::{AnyMessage, NetEvent, NodeEvent};
      16              : use http::Uri;
      17              : use postgres_backend::AuthType;
      18              : use safekeeper::SafeKeeperConf;
      19              : use safekeeper::safekeeper::{
      20              :     ProposerAcceptorMessage, SK_PROTO_VERSION_3, SafeKeeper, UNKNOWN_SERVER_VERSION,
      21              : };
      22              : use safekeeper::state::{TimelinePersistentState, TimelineState};
      23              : use safekeeper::timeline::TimelineError;
      24              : use safekeeper::wal_storage::Storage;
      25              : use safekeeper_api::ServerInfo;
      26              : use safekeeper_api::membership::Configuration;
      27              : use tracing::{debug, info_span, warn};
      28              : use utils::id::{NodeId, TenantId, TenantTimelineId, TimelineId};
      29              : use utils::lsn::Lsn;
      30              : 
      31              : use super::safekeeper_disk::{DiskStateStorage, DiskWALStorage, SafekeeperDisk, TimelineDisk};
      32              : 
      33              : struct SharedState {
      34              :     sk: SafeKeeper<DiskStateStorage, DiskWALStorage>,
      35              :     disk: Arc<TimelineDisk>,
      36              : }
      37              : 
      38              : struct GlobalMap {
      39              :     timelines: HashMap<TenantTimelineId, SharedState>,
      40              :     conf: SafeKeeperConf,
      41              :     disk: Arc<SafekeeperDisk>,
      42              : }
      43              : 
      44              : impl GlobalMap {
      45              :     /// Restores global state from disk.
      46         9623 :     fn new(disk: Arc<SafekeeperDisk>, conf: SafeKeeperConf) -> Result<Self> {
      47         9623 :         let mut timelines = HashMap::new();
      48              : 
      49         9623 :         for (&ttid, disk) in disk.timelines.lock().iter() {
      50         7223 :             debug!("loading timeline {}", ttid);
      51         7223 :             let state = disk.state.lock().clone();
      52              : 
      53         7223 :             if state.server.wal_seg_size == 0 {
      54            0 :                 bail!(TimelineError::UninitializedWalSegSize(ttid));
      55         7223 :             }
      56              : 
      57         7223 :             if state.server.pg_version == UNKNOWN_SERVER_VERSION {
      58            0 :                 bail!(TimelineError::UninitialinzedPgVersion(ttid));
      59         7223 :             }
      60              : 
      61         7223 :             if state.commit_lsn < state.local_start_lsn {
      62            0 :                 bail!(
      63            0 :                     "commit_lsn {} is smaller than local_start_lsn {}",
      64              :                     state.commit_lsn,
      65              :                     state.local_start_lsn
      66              :                 );
      67         7223 :             }
      68              : 
      69         7223 :             let control_store = DiskStateStorage::new(disk.clone());
      70         7223 :             let wal_store = DiskWALStorage::new(disk.clone(), &control_store)?;
      71              : 
      72         7223 :             let sk = SafeKeeper::new(TimelineState::new(control_store), wal_store, conf.my_id)?;
      73         7223 :             timelines.insert(
      74         7223 :                 ttid,
      75         7223 :                 SharedState {
      76         7223 :                     sk,
      77         7223 :                     disk: disk.clone(),
      78         7223 :                 },
      79              :             );
      80              :         }
      81              : 
      82         9623 :         Ok(Self {
      83         9623 :             timelines,
      84         9623 :             conf,
      85         9623 :             disk,
      86         9623 :         })
      87         9623 :     }
      88              : 
      89         1434 :     fn create(&mut self, ttid: TenantTimelineId, server_info: ServerInfo) -> Result<()> {
      90         1434 :         if self.timelines.contains_key(&ttid) {
      91            0 :             bail!("timeline {} already exists", ttid);
      92         1434 :         }
      93              : 
      94         1434 :         debug!("creating new timeline {}", ttid);
      95              : 
      96         1434 :         let commit_lsn = Lsn::INVALID;
      97         1434 :         let local_start_lsn = Lsn::INVALID;
      98              : 
      99         1434 :         let state = TimelinePersistentState::new(
     100         1434 :             &ttid,
     101         1434 :             Configuration::empty(),
     102         1434 :             server_info,
     103         1434 :             commit_lsn,
     104         1434 :             local_start_lsn,
     105            0 :         )?;
     106              : 
     107         1434 :         let disk_timeline = self.disk.put_state(&ttid, state);
     108         1434 :         let control_store = DiskStateStorage::new(disk_timeline.clone());
     109         1434 :         let wal_store = DiskWALStorage::new(disk_timeline.clone(), &control_store)?;
     110              : 
     111         1434 :         let sk = SafeKeeper::new(
     112         1434 :             TimelineState::new(control_store),
     113         1434 :             wal_store,
     114         1434 :             self.conf.my_id,
     115            0 :         )?;
     116              : 
     117         1434 :         self.timelines.insert(
     118         1434 :             ttid,
     119         1434 :             SharedState {
     120         1434 :                 sk,
     121         1434 :                 disk: disk_timeline,
     122         1434 :             },
     123              :         );
     124         1434 :         Ok(())
     125         1434 :     }
     126              : 
     127        33283 :     fn get(&mut self, ttid: &TenantTimelineId) -> &mut SharedState {
     128        33283 :         self.timelines.get_mut(ttid).expect("timeline must exist")
     129        33283 :     }
     130              : 
     131        20450 :     fn has_tli(&self, ttid: &TenantTimelineId) -> bool {
     132        20450 :         self.timelines.contains_key(ttid)
     133        20450 :     }
     134              : }
     135              : 
     136              : /// State of a single connection to walproposer.
     137              : struct ConnState {
     138              :     tcp: TCP,
     139              : 
     140              :     greeting: bool,
     141              :     ttid: TenantTimelineId,
     142              :     flush_pending: bool,
     143              : 
     144              :     runtime: tokio::runtime::Runtime,
     145              : }
     146              : 
     147         9623 : pub fn run_server(os: NodeOs, disk: Arc<SafekeeperDisk>) -> Result<()> {
     148         9623 :     let _enter = info_span!("safekeeper", id = os.id()).entered();
     149         9623 :     debug!("started server");
     150         9623 :     os.log_event("started;safekeeper".to_owned());
     151         9623 :     let conf = SafeKeeperConf {
     152         9623 :         workdir: Utf8PathBuf::from("."),
     153         9623 :         my_id: NodeId(os.id() as u64),
     154         9623 :         listen_pg_addr: String::new(),
     155         9623 :         listen_http_addr: String::new(),
     156         9623 :         listen_https_addr: None,
     157         9623 :         no_sync: false,
     158         9623 :         broker_endpoint: "/".parse::<Uri>().unwrap(),
     159         9623 :         broker_keepalive_interval: Duration::from_secs(0),
     160         9623 :         heartbeat_timeout: Duration::from_secs(0),
     161         9623 :         remote_storage: None,
     162         9623 :         max_offloader_lag_bytes: 0,
     163         9623 :         /* BEGIN_HADRON */
     164         9623 :         max_reelect_offloader_lag_bytes: 0,
     165         9623 :         max_timeline_disk_usage_bytes: 0,
     166         9623 :         /* END_HADRON */
     167         9623 :         wal_backup_enabled: false,
     168         9623 :         listen_pg_addr_tenant_only: None,
     169         9623 :         advertise_pg_addr: None,
     170         9623 :         availability_zone: None,
     171         9623 :         peer_recovery_enabled: false,
     172         9623 :         backup_parallel_jobs: 0,
     173         9623 :         auth_type: AuthType::NeonJWT,
     174         9623 :         pg_auth: None,
     175         9623 :         pg_tenant_only_auth: None,
     176         9623 :         http_auth: None,
     177         9623 :         sk_auth_token: None,
     178         9623 :         current_thread_runtime: false,
     179         9623 :         walsenders_keep_horizon: false,
     180         9623 :         partial_backup_timeout: Duration::from_secs(0),
     181         9623 :         disable_periodic_broker_push: false,
     182         9623 :         enable_offload: false,
     183         9623 :         delete_offloaded_wal: false,
     184         9623 :         control_file_save_interval: Duration::from_secs(1),
     185         9623 :         partial_backup_concurrency: 1,
     186         9623 :         eviction_min_resident: Duration::ZERO,
     187         9623 :         wal_reader_fanout: false,
     188         9623 :         max_delta_for_fanout: None,
     189         9623 :         ssl_key_file: Utf8PathBuf::from(""),
     190         9623 :         ssl_cert_file: Utf8PathBuf::from(""),
     191         9623 :         ssl_cert_reload_period: Duration::ZERO,
     192         9623 :         ssl_ca_certs: Vec::new(),
     193         9623 :         use_https_safekeeper_api: false,
     194         9623 :         enable_tls_wal_service_api: false,
     195         9623 :         force_metric_collection_on_scrape: true,
     196         9623 :         /* BEGIN_HADRON */
     197         9623 :         enable_pull_timeline_on_startup: false,
     198         9623 :         advertise_pg_addr_tenant_only: None,
     199         9623 :         hcc_base_url: None,
     200         9623 :         global_disk_check_interval: Duration::from_secs(10),
     201         9623 :         max_global_disk_usage_ratio: 0.0,
     202         9623 :         /* END_HADRON */
     203         9623 :     };
     204              : 
     205         9623 :     let mut global = GlobalMap::new(disk, conf.clone())?;
     206         9623 :     let mut conns: HashMap<usize, ConnState> = HashMap::new();
     207              : 
     208         9623 :     for (&_ttid, shared_state) in global.timelines.iter_mut() {
     209         7223 :         let flush_lsn = shared_state.sk.wal_store.flush_lsn();
     210         7223 :         let commit_lsn = shared_state.sk.state.commit_lsn;
     211         7223 :         os.log_event(format!("tli_loaded;{};{}", flush_lsn.0, commit_lsn.0));
     212         7223 :     }
     213              : 
     214         9623 :     let node_events = os.node_events();
     215         9623 :     let mut epoll_vec: Vec<Box<dyn PollSome>> = vec![];
     216         9623 :     let mut epoll_idx: Vec<usize> = vec![];
     217              : 
     218              :     // TODO: batch events processing (multiple events per tick)
     219              :     loop {
     220        79015 :         epoll_vec.clear();
     221        79015 :         epoll_idx.clear();
     222              : 
     223              :         // node events channel
     224        79015 :         epoll_vec.push(Box::new(node_events.clone()));
     225        79015 :         epoll_idx.push(0);
     226              : 
     227              :         // tcp connections
     228       304461 :         for conn in conns.values() {
     229       304461 :             epoll_vec.push(Box::new(conn.tcp.recv_chan()));
     230       304461 :             epoll_idx.push(conn.tcp.connection_id());
     231       304461 :         }
     232              : 
     233              :         // waiting for the next message
     234        79015 :         let index = executor::epoll_chans(&epoll_vec, -1).unwrap();
     235              : 
     236        79015 :         if index == 0 {
     237              :             // got a new connection
     238        25803 :             match node_events.must_recv() {
     239        25803 :                 NodeEvent::Accept(tcp) => {
     240        25803 :                     conns.insert(
     241        25803 :                         tcp.connection_id(),
     242              :                         ConnState {
     243        25803 :                             tcp,
     244              :                             greeting: false,
     245        25803 :                             ttid: TenantTimelineId::empty(),
     246              :                             flush_pending: false,
     247        25803 :                             runtime: tokio::runtime::Builder::new_current_thread().build()?,
     248              :                         },
     249              :                     );
     250              :                 }
     251            0 :                 NodeEvent::Internal(_) => unreachable!(),
     252              :             }
     253        25803 :             continue;
     254        53212 :         }
     255              : 
     256        53212 :         let connection_id = epoll_idx[index];
     257        53212 :         let conn = conns.get_mut(&connection_id).unwrap();
     258        53212 :         let mut next_event = Some(conn.tcp.recv_chan().must_recv());
     259              : 
     260              :         loop {
     261        88244 :             let event = match next_event {
     262        44925 :                 Some(event) => event,
     263        43319 :                 None => break,
     264              :             };
     265              : 
     266        44925 :             match event {
     267        29522 :                 NetEvent::Message(msg) => {
     268        29522 :                     let res = conn.process_any(msg, &mut global);
     269        29522 :                     if res.is_err() {
     270         9893 :                         let e = res.unwrap_err();
     271         9893 :                         let estr = e.to_string();
     272         9893 :                         if !estr.contains("finished processing START_REPLICATION") {
     273         9623 :                             warn!("conn {:?} error: {:?}", connection_id, e);
     274         9623 :                             panic!("unexpected error at safekeeper: {e:#}");
     275          270 :                         }
     276          270 :                         conns.remove(&connection_id);
     277          270 :                         break;
     278        19629 :                     }
     279              :                 }
     280        15403 :                 NetEvent::Closed => {
     281        15403 :                     // TODO: remove from conns?
     282        15403 :                 }
     283              :             }
     284              : 
     285        35032 :             next_event = conn.tcp.recv_chan().try_recv();
     286              :         }
     287              : 
     288       202137 :         conns.retain(|_, conn| {
     289       202137 :             let res = conn.flush(&mut global);
     290       202137 :             if res.is_err() {
     291            0 :                 debug!("conn {:?} error: {:?}", conn.tcp, res);
     292       202137 :             }
     293       202137 :             res.is_ok()
     294       202137 :         });
     295              :     }
     296            0 : }
     297              : 
     298              : impl ConnState {
     299              :     /// Process a message from the network. It can be START_REPLICATION request or a valid ProposerAcceptorMessage message.
     300        29522 :     fn process_any(&mut self, any: AnyMessage, global: &mut GlobalMap) -> Result<()> {
     301        29522 :         if let AnyMessage::Bytes(copy_data) = any {
     302        29522 :             let repl_prefix = b"START_REPLICATION ";
     303        29522 :             if !self.greeting && copy_data.starts_with(repl_prefix) {
     304          270 :                 self.process_start_replication(copy_data.slice(repl_prefix.len()..), global)?;
     305          270 :                 bail!("finished processing START_REPLICATION")
     306        29252 :             }
     307              : 
     308        29252 :             let msg = ProposerAcceptorMessage::parse(copy_data, SK_PROTO_VERSION_3)?;
     309        29252 :             debug!("got msg: {:?}", msg);
     310        29252 :             self.process(msg, global)
     311              :         } else {
     312            0 :             bail!("unexpected message, expected AnyMessage::Bytes");
     313              :         }
     314        29522 :     }
     315              : 
     316              :     /// Process START_REPLICATION request.
     317          270 :     fn process_start_replication(
     318          270 :         &mut self,
     319          270 :         copy_data: Bytes,
     320          270 :         global: &mut GlobalMap,
     321          270 :     ) -> Result<()> {
     322              :         // format is "<tenant_id> <timeline_id> <start_lsn> <end_lsn>"
     323          270 :         let str = String::from_utf8(copy_data.to_vec())?;
     324              : 
     325          270 :         let mut parts = str.split(' ');
     326          270 :         let tenant_id = parts.next().unwrap().parse::<TenantId>()?;
     327          270 :         let timeline_id = parts.next().unwrap().parse::<TimelineId>()?;
     328          270 :         let start_lsn = parts.next().unwrap().parse::<u64>()?;
     329          270 :         let end_lsn = parts.next().unwrap().parse::<u64>()?;
     330              : 
     331          270 :         let ttid = TenantTimelineId::new(tenant_id, timeline_id);
     332          270 :         let shared_state = global.get(&ttid);
     333              : 
     334              :         // read bytes from start_lsn to end_lsn
     335          270 :         let mut buf = vec![0; (end_lsn - start_lsn) as usize];
     336          270 :         shared_state.disk.wal.lock().read(start_lsn, &mut buf);
     337              : 
     338              :         // send bytes to the client
     339          270 :         self.tcp.send(AnyMessage::Bytes(Bytes::from(buf)));
     340          270 :         Ok(())
     341          270 :     }
     342              : 
     343              :     /// Get or create a timeline.
     344        20450 :     fn init_timeline(
     345        20450 :         &mut self,
     346        20450 :         ttid: TenantTimelineId,
     347        20450 :         server_info: ServerInfo,
     348        20450 :         global: &mut GlobalMap,
     349        20450 :     ) -> Result<()> {
     350        20450 :         self.ttid = ttid;
     351        20450 :         if global.has_tli(&ttid) {
     352        19016 :             return Ok(());
     353         1434 :         }
     354              : 
     355         1434 :         global.create(ttid, server_info)
     356        20450 :     }
     357              : 
     358              :     /// Process a ProposerAcceptorMessage.
     359        29252 :     fn process(&mut self, msg: ProposerAcceptorMessage, global: &mut GlobalMap) -> Result<()> {
     360        29252 :         if !self.greeting {
     361        20450 :             self.greeting = true;
     362              : 
     363        20450 :             match msg {
     364        20450 :                 ProposerAcceptorMessage::Greeting(ref greeting) => {
     365        20450 :                     tracing::info!(
     366            0 :                         "start handshake with walproposer {:?} {:?}",
     367              :                         self.tcp,
     368              :                         greeting
     369              :                     );
     370        20450 :                     let server_info = ServerInfo {
     371        20450 :                         pg_version: greeting.pg_version,
     372        20450 :                         system_id: greeting.system_id,
     373        20450 :                         wal_seg_size: greeting.wal_seg_size,
     374        20450 :                     };
     375        20450 :                     let ttid = TenantTimelineId::new(greeting.tenant_id, greeting.timeline_id);
     376        20450 :                     self.init_timeline(ttid, server_info, global)?
     377              :                 }
     378              :                 _ => {
     379            0 :                     bail!("unexpected message {msg:?} instead of greeting");
     380              :                 }
     381              :             }
     382         8802 :         }
     383              : 
     384        29252 :         let tli = global.get(&self.ttid);
     385              : 
     386        29252 :         match msg {
     387         4493 :             ProposerAcceptorMessage::AppendRequest(append_request) => {
     388         4493 :                 self.flush_pending = true;
     389         4493 :                 self.process_sk_msg(
     390         4493 :                     tli,
     391         4493 :                     &ProposerAcceptorMessage::NoFlushAppendRequest(append_request),
     392            0 :                 )?;
     393              :             }
     394        24759 :             other => {
     395        24759 :                 self.process_sk_msg(tli, &other)?;
     396              :             }
     397              :         }
     398              : 
     399        29252 :         Ok(())
     400        29252 :     }
     401              : 
     402              :     /// Process FlushWAL if needed.
     403       202137 :     fn flush(&mut self, global: &mut GlobalMap) -> Result<()> {
     404              :         // TODO: try to add extra flushes in simulation, to verify that extra flushes don't break anything
     405       202137 :         if !self.flush_pending {
     406       198376 :             return Ok(());
     407         3761 :         }
     408         3761 :         self.flush_pending = false;
     409         3761 :         let shared_state = global.get(&self.ttid);
     410         3761 :         self.process_sk_msg(shared_state, &ProposerAcceptorMessage::FlushWAL)
     411       202137 :     }
     412              : 
     413              :     /// Make safekeeper process a message and send a reply to the TCP
     414        33013 :     fn process_sk_msg(
     415        33013 :         &mut self,
     416        33013 :         shared_state: &mut SharedState,
     417        33013 :         msg: &ProposerAcceptorMessage,
     418        33013 :     ) -> Result<()> {
     419        33013 :         let mut reply = self.runtime.block_on(shared_state.sk.process_msg(msg))?;
     420        33013 :         if let Some(reply) = &mut reply {
     421              :             // TODO: if this is AppendResponse, fill in proper hot standby feedback and disk consistent lsn
     422              : 
     423        27484 :             let mut buf = BytesMut::with_capacity(128);
     424        27484 :             reply.serialize(&mut buf, SK_PROTO_VERSION_3)?;
     425              : 
     426        27484 :             self.tcp.send(AnyMessage::Bytes(buf.into()));
     427         5529 :         }
     428        33013 :         Ok(())
     429        33013 :     }
     430              : }
     431              : 
     432              : impl Drop for ConnState {
     433        25777 :     fn drop(&mut self) {
     434        25777 :         debug!("dropping conn: {:?}", self.tcp);
     435        25777 :         if !std::thread::panicking() {
     436          270 :             self.tcp.close();
     437        25507 :         }
     438              :         // TODO: clean up non-fsynced WAL
     439        25777 :     }
     440              : }
        

Generated by: LCOV version 2.1-beta