LCOV - code coverage report
Current view: top level - storage_broker/src/bin - storage_broker.rs (source / functions) Coverage Total Hit
Test: 5fe7fa8d483b39476409aee736d6d5e32728bfac.info Lines: 37.9 % 560 212
Test Date: 2025-03-12 16:10:49 Functions: 29.2 % 65 19

            Line data    Source code
       1              : //! Simple pub-sub based on grpc (tonic) and Tokio broadcast channel for storage
       2              : //! nodes messaging.
       3              : //!
       4              : //! Subscriptions to 1) single timeline 2) all timelines are possible. We could
       5              : //! add subscription to the set of timelines to save grpc streams, but testing
       6              : //! shows many individual streams is also ok.
       7              : //!
       8              : //! Message is dropped if subscriber can't consume it, not affecting other
       9              : //! subscribers.
      10              : //!
      11              : //! Only safekeeper message is supported, but it is not hard to add something
      12              : //! else with generics.
      13              : use std::collections::HashMap;
      14              : use std::convert::Infallible;
      15              : use std::net::SocketAddr;
      16              : use std::pin::Pin;
      17              : use std::sync::Arc;
      18              : use std::time::Duration;
      19              : 
      20              : use clap::{Parser, command};
      21              : use futures_core::Stream;
      22              : use futures_util::StreamExt;
      23              : use http_body_util::Full;
      24              : use hyper::body::Incoming;
      25              : use hyper::header::CONTENT_TYPE;
      26              : use hyper::service::service_fn;
      27              : use hyper::{Method, StatusCode};
      28              : use hyper_util::rt::{TokioExecutor, TokioIo, TokioTimer};
      29              : use metrics::{Encoder, TextEncoder};
      30              : use parking_lot::RwLock;
      31              : use storage_broker::metrics::{
      32              :     BROADCAST_DROPPED_MESSAGES_TOTAL, BROADCASTED_MESSAGES_TOTAL, NUM_PUBS, NUM_SUBS_ALL,
      33              :     NUM_SUBS_TIMELINE, PROCESSED_MESSAGES_TOTAL, PUBLISHED_ONEOFF_MESSAGES_TOTAL,
      34              : };
      35              : use storage_broker::proto::broker_service_server::{BrokerService, BrokerServiceServer};
      36              : use storage_broker::proto::subscribe_safekeeper_info_request::SubscriptionKey as ProtoSubscriptionKey;
      37              : use storage_broker::proto::{
      38              :     FilterTenantTimelineId, MessageType, SafekeeperDiscoveryRequest, SafekeeperDiscoveryResponse,
      39              :     SafekeeperTimelineInfo, SubscribeByFilterRequest, SubscribeSafekeeperInfoRequest, TypedMessage,
      40              : };
      41              : use storage_broker::{DEFAULT_KEEPALIVE_INTERVAL, DEFAULT_LISTEN_ADDR, parse_proto_ttid};
      42              : use tokio::net::TcpListener;
      43              : use tokio::sync::broadcast;
      44              : use tokio::sync::broadcast::error::RecvError;
      45              : use tokio::time;
      46              : use tonic::body::{self, BoxBody, empty_body};
      47              : use tonic::codegen::Service;
      48              : use tonic::{Code, Request, Response, Status};
      49              : use tracing::*;
      50              : use utils::id::TenantTimelineId;
      51              : use utils::logging::{self, LogFormat};
      52              : use utils::sentry_init::init_sentry;
      53              : use utils::signals::ShutdownSignals;
      54              : use utils::{project_build_tag, project_git_version};
      55              : 
      56              : project_git_version!(GIT_VERSION);
      57              : project_build_tag!(BUILD_TAG);
      58              : 
      59              : const DEFAULT_CHAN_SIZE: usize = 32;
      60              : const DEFAULT_ALL_KEYS_CHAN_SIZE: usize = 16384;
      61              : 
      62              : #[derive(Parser, Debug)]
      63              : #[command(version = GIT_VERSION, about = "Broker for neon storage nodes communication", long_about = None)]
      64              : struct Args {
      65              :     /// Endpoint to listen on.
      66              :     #[arg(short, long, default_value = DEFAULT_LISTEN_ADDR)]
      67            0 :     listen_addr: SocketAddr,
      68              :     /// Size of the queue to the per timeline subscriber.
      69            0 :     #[arg(long, default_value_t = DEFAULT_CHAN_SIZE)]
      70            0 :     timeline_chan_size: usize,
      71              :     /// Size of the queue to the all keys subscriber.
      72            0 :     #[arg(long, default_value_t = DEFAULT_ALL_KEYS_CHAN_SIZE)]
      73            0 :     all_keys_chan_size: usize,
      74              :     /// HTTP/2 keepalive interval.
      75              :     #[arg(long, value_parser= humantime::parse_duration, default_value = DEFAULT_KEEPALIVE_INTERVAL)]
      76            0 :     http2_keepalive_interval: Duration,
      77              :     /// Format for logging, either 'plain' or 'json'.
      78              :     #[arg(long, default_value = "plain")]
      79            0 :     log_format: String,
      80              : }
      81              : 
      82              : /// Id of publisher for registering in maps
      83              : type PubId = u64;
      84              : 
      85              : /// Id of subscriber for registering in maps
      86              : type SubId = u64;
      87              : 
      88              : /// Single enum type for all messages.
      89              : #[derive(Clone, Debug, PartialEq)]
      90              : #[allow(clippy::enum_variant_names)]
      91              : enum Message {
      92              :     SafekeeperTimelineInfo(SafekeeperTimelineInfo),
      93              :     SafekeeperDiscoveryRequest(SafekeeperDiscoveryRequest),
      94              :     SafekeeperDiscoveryResponse(SafekeeperDiscoveryResponse),
      95              : }
      96              : 
      97              : impl Message {
      98              :     /// Convert proto message to internal message.
      99            0 :     pub fn from(proto_msg: TypedMessage) -> Result<Self, Status> {
     100            0 :         match proto_msg.r#type() {
     101              :             MessageType::SafekeeperTimelineInfo => Ok(Message::SafekeeperTimelineInfo(
     102            0 :                 proto_msg.safekeeper_timeline_info.ok_or_else(|| {
     103            0 :                     Status::new(Code::InvalidArgument, "missing safekeeper_timeline_info")
     104            0 :                 })?,
     105              :             )),
     106              :             MessageType::SafekeeperDiscoveryRequest => Ok(Message::SafekeeperDiscoveryRequest(
     107            0 :                 proto_msg.safekeeper_discovery_request.ok_or_else(|| {
     108            0 :                     Status::new(
     109            0 :                         Code::InvalidArgument,
     110            0 :                         "missing safekeeper_discovery_request",
     111            0 :                     )
     112            0 :                 })?,
     113              :             )),
     114              :             MessageType::SafekeeperDiscoveryResponse => Ok(Message::SafekeeperDiscoveryResponse(
     115            0 :                 proto_msg.safekeeper_discovery_response.ok_or_else(|| {
     116            0 :                     Status::new(
     117            0 :                         Code::InvalidArgument,
     118            0 :                         "missing safekeeper_discovery_response",
     119            0 :                     )
     120            0 :                 })?,
     121              :             )),
     122            0 :             MessageType::Unknown => Err(Status::new(
     123            0 :                 Code::InvalidArgument,
     124            0 :                 format!("invalid message type: {:?}", proto_msg.r#type),
     125            0 :             )),
     126              :         }
     127            0 :     }
     128              : 
     129              :     /// Get the tenant_timeline_id from the message.
     130            2 :     pub fn tenant_timeline_id(&self) -> Result<Option<TenantTimelineId>, Status> {
     131            2 :         match self {
     132            2 :             Message::SafekeeperTimelineInfo(msg) => Ok(msg
     133            2 :                 .tenant_timeline_id
     134            2 :                 .as_ref()
     135            2 :                 .map(parse_proto_ttid)
     136            2 :                 .transpose()?),
     137            0 :             Message::SafekeeperDiscoveryRequest(msg) => Ok(msg
     138            0 :                 .tenant_timeline_id
     139            0 :                 .as_ref()
     140            0 :                 .map(parse_proto_ttid)
     141            0 :                 .transpose()?),
     142            0 :             Message::SafekeeperDiscoveryResponse(msg) => Ok(msg
     143            0 :                 .tenant_timeline_id
     144            0 :                 .as_ref()
     145            0 :                 .map(parse_proto_ttid)
     146            0 :                 .transpose()?),
     147              :         }
     148            2 :     }
     149              : 
     150              :     /// Convert internal message to the protobuf struct.
     151            0 :     pub fn as_typed_message(&self) -> TypedMessage {
     152            0 :         let mut res = TypedMessage {
     153            0 :             r#type: self.message_type() as i32,
     154            0 :             ..Default::default()
     155            0 :         };
     156            0 :         match self {
     157            0 :             Message::SafekeeperTimelineInfo(msg) => {
     158            0 :                 res.safekeeper_timeline_info = Some(msg.clone())
     159              :             }
     160            0 :             Message::SafekeeperDiscoveryRequest(msg) => {
     161            0 :                 res.safekeeper_discovery_request = Some(msg.clone())
     162              :             }
     163            0 :             Message::SafekeeperDiscoveryResponse(msg) => {
     164            0 :                 res.safekeeper_discovery_response = Some(msg.clone())
     165              :             }
     166              :         }
     167            0 :         res
     168            0 :     }
     169              : 
     170              :     /// Get the message type.
     171            0 :     pub fn message_type(&self) -> MessageType {
     172            0 :         match self {
     173            0 :             Message::SafekeeperTimelineInfo(_) => MessageType::SafekeeperTimelineInfo,
     174            0 :             Message::SafekeeperDiscoveryRequest(_) => MessageType::SafekeeperDiscoveryRequest,
     175            0 :             Message::SafekeeperDiscoveryResponse(_) => MessageType::SafekeeperDiscoveryResponse,
     176              :         }
     177            0 :     }
     178              : }
     179              : 
     180              : #[derive(Copy, Clone, Debug)]
     181              : enum SubscriptionKey {
     182              :     All,
     183              :     Timeline(TenantTimelineId),
     184              : }
     185              : 
     186              : impl SubscriptionKey {
     187              :     /// Parse protobuf subkey (protobuf doesn't have fixed size bytes, we get vectors).
     188            0 :     pub fn from_proto_subscription_key(key: ProtoSubscriptionKey) -> Result<Self, Status> {
     189            0 :         match key {
     190            0 :             ProtoSubscriptionKey::All(_) => Ok(SubscriptionKey::All),
     191            0 :             ProtoSubscriptionKey::TenantTimelineId(proto_ttid) => {
     192            0 :                 Ok(SubscriptionKey::Timeline(parse_proto_ttid(&proto_ttid)?))
     193              :             }
     194              :         }
     195            0 :     }
     196              : 
     197              :     /// Parse from FilterTenantTimelineId
     198            0 :     pub fn from_proto_filter_tenant_timeline_id(
     199            0 :         opt: Option<&FilterTenantTimelineId>,
     200            0 :     ) -> Result<Self, Status> {
     201            0 :         if opt.is_none() {
     202            0 :             return Ok(SubscriptionKey::All);
     203            0 :         }
     204            0 : 
     205            0 :         let f = opt.unwrap();
     206            0 :         if !f.enabled {
     207            0 :             return Ok(SubscriptionKey::All);
     208            0 :         }
     209              : 
     210            0 :         let ttid =
     211            0 :             parse_proto_ttid(f.tenant_timeline_id.as_ref().ok_or_else(|| {
     212            0 :                 Status::new(Code::InvalidArgument, "missing tenant_timeline_id")
     213            0 :             })?)?;
     214            0 :         Ok(SubscriptionKey::Timeline(ttid))
     215            0 :     }
     216              : }
     217              : 
     218              : /// Channel to timeline subscribers.
     219              : struct ChanToTimelineSub {
     220              :     chan: broadcast::Sender<Message>,
     221              :     /// Tracked separately to know when delete the shmem entry. receiver_count()
     222              :     /// is unhandy for that as unregistering and dropping the receiver side
     223              :     /// happens at different moments.
     224              :     num_subscribers: u64,
     225              : }
     226              : 
     227              : struct SharedState {
     228              :     next_pub_id: PubId,
     229              :     num_pubs: i64,
     230              :     next_sub_id: SubId,
     231              :     num_subs_to_timelines: i64,
     232              :     chans_to_timeline_subs: HashMap<TenantTimelineId, ChanToTimelineSub>,
     233              :     num_subs_to_all: i64,
     234              :     chan_to_all_subs: broadcast::Sender<Message>,
     235              : }
     236              : 
     237              : impl SharedState {
     238            1 :     pub fn new(all_keys_chan_size: usize) -> Self {
     239            1 :         SharedState {
     240            1 :             next_pub_id: 0,
     241            1 :             num_pubs: 0,
     242            1 :             next_sub_id: 0,
     243            1 :             num_subs_to_timelines: 0,
     244            1 :             chans_to_timeline_subs: HashMap::new(),
     245            1 :             num_subs_to_all: 0,
     246            1 :             chan_to_all_subs: broadcast::channel(all_keys_chan_size).0,
     247            1 :         }
     248            1 :     }
     249              : 
     250              :     // Register new publisher.
     251            1 :     pub fn register_publisher(&mut self) -> PubId {
     252            1 :         let pub_id = self.next_pub_id;
     253            1 :         self.next_pub_id += 1;
     254            1 :         self.num_pubs += 1;
     255            1 :         NUM_PUBS.set(self.num_pubs);
     256            1 :         pub_id
     257            1 :     }
     258              : 
     259              :     // Unregister publisher.
     260            1 :     pub fn unregister_publisher(&mut self) {
     261            1 :         self.num_pubs -= 1;
     262            1 :         NUM_PUBS.set(self.num_pubs);
     263            1 :     }
     264              : 
     265              :     // Register new subscriber.
     266            2 :     pub fn register_subscriber(
     267            2 :         &mut self,
     268            2 :         sub_key: SubscriptionKey,
     269            2 :         timeline_chan_size: usize,
     270            2 :     ) -> (SubId, broadcast::Receiver<Message>) {
     271            2 :         let sub_id = self.next_sub_id;
     272            2 :         self.next_sub_id += 1;
     273            2 :         let sub_rx = match sub_key {
     274              :             SubscriptionKey::All => {
     275            1 :                 self.num_subs_to_all += 1;
     276            1 :                 NUM_SUBS_ALL.set(self.num_subs_to_all);
     277            1 :                 self.chan_to_all_subs.subscribe()
     278              :             }
     279            1 :             SubscriptionKey::Timeline(ttid) => {
     280            1 :                 self.num_subs_to_timelines += 1;
     281            1 :                 NUM_SUBS_TIMELINE.set(self.num_subs_to_timelines);
     282            1 :                 // Create new broadcast channel for this key, or subscriber to
     283            1 :                 // the existing one.
     284            1 :                 let chan_to_timeline_sub =
     285            1 :                     self.chans_to_timeline_subs
     286            1 :                         .entry(ttid)
     287            1 :                         .or_insert(ChanToTimelineSub {
     288            1 :                             chan: broadcast::channel(timeline_chan_size).0,
     289            1 :                             num_subscribers: 0,
     290            1 :                         });
     291            1 :                 chan_to_timeline_sub.num_subscribers += 1;
     292            1 :                 chan_to_timeline_sub.chan.subscribe()
     293              :             }
     294              :         };
     295            2 :         (sub_id, sub_rx)
     296            2 :     }
     297              : 
     298              :     // Unregister the subscriber.
     299            2 :     pub fn unregister_subscriber(&mut self, sub_key: SubscriptionKey) {
     300            2 :         match sub_key {
     301            1 :             SubscriptionKey::All => {
     302            1 :                 self.num_subs_to_all -= 1;
     303            1 :                 NUM_SUBS_ALL.set(self.num_subs_to_all);
     304            1 :             }
     305            1 :             SubscriptionKey::Timeline(ttid) => {
     306            1 :                 self.num_subs_to_timelines -= 1;
     307            1 :                 NUM_SUBS_TIMELINE.set(self.num_subs_to_timelines);
     308            1 : 
     309            1 :                 // Remove from the map, destroying the channel, if we are the
     310            1 :                 // last subscriber to this timeline.
     311            1 : 
     312            1 :                 // Missing entry is a bug; we must have registered.
     313            1 :                 let chan_to_timeline_sub = self
     314            1 :                     .chans_to_timeline_subs
     315            1 :                     .get_mut(&ttid)
     316            1 :                     .expect("failed to find sub entry in shmem during unregister");
     317            1 :                 chan_to_timeline_sub.num_subscribers -= 1;
     318            1 :                 if chan_to_timeline_sub.num_subscribers == 0 {
     319            1 :                     self.chans_to_timeline_subs.remove(&ttid);
     320            1 :                 }
     321              :             }
     322              :         }
     323            2 :     }
     324              : }
     325              : 
     326              : // SharedState wrapper.
     327              : #[derive(Clone)]
     328              : struct Registry {
     329              :     shared_state: Arc<RwLock<SharedState>>,
     330              :     timeline_chan_size: usize,
     331              : }
     332              : 
     333              : impl Registry {
     334              :     // Register new publisher in shared state.
     335            1 :     pub fn register_publisher(&self, remote_addr: SocketAddr) -> Publisher {
     336            1 :         let pub_id = self.shared_state.write().register_publisher();
     337            1 :         info!("publication started id={} addr={:?}", pub_id, remote_addr);
     338            1 :         Publisher {
     339            1 :             id: pub_id,
     340            1 :             registry: self.clone(),
     341            1 :             remote_addr,
     342            1 :         }
     343            1 :     }
     344              : 
     345            1 :     pub fn unregister_publisher(&self, publisher: &Publisher) {
     346            1 :         self.shared_state.write().unregister_publisher();
     347            1 :         info!(
     348            0 :             "publication ended id={} addr={:?}",
     349              :             publisher.id, publisher.remote_addr
     350              :         );
     351            1 :     }
     352              : 
     353              :     // Register new subscriber in shared state.
     354            2 :     pub fn register_subscriber(
     355            2 :         &self,
     356            2 :         sub_key: SubscriptionKey,
     357            2 :         remote_addr: SocketAddr,
     358            2 :     ) -> Subscriber {
     359            2 :         let (sub_id, sub_rx) = self
     360            2 :             .shared_state
     361            2 :             .write()
     362            2 :             .register_subscriber(sub_key, self.timeline_chan_size);
     363            2 :         info!(
     364            0 :             "subscription started id={}, key={:?}, addr={:?}",
     365              :             sub_id, sub_key, remote_addr
     366              :         );
     367            2 :         Subscriber {
     368            2 :             id: sub_id,
     369            2 :             key: sub_key,
     370            2 :             sub_rx,
     371            2 :             registry: self.clone(),
     372            2 :             remote_addr,
     373            2 :         }
     374            2 :     }
     375              : 
     376              :     // Unregister the subscriber
     377            2 :     pub fn unregister_subscriber(&self, subscriber: &Subscriber) {
     378            2 :         self.shared_state
     379            2 :             .write()
     380            2 :             .unregister_subscriber(subscriber.key);
     381            2 :         info!(
     382            0 :             "subscription ended id={}, key={:?}, addr={:?}",
     383              :             subscriber.id, subscriber.key, subscriber.remote_addr
     384              :         );
     385            2 :     }
     386              : 
     387              :     /// Send msg to relevant subscribers.
     388            2 :     pub fn send_msg(&self, msg: &Message) -> Result<(), Status> {
     389            2 :         PROCESSED_MESSAGES_TOTAL.inc();
     390            2 : 
     391            2 :         // send message to subscribers for everything
     392            2 :         let shared_state = self.shared_state.read();
     393            2 :         // Err means there is no subscribers, it is fine.
     394            2 :         shared_state.chan_to_all_subs.send(msg.clone()).ok();
     395              : 
     396              :         // send message to per timeline subscribers, if there is ttid
     397            2 :         let ttid = msg.tenant_timeline_id()?;
     398            2 :         if let Some(ttid) = ttid {
     399            2 :             if let Some(subs) = shared_state.chans_to_timeline_subs.get(&ttid) {
     400            1 :                 // Err can't happen here, as tx is destroyed only after removing
     401            1 :                 // from the map the last subscriber along with tx.
     402            1 :                 subs.chan
     403            1 :                     .send(msg.clone())
     404            1 :                     .expect("rx is still in the map with zero subscribers");
     405            1 :             }
     406            0 :         }
     407            2 :         Ok(())
     408            2 :     }
     409              : }
     410              : 
     411              : // Private subscriber state.
     412              : struct Subscriber {
     413              :     id: SubId,
     414              :     key: SubscriptionKey,
     415              :     // Subscriber receives messages from publishers here.
     416              :     sub_rx: broadcast::Receiver<Message>,
     417              :     // to unregister itself from shared state in Drop
     418              :     registry: Registry,
     419              :     // for logging
     420              :     remote_addr: SocketAddr,
     421              : }
     422              : 
     423              : impl Drop for Subscriber {
     424            2 :     fn drop(&mut self) {
     425            2 :         self.registry.unregister_subscriber(self);
     426            2 :     }
     427              : }
     428              : 
     429              : // Private publisher state
     430              : struct Publisher {
     431              :     id: PubId,
     432              :     registry: Registry,
     433              :     // for logging
     434              :     remote_addr: SocketAddr,
     435              : }
     436              : 
     437              : impl Publisher {
     438              :     /// Send msg to relevant subscribers.
     439            2 :     pub fn send_msg(&mut self, msg: &Message) -> Result<(), Status> {
     440            2 :         self.registry.send_msg(msg)
     441            2 :     }
     442              : }
     443              : 
     444              : impl Drop for Publisher {
     445            1 :     fn drop(&mut self) {
     446            1 :         self.registry.unregister_publisher(self);
     447            1 :     }
     448              : }
     449              : 
     450              : struct Broker {
     451              :     registry: Registry,
     452              : }
     453              : 
     454              : #[tonic::async_trait]
     455              : impl BrokerService for Broker {
     456            0 :     async fn publish_safekeeper_info(
     457            0 :         &self,
     458            0 :         request: Request<tonic::Streaming<SafekeeperTimelineInfo>>,
     459            0 :     ) -> Result<Response<()>, Status> {
     460            0 :         let &RemoteAddr(remote_addr) = request
     461            0 :             .extensions()
     462            0 :             .get()
     463            0 :             .expect("RemoteAddr inserted by handler");
     464            0 :         let mut publisher = self.registry.register_publisher(remote_addr);
     465            0 : 
     466            0 :         let mut stream = request.into_inner();
     467              : 
     468              :         loop {
     469            0 :             match stream.next().await {
     470            0 :                 Some(Ok(msg)) => publisher.send_msg(&Message::SafekeeperTimelineInfo(msg))?,
     471            0 :                 Some(Err(e)) => return Err(e), // grpc error from the stream
     472            0 :                 None => break,                 // closed stream
     473            0 :             }
     474            0 :         }
     475            0 : 
     476            0 :         Ok(Response::new(()))
     477            0 :     }
     478              : 
     479              :     type SubscribeSafekeeperInfoStream =
     480              :         Pin<Box<dyn Stream<Item = Result<SafekeeperTimelineInfo, Status>> + Send + 'static>>;
     481              : 
     482            0 :     async fn subscribe_safekeeper_info(
     483            0 :         &self,
     484            0 :         request: Request<SubscribeSafekeeperInfoRequest>,
     485            0 :     ) -> Result<Response<Self::SubscribeSafekeeperInfoStream>, Status> {
     486            0 :         let &RemoteAddr(remote_addr) = request
     487            0 :             .extensions()
     488            0 :             .get()
     489            0 :             .expect("RemoteAddr inserted by handler");
     490            0 :         let proto_key = request
     491            0 :             .into_inner()
     492            0 :             .subscription_key
     493            0 :             .ok_or_else(|| Status::new(Code::InvalidArgument, "missing subscription key"))?;
     494            0 :         let sub_key = SubscriptionKey::from_proto_subscription_key(proto_key)?;
     495            0 :         let mut subscriber = self.registry.register_subscriber(sub_key, remote_addr);
     496            0 : 
     497            0 :         // transform rx into stream with item = Result, as method result demands
     498            0 :         let output = async_stream::try_stream! {
     499            0 :             let mut warn_interval = time::interval(Duration::from_millis(1000));
     500            0 :             let mut missed_msgs: u64 = 0;
     501            0 :             loop {
     502            0 :                 match subscriber.sub_rx.recv().await {
     503            0 :                     Ok(info) => {
     504            0 :                         match info {
     505            0 :                             Message::SafekeeperTimelineInfo(info) => yield info,
     506            0 :                             _ => {},
     507            0 :                         }
     508            0 :                         BROADCASTED_MESSAGES_TOTAL.inc();
     509            0 :                     },
     510            0 :                     Err(RecvError::Lagged(skipped_msg)) => {
     511            0 :                         BROADCAST_DROPPED_MESSAGES_TOTAL.inc_by(skipped_msg);
     512            0 :                         missed_msgs += skipped_msg;
     513            0 :                         if (futures::poll!(Box::pin(warn_interval.tick()))).is_ready() {
     514            0 :                             warn!("subscription id={}, key={:?} addr={:?} dropped {} messages, channel is full",
     515            0 :                                 subscriber.id, subscriber.key, subscriber.remote_addr, missed_msgs);
     516            0 :                             missed_msgs = 0;
     517            0 :                         }
     518            0 :                     }
     519            0 :                     Err(RecvError::Closed) => {
     520            0 :                         // can't happen, we never drop the channel while there is a subscriber
     521            0 :                         Err(Status::new(Code::Internal, "channel unexpectantly closed"))?;
     522            0 :                     }
     523            0 :                 }
     524            0 :             }
     525            0 :         };
     526            0 : 
     527            0 :         Ok(Response::new(
     528            0 :             Box::pin(output) as Self::SubscribeSafekeeperInfoStream
     529            0 :         ))
     530            0 :     }
     531              : 
     532              :     type SubscribeByFilterStream =
     533              :         Pin<Box<dyn Stream<Item = Result<TypedMessage, Status>> + Send + 'static>>;
     534              : 
     535              :     /// Subscribe to all messages, limited by a filter.
     536            0 :     async fn subscribe_by_filter(
     537            0 :         &self,
     538            0 :         request: Request<SubscribeByFilterRequest>,
     539            0 :     ) -> std::result::Result<Response<Self::SubscribeByFilterStream>, Status> {
     540            0 :         let &RemoteAddr(remote_addr) = request
     541            0 :             .extensions()
     542            0 :             .get()
     543            0 :             .expect("RemoteAddr inserted by handler");
     544            0 :         let proto_filter = request.into_inner();
     545            0 :         let ttid_filter = proto_filter.tenant_timeline_id.as_ref();
     546              : 
     547            0 :         let sub_key = SubscriptionKey::from_proto_filter_tenant_timeline_id(ttid_filter)?;
     548            0 :         let types_set = proto_filter
     549            0 :             .types
     550            0 :             .iter()
     551            0 :             .map(|t| t.r#type)
     552            0 :             .collect::<std::collections::HashSet<_>>();
     553            0 : 
     554            0 :         let mut subscriber = self.registry.register_subscriber(sub_key, remote_addr);
     555            0 : 
     556            0 :         // transform rx into stream with item = Result, as method result demands
     557            0 :         let output = async_stream::try_stream! {
     558            0 :             let mut warn_interval = time::interval(Duration::from_millis(1000));
     559            0 :             let mut missed_msgs: u64 = 0;
     560            0 :             loop {
     561            0 :                 match subscriber.sub_rx.recv().await {
     562            0 :                     Ok(msg) => {
     563            0 :                         let msg_type = msg.message_type() as i32;
     564            0 :                         if types_set.contains(&msg_type) {
     565            0 :                             yield msg.as_typed_message();
     566            0 :                             BROADCASTED_MESSAGES_TOTAL.inc();
     567            0 :                         }
     568            0 :                     },
     569            0 :                     Err(RecvError::Lagged(skipped_msg)) => {
     570            0 :                         BROADCAST_DROPPED_MESSAGES_TOTAL.inc_by(skipped_msg);
     571            0 :                         missed_msgs += skipped_msg;
     572            0 :                         if (futures::poll!(Box::pin(warn_interval.tick()))).is_ready() {
     573            0 :                             warn!("subscription id={}, key={:?} addr={:?} dropped {} messages, channel is full",
     574            0 :                                 subscriber.id, subscriber.key, subscriber.remote_addr, missed_msgs);
     575            0 :                             missed_msgs = 0;
     576            0 :                         }
     577            0 :                     }
     578            0 :                     Err(RecvError::Closed) => {
     579            0 :                         // can't happen, we never drop the channel while there is a subscriber
     580            0 :                         Err(Status::new(Code::Internal, "channel unexpectantly closed"))?;
     581            0 :                     }
     582            0 :                 }
     583            0 :             }
     584            0 :         };
     585            0 : 
     586            0 :         Ok(Response::new(
     587            0 :             Box::pin(output) as Self::SubscribeByFilterStream
     588            0 :         ))
     589            0 :     }
     590              : 
     591              :     /// Publish one message.
     592            0 :     async fn publish_one(
     593            0 :         &self,
     594            0 :         request: Request<TypedMessage>,
     595            0 :     ) -> std::result::Result<Response<()>, Status> {
     596            0 :         let msg = Message::from(request.into_inner())?;
     597            0 :         PUBLISHED_ONEOFF_MESSAGES_TOTAL.inc();
     598            0 :         self.registry.send_msg(&msg)?;
     599            0 :         Ok(Response::new(()))
     600            0 :     }
     601              : }
     602              : 
     603              : // We serve only metrics and healthcheck through http1.
     604            0 : async fn http1_handler(
     605            0 :     req: hyper::Request<Incoming>,
     606            0 : ) -> Result<hyper::Response<BoxBody>, Infallible> {
     607            0 :     let resp = match (req.method(), req.uri().path()) {
     608            0 :         (&Method::GET, "/metrics") => {
     609            0 :             let mut buffer = vec![];
     610            0 :             let metrics = metrics::gather();
     611            0 :             let encoder = TextEncoder::new();
     612            0 :             encoder.encode(&metrics, &mut buffer).unwrap();
     613            0 : 
     614            0 :             hyper::Response::builder()
     615            0 :                 .status(StatusCode::OK)
     616            0 :                 .header(CONTENT_TYPE, encoder.format_type())
     617            0 :                 .body(body::boxed(Full::new(bytes::Bytes::from(buffer))))
     618            0 :                 .unwrap()
     619              :         }
     620            0 :         (&Method::GET, "/status") => hyper::Response::builder()
     621            0 :             .status(StatusCode::OK)
     622            0 :             .body(empty_body())
     623            0 :             .unwrap(),
     624            0 :         _ => hyper::Response::builder()
     625            0 :             .status(StatusCode::NOT_FOUND)
     626            0 :             .body(empty_body())
     627            0 :             .unwrap(),
     628              :     };
     629            0 :     Ok(resp)
     630            0 : }
     631              : 
     632              : #[derive(Clone, Copy)]
     633              : struct RemoteAddr(SocketAddr);
     634              : 
     635              : #[tokio::main]
     636            0 : async fn main() -> Result<(), Box<dyn std::error::Error>> {
     637            0 :     let args = Args::parse();
     638            0 : 
     639            0 :     // important to keep the order of:
     640            0 :     // 1. init logging
     641            0 :     // 2. tracing panic hook
     642            0 :     // 3. sentry
     643            0 :     logging::init(
     644            0 :         LogFormat::from_config(&args.log_format)?,
     645            0 :         logging::TracingErrorLayerEnablement::Disabled,
     646            0 :         logging::Output::Stdout,
     647            0 :     )?;
     648            0 :     logging::replace_panic_hook_with_tracing_panic_hook().forget();
     649            0 :     // initialize sentry if SENTRY_DSN is provided
     650            0 :     let _sentry_guard = init_sentry(Some(GIT_VERSION.into()), &[]);
     651            0 :     info!("version: {GIT_VERSION} build_tag: {BUILD_TAG}");
     652            0 :     metrics::set_build_info_metric(GIT_VERSION, BUILD_TAG);
     653            0 : 
     654            0 :     // On any shutdown signal, log receival and exit.
     655            0 :     std::thread::spawn(move || {
     656            0 :         ShutdownSignals::handle(|signal| {
     657            0 :             info!("received {}, terminating", signal.name());
     658            0 :             std::process::exit(0);
     659            0 :         })
     660            0 :     });
     661            0 : 
     662            0 :     let registry = Registry {
     663            0 :         shared_state: Arc::new(RwLock::new(SharedState::new(args.all_keys_chan_size))),
     664            0 :         timeline_chan_size: args.timeline_chan_size,
     665            0 :     };
     666            0 :     let storage_broker_impl = Broker {
     667            0 :         registry: registry.clone(),
     668            0 :     };
     669            0 :     let storage_broker_server = BrokerServiceServer::new(storage_broker_impl);
     670            0 : 
     671            0 :     // grpc is served along with http1 for metrics on a single port, hence we
     672            0 :     // don't use tonic's Server.
     673            0 :     let tcp_listener = TcpListener::bind(&args.listen_addr).await?;
     674            0 :     info!("listening on {}", &args.listen_addr);
     675            0 :     loop {
     676            0 :         let (stream, addr) = match tcp_listener.accept().await {
     677            0 :             Ok(v) => v,
     678            0 :             Err(e) => {
     679            0 :                 info!("couldn't accept connection: {e}");
     680            0 :                 continue;
     681            0 :             }
     682            0 :         };
     683            0 : 
     684            0 :         let mut builder = hyper_util::server::conn::auto::Builder::new(TokioExecutor::new());
     685            0 :         builder.http1().timer(TokioTimer::new());
     686            0 :         builder
     687            0 :             .http2()
     688            0 :             .timer(TokioTimer::new())
     689            0 :             .keep_alive_interval(Some(args.http2_keepalive_interval))
     690            0 :             // This matches the tonic server default. It allows us to support production-like workloads.
     691            0 :             .max_concurrent_streams(None);
     692            0 : 
     693            0 :         let storage_broker_server_cloned = storage_broker_server.clone();
     694            0 :         let remote_addr = RemoteAddr(addr);
     695            0 :         let service_fn_ = async move {
     696            0 :             service_fn(move |mut req| {
     697            0 :                 // That's what tonic's MakeSvc.call does to pass conninfo to
     698            0 :                 // the request handler (and where its request.remote_addr()
     699            0 :                 // expects it to find).
     700            0 :                 req.extensions_mut().insert(remote_addr);
     701            0 : 
     702            0 :                 // Technically this second clone is not needed, but consume
     703            0 :                 // by async block is apparently unavoidable. BTW, error
     704            0 :                 // message is enigmatic, see
     705            0 :                 // https://github.com/rust-lang/rust/issues/68119
     706            0 :                 //
     707            0 :                 // We could get away without async block at all, but then we
     708            0 :                 // need to resort to futures::Either to merge the result,
     709            0 :                 // which doesn't caress an eye as well.
     710            0 :                 let mut storage_broker_server_svc = storage_broker_server_cloned.clone();
     711            0 :                 async move {
     712            0 :                     if req.headers().get("content-type").map(|x| x.as_bytes())
     713            0 :                         == Some(b"application/grpc")
     714            0 :                     {
     715            0 :                         let res_resp = storage_broker_server_svc.call(req).await;
     716            0 :                         // Grpc and http1 handlers have slightly different
     717            0 :                         // Response types: it is UnsyncBoxBody for the
     718            0 :                         // former one (not sure why) and plain hyper::Body
     719            0 :                         // for the latter. Both implement HttpBody though,
     720            0 :                         // and `Either` is used to merge them.
     721            0 :                         res_resp.map(|resp| resp.map(http_body_util::Either::Left))
     722            0 :                     } else {
     723            0 :                         let res_resp = http1_handler(req).await;
     724            0 :                         res_resp.map(|resp| resp.map(http_body_util::Either::Right))
     725            0 :                     }
     726            0 :                 }
     727            0 :             })
     728            0 :         }
     729            0 :         .await;
     730            0 : 
     731            0 :         tokio::task::spawn(async move {
     732            0 :             let res = builder
     733            0 :                 .serve_connection(TokioIo::new(stream), service_fn_)
     734            0 :                 .await;
     735            0 : 
     736            0 :             if let Err(e) = res {
     737            0 :                 info!("error serving connection from {addr}: {e}");
     738            0 :             }
     739            0 :         });
     740            0 :     }
     741            0 : }
     742              : 
     743              : #[cfg(test)]
     744              : mod tests {
     745              :     use storage_broker::proto::TenantTimelineId as ProtoTenantTimelineId;
     746              :     use tokio::sync::broadcast::error::TryRecvError;
     747              :     use utils::id::{TenantId, TimelineId};
     748              : 
     749              :     use super::*;
     750              : 
     751            2 :     fn msg(timeline_id: Vec<u8>) -> Message {
     752            2 :         Message::SafekeeperTimelineInfo(SafekeeperTimelineInfo {
     753            2 :             safekeeper_id: 1,
     754            2 :             tenant_timeline_id: Some(ProtoTenantTimelineId {
     755            2 :                 tenant_id: vec![0x00; 16],
     756            2 :                 timeline_id,
     757            2 :             }),
     758            2 :             term: 0,
     759            2 :             last_log_term: 0,
     760            2 :             flush_lsn: 1,
     761            2 :             commit_lsn: 2,
     762            2 :             backup_lsn: 3,
     763            2 :             remote_consistent_lsn: 4,
     764            2 :             peer_horizon_lsn: 5,
     765            2 :             safekeeper_connstr: "neon-1-sk-1.local:7676".to_owned(),
     766            2 :             http_connstr: "neon-1-sk-1.local:7677".to_owned(),
     767            2 :             local_start_lsn: 0,
     768            2 :             availability_zone: None,
     769            2 :             standby_horizon: 0,
     770            2 :         })
     771            2 :     }
     772              : 
     773            3 :     fn tli_from_u64(i: u64) -> Vec<u8> {
     774            3 :         let mut timeline_id = vec![0xFF; 8];
     775            3 :         timeline_id.extend_from_slice(&i.to_be_bytes());
     776            3 :         timeline_id
     777            3 :     }
     778              : 
     779            3 :     fn mock_addr() -> SocketAddr {
     780            3 :         "127.0.0.1:8080".parse().unwrap()
     781            3 :     }
     782              : 
     783              :     #[tokio::test]
     784            1 :     async fn test_registry() {
     785            1 :         let registry = Registry {
     786            1 :             shared_state: Arc::new(RwLock::new(SharedState::new(16))),
     787            1 :             timeline_chan_size: 16,
     788            1 :         };
     789            1 : 
     790            1 :         // subscribe to timeline 2
     791            1 :         let ttid_2 = TenantTimelineId {
     792            1 :             tenant_id: TenantId::from_slice(&[0x00; 16]).unwrap(),
     793            1 :             timeline_id: TimelineId::from_slice(&tli_from_u64(2)).unwrap(),
     794            1 :         };
     795            1 :         let sub_key_2 = SubscriptionKey::Timeline(ttid_2);
     796            1 :         let mut subscriber_2 = registry.register_subscriber(sub_key_2, mock_addr());
     797            1 :         let mut subscriber_all = registry.register_subscriber(SubscriptionKey::All, mock_addr());
     798            1 : 
     799            1 :         // send two messages with different keys
     800            1 :         let msg_1 = msg(tli_from_u64(1));
     801            1 :         let msg_2 = msg(tli_from_u64(2));
     802            1 :         let mut publisher = registry.register_publisher(mock_addr());
     803            1 :         publisher.send_msg(&msg_1).expect("failed to send msg");
     804            1 :         publisher.send_msg(&msg_2).expect("failed to send msg");
     805            1 : 
     806            1 :         // msg with key 2 should arrive to subscriber_2
     807            1 :         assert_eq!(subscriber_2.sub_rx.try_recv().unwrap(), msg_2);
     808            1 : 
     809            1 :         // but nothing more
     810            1 :         assert_eq!(
     811            1 :             subscriber_2.sub_rx.try_recv().unwrap_err(),
     812            1 :             TryRecvError::Empty
     813            1 :         );
     814            1 : 
     815            1 :         // subscriber_all should receive both messages
     816            1 :         assert_eq!(subscriber_all.sub_rx.try_recv().unwrap(), msg_1);
     817            1 :         assert_eq!(subscriber_all.sub_rx.try_recv().unwrap(), msg_2);
     818            1 :         assert_eq!(
     819            1 :             subscriber_all.sub_rx.try_recv().unwrap_err(),
     820            1 :             TryRecvError::Empty
     821            1 :         );
     822            1 :     }
     823              : }
        

Generated by: LCOV version 2.1-beta