TLA  Line data    Source code

       1                 : //!
       2                 : //!   WAL service listens for client connections and
       3                 : //!   receive WAL from wal_proposer and send it to WAL receivers
       4                 : //!
       5                 : use anyhow::{Context, Result};
       6                 : use postgres_backend::QueryError;
       7                 : use std::{future, time::Duration};
       8                 : use tokio::net::TcpStream;
       9                 : use tokio_io_timeout::TimeoutReader;
      10                 : use tracing::*;
      11                 : use utils::{auth::Scope, measured_stream::MeasuredStream};
      12                 : 
      13                 : use crate::handler::SafekeeperPostgresHandler;
      14                 : use crate::metrics::TrafficMetrics;
      15                 : use crate::SafeKeeperConf;
      16                 : use postgres_backend::{AuthType, PostgresBackend};
      17                 : 
      18                 : /// Accept incoming TCP connections and spawn them into a background thread.
      19                 : /// allowed_auth_scope is either SafekeeperData (wide JWT tokens giving access
      20                 : /// to any tenant are allowed) or Tenant (only tokens giving access to specific
      21                 : /// tenant are allowed). Doesn't matter if auth is disabled in conf.
      22 CBC        1000 : pub async fn task_main(
      23            1000 :     conf: SafeKeeperConf,
      24            1000 :     pg_listener: std::net::TcpListener,
      25            1000 :     allowed_auth_scope: Scope,
      26            1000 : ) -> anyhow::Result<()> {
      27                 :     // Tokio's from_std won't do this for us, per its comment.
      28            1000 :     pg_listener.set_nonblocking(true)?;
      29                 : 
      30            1000 :     let listener = tokio::net::TcpListener::from_std(pg_listener)?;
      31            1000 :     let mut connection_count: ConnectionCount = 0;
      32                 : 
      33                 :     loop {
      34            4561 :         let (socket, peer_addr) = listener.accept().await.context("accept")?;
      35 UBC           0 :         debug!("accepted connection from {}", peer_addr);
      36 CBC        3561 :         let conf = conf.clone();
      37            3561 :         let conn_id = issue_connection_id(&mut connection_count);
      38            3561 : 
      39            3561 :         tokio::spawn(async move {
      40            3561 :             if let Err(err) = handle_socket(socket, conf, conn_id, allowed_auth_scope)
      41            3561 :                 .instrument(info_span!("", cid = %conn_id))
      42         8036854 :                 .await
      43                 :             {
      44              59 :                 error!("connection handler exited: {}", err);
      45            3105 :             }
      46            3561 :         });
      47                 :     }
      48 UBC           0 : }
      49                 : 
      50                 : /// This is run by `task_main` above, inside a background thread.
      51                 : ///
      52 CBC        3561 : async fn handle_socket(
      53            3561 :     socket: TcpStream,
      54            3561 :     conf: SafeKeeperConf,
      55            3561 :     conn_id: ConnectionId,
      56            3561 :     allowed_auth_scope: Scope,
      57            3561 : ) -> Result<(), QueryError> {
      58            3561 :     socket.set_nodelay(true)?;
      59            3561 :     let peer_addr = socket.peer_addr()?;
      60                 : 
      61                 :     // Set timeout on reading from the socket. It prevents hanged up connection
      62                 :     // if client suddenly disappears. Note that TCP_KEEPALIVE is not enabled by
      63                 :     // default, and tokio doesn't provide ability to set it out of the box.
      64            3561 :     let mut socket = TimeoutReader::new(socket);
      65            3561 :     let wal_service_timeout = Duration::from_secs(60 * 10);
      66            3561 :     socket.set_timeout(Some(wal_service_timeout));
      67            3561 :     // pin! is here because TimeoutReader (due to storing sleep future inside)
      68            3561 :     // is not Unpin, and all pgbackend/framed/tokio dependencies require stream
      69            3561 :     // to be Unpin. Which is reasonable, as indeed something like TimeoutReader
      70            3561 :     // shouldn't be moved.
      71            3561 :     tokio::pin!(socket);
      72            3561 : 
      73            3561 :     let traffic_metrics = TrafficMetrics::new();
      74            3561 :     if let Some(current_az) = conf.availability_zone.as_deref() {
      75            3561 :         traffic_metrics.set_sk_az(current_az);
      76            3561 :     }
      77                 : 
      78            3561 :     let socket = MeasuredStream::new(
      79            3561 :         socket,
      80         3539877 :         |cnt| {
      81         3539877 :             traffic_metrics.observe_read(cnt);
      82         3539877 :         },
      83         3265120 :         |cnt| {
      84         3265120 :             traffic_metrics.observe_write(cnt);
      85         3265120 :         },
      86            3561 :     );
      87                 : 
      88            3561 :     let auth_key = match allowed_auth_scope {
      89            2854 :         Scope::Tenant => conf.pg_tenant_only_auth.clone(),
      90             707 :         _ => conf.pg_auth.clone(),
      91                 :     };
      92            3561 :     let auth_type = match auth_key {
      93            3418 :         None => AuthType::Trust,
      94             143 :         Some(_) => AuthType::NeonJWT,
      95                 :     };
      96            3561 :     let auth_pair = auth_key.map(|key| (allowed_auth_scope, key));
      97            3561 :     let mut conn_handler =
      98            3561 :         SafekeeperPostgresHandler::new(conf, conn_id, Some(traffic_metrics.clone()), auth_pair);
      99            3561 :     let pgbackend = PostgresBackend::new_from_io(socket, peer_addr, auth_type, None)?;
     100                 :     // libpq protocol between safekeeper and walproposer / pageserver
     101                 :     // We don't use shutdown.
     102            3561 :     pgbackend
     103            3561 :         .run(&mut conn_handler, future::pending::<()>)
     104         8036854 :         .await
     105            3164 : }
     106                 : 
     107                 : /// Unique WAL service connection ids are logged in spans for observability.
     108                 : pub type ConnectionId = u32;
     109                 : pub type ConnectionCount = u32;
     110                 : 
     111            3561 : pub fn issue_connection_id(count: &mut ConnectionCount) -> ConnectionId {
     112            3561 :     *count = count.wrapping_add(1);
     113            3561 :     *count
     114            3561 : }