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

Generated by: LCOV version 2.1-beta