reifydb_sub_server/
accept.rs

1// SPDX-License-Identifier: Apache-2.0
2// Copyright (c) 2026 ReifyDB
3
4use std::{
5	io,
6	net::SocketAddr,
7	pin::Pin,
8	sync::Arc,
9	task::{Context, Poll},
10};
11
12use libc::{EMFILE, ENFILE};
13use reifydb_value::value::duration::Duration;
14use tokio::{
15	io::{AsyncRead, AsyncWrite, ReadBuf},
16	net::{TcpListener, TcpStream},
17	sync::{OwnedSemaphorePermit, Semaphore},
18	time::sleep,
19};
20use tracing::warn;
21
22pub const BACKOFF_MIN: Duration = Duration::from_milliseconds_const(5);
23
24pub const BACKOFF_MAX: Duration = Duration::from_milliseconds_const(1000);
25
26pub fn is_fd_exhausted(e: &io::Error) -> bool {
27	matches!(e.raw_os_error(), Some(code) if code == EMFILE || code == ENFILE)
28}
29
30pub async fn backoff_after_accept_error(e: &io::Error, backoff: &mut Duration, name: &str) {
31	if is_fd_exhausted(e) {
32		warn!("{name}: accept failed, file descriptors exhausted: {e}; retrying in {backoff:?}");
33	} else {
34		warn!("{name}: accept error: {e}; retrying in {backoff:?}");
35	}
36	sleep(backoff.to_std()).await;
37	*backoff = backoff.try_mul(2).map_or(BACKOFF_MAX, |doubled| doubled.min(BACKOFF_MAX));
38}
39
40pub async fn accept_admitted(
41	listener: &TcpListener,
42	semaphore: &Arc<Semaphore>,
43	name: &str,
44) -> (TcpStream, OwnedSemaphorePermit, SocketAddr) {
45	let mut backoff = BACKOFF_MIN;
46	loop {
47		match listener.accept().await {
48			Ok((stream, peer)) => match Arc::clone(semaphore).try_acquire_owned() {
49				Ok(permit) => {
50					if let Err(e) = stream.set_nodelay(true) {
51						warn!("{name}: failed to set TCP_NODELAY: {e}");
52					}
53					return (stream, permit, peer);
54				}
55				Err(_) => {
56					warn!("{name}: connection limit reached, rejecting {peer}");
57					drop(stream);
58				}
59			},
60			Err(e) => backoff_after_accept_error(&e, &mut backoff, name).await,
61		}
62	}
63}
64
65pub struct PermittedStream {
66	inner: TcpStream,
67	_permit: OwnedSemaphorePermit,
68}
69
70impl PermittedStream {
71	pub fn new(inner: TcpStream, permit: OwnedSemaphorePermit) -> Self {
72		Self {
73			inner,
74			_permit: permit,
75		}
76	}
77}
78
79impl AsyncRead for PermittedStream {
80	fn poll_read(mut self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut ReadBuf<'_>) -> Poll<io::Result<()>> {
81		Pin::new(&mut self.inner).poll_read(cx, buf)
82	}
83}
84
85impl AsyncWrite for PermittedStream {
86	fn poll_write(mut self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &[u8]) -> Poll<io::Result<usize>> {
87		Pin::new(&mut self.inner).poll_write(cx, buf)
88	}
89
90	fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
91		Pin::new(&mut self.inner).poll_flush(cx)
92	}
93
94	fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
95		Pin::new(&mut self.inner).poll_shutdown(cx)
96	}
97
98	fn poll_write_vectored(
99		mut self: Pin<&mut Self>,
100		cx: &mut Context<'_>,
101		bufs: &[io::IoSlice<'_>],
102	) -> Poll<io::Result<usize>> {
103		Pin::new(&mut self.inner).poll_write_vectored(cx, bufs)
104	}
105
106	fn is_write_vectored(&self) -> bool {
107		self.inner.is_write_vectored()
108	}
109}
110
111#[cfg(test)]
112mod tests {
113	use std::io::Error;
114
115	use libc::ECONNABORTED;
116
117	use super::*;
118
119	#[test]
120	fn detects_fd_exhaustion_errors() {
121		assert!(is_fd_exhausted(&Error::from_raw_os_error(EMFILE)));
122		assert!(is_fd_exhausted(&Error::from_raw_os_error(ENFILE)));
123	}
124
125	#[test]
126	fn ignores_other_errors() {
127		assert!(!is_fd_exhausted(&Error::from_raw_os_error(ECONNABORTED)));
128		// An error with no OS code (e.g. a synthesized one) must not be treated as fd exhaustion.
129		assert!(!is_fd_exhausted(&Error::other("boom")));
130	}
131}
reifydb_sub_server/accept.rs

reifydb_sub_server/
accept.rs
