tako-rs-server 2.5.0

Internal server bootstrap for tako-rs. Use the `tako-rs` umbrella crate instead.
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132

//! Raw Unix domain socket serve loops on the compio runtime.

use std::future::Future;
use std::path::Path;
use std::pin::Pin;
use std::sync::Arc;
use std::time::Duration;

use futures_util::FutureExt;
use futures_util::StreamExt;
use futures_util::future::Either;
use futures_util::stream::FuturesUnordered;
use tako_rs_core::conn_info::UnixPeerAddr;

use crate::server_compio::unix::bind_listener;
use crate::server_compio::unix::peer_addr;

/// Runs `handler` for every accepted connection until an accept error, which it returns.
pub async fn serve_unix<F>(path: impl AsRef<Path>, handler: F) -> std::io::Result<()>
where
  F: Fn(compio::net::UnixStream, UnixPeerAddr) -> Pin<Box<dyn Future<Output = std::io::Result<()>>>>
    + Send
    + Sync
    + 'static,
{
  let path = path.as_ref();
  let listener = bind_listener(path).await?;
  tracing::info!("Unix socket server listening on {}", path.display());

  #[cfg(feature = "signals")]
  tako_rs_core::signals::transport::emit_server_started(&path.to_string_lossy(), "unix", false)
    .await;

  let handler = Arc::new(handler);

  loop {
    let (stream, addr) = listener.accept().await?;
    let peer = peer_addr(addr.as_pathname());
    let handler = Arc::clone(&handler);

    compio::runtime::spawn(async move {
      if let Err(e) = handler(stream, peer).await {
        tracing::error!("Unix socket connection error: {e}");
      }
    })
    .detach();
  }
}

/// Stops accepting once `signal` resolves and drains connections for up to 30 seconds; use
/// [`serve_unix_with_shutdown_and_drain`] for another bound.
pub async fn serve_unix_with_shutdown<F, S>(
  path: impl AsRef<Path>,
  handler: F,
  signal: S,
) -> std::io::Result<()>
where
  F: Fn(compio::net::UnixStream, UnixPeerAddr) -> Pin<Box<dyn Future<Output = std::io::Result<()>>>>
    + Send
    + Sync
    + 'static,
  S: Future<Output = ()> + 'static,
{
  serve_unix_with_shutdown_and_drain(path, handler, signal, Duration::from_secs(30)).await
}

/// Like [`serve_unix_with_shutdown`] with an explicit drain timeout.
pub async fn serve_unix_with_shutdown_and_drain<F, S>(
  path: impl AsRef<Path>,
  handler: F,
  signal: S,
  drain_timeout: Duration,
) -> std::io::Result<()>
where
  F: Fn(compio::net::UnixStream, UnixPeerAddr) -> Pin<Box<dyn Future<Output = std::io::Result<()>>>>
    + Send
    + Sync
    + 'static,
  S: Future<Output = ()> + 'static,
{
  let path = path.as_ref();
  let listener = bind_listener(path).await?;
  tracing::info!("Unix socket server listening on {}", path.display());

  #[cfg(feature = "signals")]
  tako_rs_core::signals::transport::emit_server_started(&path.to_string_lossy(), "unix", false)
    .await;

  let handler = Arc::new(handler);
  let mut connections = FuturesUnordered::new();
  let mut signal = std::pin::pin!(signal);

  loop {
    let accept = std::pin::pin!(listener.accept());
    match futures_util::future::select(accept, &mut signal).await {
      Either::Left((result, _)) => {
        let (stream, addr) = result?;
        let peer = peer_addr(addr.as_pathname());
        let handler = Arc::clone(&handler);
        connections.push(compio::runtime::spawn(async move {
          if let Err(e) = handler(stream, peer).await {
            tracing::error!("Unix socket connection error: {e}");
          }
        }));
        while connections.next().now_or_never().flatten().is_some() {}
      }
      Either::Right(_) => {
        tracing::info!(
          "Unix socket server shutting down, draining {} connections",
          connections.len()
        );
        break;
      }
    }
  }

  if compio::time::timeout(drain_timeout, async {
    while connections.next().await.is_some() {}
  })
  .await
  .is_err()
  {
    for connection in connections {
      connection.cancel().await;
    }
  }

  #[cfg(feature = "signals")]
  tako_rs_core::signals::transport::emit_server_stopped(&path.to_string_lossy(), "unix", false)
    .await;
  Ok(())
}