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

Generated by: LCOV version 2.1-beta