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

Generated by: LCOV version 2.1-beta