use std::{
io,
net::{TcpListener, TcpStream},
sync::{
Arc,
atomic::{
AtomicU64,
Ordering::{AcqRel, Acquire, Release},
},
},
thread::sleep,
time::Duration,
};
use parking_lot::Mutex as ParkingMutex;
use super::{
garnet_server_base::GarnetServerBase,
i_garnet_server::{ClusterSessionFace, MessageConsumerFace, ServerEnumerate, WireFormat},
};
pub const INITIAL_ACCEPT_BACKOFF_MS: u64 = 100;
pub const MAX_ACCEPT_BACKOFF_MS: u64 = 5000;
pub const LISTEN_BACKLOG: i32 = 512;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum AcceptErrorTier {
FatalClean,
Backoff,
Transient,
}
pub struct AcceptedConnection {
pub handler_id: u64,
pub stream: TcpStream,
pub remote_endpoint: String,
}
pub struct GarnetServerTcp {
base: Arc<GarnetServerBase>,
listener: ParkingMutex<Option<TcpListener>>,
network_connection_limit: i32,
network_send_throttle_max: usize,
accept_backoff_ms: AtomicU64,
}
impl GarnetServerTcp {
pub fn new(
endpoint: &str,
network_buffer_size: usize,
network_send_throttle_max: usize,
network_connection_limit: i32,
) -> Self {
Self {
base: Arc::new(GarnetServerBase::new(endpoint, network_buffer_size)),
listener: ParkingMutex::new(None),
network_connection_limit,
network_send_throttle_max: network_send_throttle_max.max(1),
accept_backoff_ms: AtomicU64::new(INITIAL_ACCEPT_BACKOFF_MS),
}
}
pub fn base(&self) -> &Arc<GarnetServerBase> {
&self.base
}
pub fn network_send_throttle_max(&self) -> usize {
self.network_send_throttle_max
}
pub fn accept_backoff_ms(&self) -> u64 {
self.accept_backoff_ms.load(Acquire)
}
pub fn active_consumers(&self) -> Vec<Arc<dyn MessageConsumerFace>> {
self.base.active_consumers()
}
pub fn active_cluster_sessions(&self) -> Vec<Arc<dyn ClusterSessionFace>> {
self.base.active_cluster_sessions()
}
pub fn start(&self) -> io::Result<()> {
let listener = TcpListener::bind(self.base.endpoint())?;
*self.listener.lock() = Some(listener);
Ok(())
}
pub fn close(&self) {
*self.listener.lock() = None;
}
pub fn dispose(&self) {
self.close();
self.base.dispose();
}
pub fn accept_once(&self) -> bool {
let listener = {
let guard = self.listener.lock();
match guard.as_ref() {
Some(listener) => match listener.try_clone() {
Ok(cloned) => cloned,
Err(error) => return self.handle_accept_error(&error),
},
None => return false,
}
};
match listener.accept() {
Ok((stream, peer)) => {
let _ = self.handle_new_connection(stream, peer.to_string());
true
}
Err(error) => self.handle_accept_error(&error),
}
}
pub fn handle_accept_error(&self, error: &io::Error) -> bool {
match accept_error_tier(error) {
AcceptErrorTier::FatalClean => false,
AcceptErrorTier::Backoff => {
let backoff_ms = self.grow_backoff();
log::warn!("Accept backoff ({backoff_ms}ms) due to resource pressure: {error}");
sleep(Duration::from_millis(backoff_ms));
true
}
AcceptErrorTier::Transient => {
log::debug!("Transient accept error, continuing: {error}");
true
}
}
}
pub fn handle_new_connection(
&self,
stream: TcpStream,
remote_endpoint: String,
) -> Option<AcceptedConnection> {
if stream.set_nodelay(true).is_err() {
return None;
}
self
.accept_backoff_ms
.store(INITIAL_ACCEPT_BACKOFF_MS, Release);
let session = self.try_create_message_consumer(remote_endpoint.as_bytes())?;
let handler_id = self
.base
.admit_handler(session, self.network_connection_limit)?;
self.base.increment_connections_received();
Some(AcceptedConnection {
handler_id,
stream,
remote_endpoint,
})
}
fn grow_backoff(&self) -> u64 {
let backoff_ms = self.accept_backoff_ms.load(Acquire);
let next = (backoff_ms * 2).min(MAX_ACCEPT_BACKOFF_MS);
self.accept_backoff_ms.store(next, Release);
backoff_ms
}
pub fn try_create_message_consumer(&self, _bytes: &[u8]) -> Option<Arc<dyn MessageConsumerFace>> {
let provider = self.base.find_session_provider(WireFormat::Ascii)?;
self
.base
.add_session(WireFormat::Ascii, provider.as_ref(), next_sender_id())
}
pub fn dispose_message_consumer(&self, handler_id: u64) -> bool {
self.base.dispose_message_consumer(handler_id)
}
}
impl ServerEnumerate for GarnetServerTcp {
fn active_consumers(&self) -> Vec<Arc<dyn MessageConsumerFace>> {
self.base.active_consumers()
}
fn active_cluster_sessions(&self) -> Vec<Arc<dyn ClusterSessionFace>> {
self.base.active_cluster_sessions()
}
}
impl super::i_garnet_server::GarnetServer for GarnetServerTcp {
fn register(
&self,
wire_format: WireFormat,
backend_provider: Arc<dyn super::i_garnet_server::SessionProviderFace>,
) -> Result<(), super::i_garnet_server::ServerError> {
self.base.register(wire_format, backend_provider)
}
fn unregister(
&self,
wire_format: WireFormat,
) -> Option<Arc<dyn super::i_garnet_server::SessionProviderFace>> {
self.base.unregister(wire_format)
}
fn get_session_providers(
&self,
) -> Vec<(
WireFormat,
Arc<dyn super::i_garnet_server::SessionProviderFace>,
)> {
self.base.get_session_providers()
}
fn add_session(
&self,
wire_format: WireFormat,
backend_provider: &dyn super::i_garnet_server::SessionProviderFace,
network_sender_id: u64,
) -> Option<Arc<dyn MessageConsumerFace>> {
self
.base
.add_session(wire_format, backend_provider, network_sender_id)
}
fn start(&self) -> io::Result<()> {
GarnetServerTcp::start(self)
}
fn close(&self) {
GarnetServerTcp::close(self);
}
fn dispose(&self) {
GarnetServerTcp::dispose(self);
}
}
fn next_sender_id() -> u64 {
static SENDER_ID: AtomicU64 = AtomicU64::new(1);
SENDER_ID.fetch_add(1, AcqRel)
}
fn accept_error_tier(error: &io::Error) -> AcceptErrorTier {
use std::io::ErrorKind::{
BrokenPipe, ConnectionAborted, ConnectionReset, Interrupted, OutOfMemory, WouldBlock,
};
const EMFILE: i32 = 24;
const ENFILE: i32 = 23;
match error.kind() {
BrokenPipe => AcceptErrorTier::FatalClean,
OutOfMemory | WouldBlock => AcceptErrorTier::Backoff,
ConnectionAborted | ConnectionReset | Interrupted => AcceptErrorTier::Transient,
_ => {
#[cfg(unix)]
if matches!(error.raw_os_error(), Some(EMFILE | ENFILE)) {
return AcceptErrorTier::Backoff;
}
AcceptErrorTier::Transient
}
}
}
#[cfg(test)]
mod tests {
use std::{io::ErrorKind, net::SocketAddr};
use super::{
super::garnet_server_base::{DISPOSED_HANDLER_COUNT, GarnetServerBase},
*,
};
struct TestConsumer;
impl MessageConsumerFace for TestConsumer {
fn dispose(&self) {}
fn attach_server(&self, _server: Arc<dyn ServerEnumerate>) {}
}
struct TestProvider;
impl super::super::i_garnet_server::SessionProviderFace for TestProvider {
fn get_session(
&self,
_wire_format: WireFormat,
_network_sender_id: u64,
) -> Option<Arc<dyn MessageConsumerFace>> {
Some(Arc::new(TestConsumer))
}
}
fn server(limit: i32) -> (GarnetServerTcp, SocketAddr) {
let server = GarnetServerTcp::new("127.0.0.1:0", 0, 8, limit);
server
.base()
.register(WireFormat::Ascii, Arc::new(TestProvider))
.expect("注册成功");
server.start().expect("绑定成功");
let addr = server
.listener
.lock()
.as_ref()
.expect("监听在册")
.local_addr()
.expect("本地地址可得");
(server, addr)
}
#[test]
fn start_accept_and_dispose_roundtrip() {
let (server, addr) = server(-1);
let client = TcpStream::connect(addr).expect("回环连接成功");
drop(client);
assert!(server.accept_once(), "接受成功");
assert_eq!(server.base().get_conn_active(), 1);
assert_eq!(server.base().total_connections_received(), 1);
let handler_ids = server.base().active_handler_ids();
assert_eq!(handler_ids.len(), 1);
assert!(server.dispose_message_consumer(handler_ids[0]));
assert_eq!(server.base().get_conn_active(), 0);
assert_eq!(server.base().total_connections_disposed(), 1);
assert!(!server.dispose_message_consumer(handler_ids[0]));
server.close();
assert!(!server.accept_once());
}
#[test]
fn handle_new_connection_hands_stream_to_host_pump() {
let (server, addr) = server(-1);
let client = TcpStream::connect(addr).expect("回环连接成功");
drop(client);
let listener = server
.listener
.lock()
.as_ref()
.expect("监听在册")
.try_clone()
.expect("克隆成功");
let (stream, peer) = listener.accept().expect("接受成功");
let accepted = server
.handle_new_connection(stream, peer.to_string())
.expect("登记成功");
assert_eq!(server.base().get_conn_active(), 1);
assert!(server.dispose_message_consumer(accepted.handler_id));
drop(accepted.stream);
}
#[test]
fn connection_limit_rejects_new_connections() {
let (server, addr) = server(1);
let client = TcpStream::connect(addr).expect("回环连接成功");
drop(client);
assert!(server.accept_once());
assert_eq!(server.base().get_conn_active(), 1);
let second = TcpStream::connect(addr).expect("第二连接成功");
drop(second);
assert!(server.accept_once());
assert_eq!(server.base().get_conn_active(), 1);
}
#[test]
fn accept_backoff_grows_to_max() {
let server = GarnetServerTcp::new("ep", 0, 8, -1);
assert_eq!(server.accept_backoff_ms(), INITIAL_ACCEPT_BACKOFF_MS);
assert_eq!(server.grow_backoff(), 100);
assert_eq!(server.accept_backoff_ms(), 200);
assert_eq!(server.grow_backoff(), 200);
assert_eq!(server.grow_backoff(), 400);
for _ in 0..24 {
server.grow_backoff();
}
assert_eq!(server.accept_backoff_ms(), MAX_ACCEPT_BACKOFF_MS);
assert_eq!(server.grow_backoff(), MAX_ACCEPT_BACKOFF_MS);
assert_eq!(server.accept_backoff_ms(), MAX_ACCEPT_BACKOFF_MS);
}
#[test]
fn successful_accept_resets_backoff() {
let (server, addr) = server(-1);
assert!(server.handle_accept_error(&io::Error::from_raw_os_error(24)));
assert_eq!(server.accept_backoff_ms(), 2 * INITIAL_ACCEPT_BACKOFF_MS);
let client = TcpStream::connect(addr).expect("回环连接成功");
drop(client);
assert!(server.accept_once());
assert_eq!(server.accept_backoff_ms(), INITIAL_ACCEPT_BACKOFF_MS);
}
#[test]
fn fatal_clean_tier_stops_accept_loop() {
let server = GarnetServerTcp::new("ep", 0, 8, -1);
let shutdown = io::Error::new(ErrorKind::BrokenPipe, "shutdown");
assert!(!server.handle_accept_error(&shutdown));
}
#[test]
fn transient_tiers_keep_loop_alive() {
let server = GarnetServerTcp::new("ep", 0, 8, -1);
for kind in [ErrorKind::Interrupted, ErrorKind::ConnectionAborted] {
let error = io::Error::new(kind, "transient");
assert!(server.handle_accept_error(&error), "{kind:?} 应继续");
assert_eq!(server.accept_backoff_ms(), INITIAL_ACCEPT_BACKOFF_MS);
}
}
#[test]
fn resource_pressure_tier_backs_off() {
let server = GarnetServerTcp::new("ep", 0, 8, -1);
let exhausted = io::Error::from_raw_os_error(24);
assert!(server.handle_accept_error(&exhausted));
assert_eq!(server.accept_backoff_ms(), 2 * INITIAL_ACCEPT_BACKOFF_MS);
}
#[test]
fn dispose_closes_listener_then_drains() {
let (server, addr) = server(-1);
let client = TcpStream::connect(addr).expect("回环连接成功");
drop(client);
assert!(server.accept_once());
assert_eq!(server.base().get_conn_active(), 1);
server.dispose();
assert!(server.listener.lock().is_none());
assert_eq!(server.base().get_conn_active(), 0);
assert!(
server
.base()
.register_handler(Arc::new(TestConsumer))
.is_none()
);
server.dispose();
}
#[test]
fn base_served_by_tcp_implementation() {
let (server, _addr) = server(-1);
let base: &GarnetServerBase = server.base();
assert_eq!(base.get_conn_active(), 0);
let _ = DISPOSED_HANDLER_COUNT;
}
}