LCOV - code coverage report
Current view: top level - safekeeper/tests/walproposer_sim - safekeeper.rs (source / functions) Coverage Total Hit
Test: 190869232aac3a234374e5bb62582e91cf5f5818.info Lines: 92.9 % 283 263
Test Date: 2024-02-23 13:21:27 Functions: 79.4 % 63 50

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

Generated by: LCOV version 2.1-beta