#![allow(clippy::expect_used)]
use std::fmt;
use std::marker::PhantomData;
use std::net::{SocketAddr, ToSocketAddrs};
use crate::connection::{ConnectionConfig, ConnectionError, IngressStats, State};
use crate::connection_state::ConnectionState;
use crate::observability::LatencyHistograms;
use crate::proactor::{Completion, Proactor, ProactorConfig, ProactorError};
use crate::ws::{
ConnState as WsConnState, DataEvent as WsDataEvent, Event as WsEvent,
MarkedDataEvent as WsMarkedDataEvent,
};
const CONN_ID_MASK: u64 = 0x0FFF_FFFF;
pub const DEFAULT_POOL_INITIAL_CONN_CAPACITY: usize = 0;
pub const DEFAULT_POOL_COMPLETION_BATCH_CAPACITY: usize = 16;
#[derive(Debug, Clone, Copy)]
pub struct PoolConfig {
pub proactor: ProactorConfig,
pub initial_conn_capacity: usize,
pub completion_batch_capacity: usize,
pub post_progress_spin_iters: usize,
}
impl PoolConfig {
#[must_use]
pub const fn new(proactor: ProactorConfig) -> Self {
Self {
proactor,
initial_conn_capacity: DEFAULT_POOL_INITIAL_CONN_CAPACITY,
completion_batch_capacity: DEFAULT_POOL_COMPLETION_BATCH_CAPACITY,
post_progress_spin_iters: 0,
}
}
#[must_use]
pub const fn with_initial_conn_capacity(mut self, capacity: usize) -> Self {
self.initial_conn_capacity = capacity;
self
}
#[must_use]
pub const fn with_completion_batch_capacity(mut self, capacity: usize) -> Self {
self.completion_batch_capacity = capacity;
self
}
#[must_use]
pub const fn with_post_progress_spin_iters(mut self, iters: usize) -> Self {
self.post_progress_spin_iters = iters;
self
}
}
impl Default for PoolConfig {
fn default() -> Self {
Self::new(ProactorConfig::default())
}
}
#[derive(Debug, Clone, Copy, Eq, PartialEq, Hash)]
pub struct ConnHandle(u32);
impl ConnHandle {
#[inline]
#[must_use]
pub const fn as_u32(self) -> u32 {
self.0
}
}
pub struct Pool {
proactor: Proactor,
conns: Vec<Option<ConnectionState>>,
active_count: u32,
next_conn_id: u32,
next_bgid: u16,
completions_buf: Vec<Completion>,
post_progress_spin_iters: usize,
_not_send: PhantomData<*const ()>,
}
impl std::fmt::Debug for Pool {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("Pool")
.field("proactor", &self.proactor)
.field("active_count", &self.active_count)
.field("slot_len", &self.conns.len())
.field("slot_capacity", &self.conns.capacity())
.field("next_conn_id", &self.next_conn_id)
.field("next_bgid", &self.next_bgid)
.finish()
}
}
impl Pool {
pub fn new(cfg: PoolConfig) -> Result<Self, ProactorError> {
let proactor = Proactor::new(cfg.proactor)?;
Ok(Self {
proactor,
conns: Vec::with_capacity(cfg.initial_conn_capacity),
active_count: 0,
next_conn_id: 0,
next_bgid: 0,
completions_buf: Vec::with_capacity(cfg.completion_batch_capacity),
post_progress_spin_iters: cfg.post_progress_spin_iters,
_not_send: PhantomData,
})
}
pub fn connect_blocking(
&mut self,
cfg: ConnectionConfig,
) -> Result<ConnHandle, ConnectionError> {
let addr = resolve_addr(&cfg)?;
self.connect_blocking_to(cfg, addr)
}
pub fn connect_blocking_to(
&mut self,
cfg: ConnectionConfig,
addr: SocketAddr,
) -> Result<ConnHandle, ConnectionError> {
let handle = self.submit_connect_to(cfg, addr)?;
let conn_id = handle.0;
match self.drive_conn_until_open(conn_id) {
Ok(()) => Ok(handle),
Err(e) => {
self.drop_slot(conn_id);
Err(e)
}
}
}
pub fn submit_connect(&mut self, cfg: ConnectionConfig) -> Result<ConnHandle, ConnectionError> {
let addr = resolve_addr(&cfg)?;
self.submit_connect_to(cfg, addr)
}
pub fn submit_connect_to(
&mut self,
mut cfg: ConnectionConfig,
addr: SocketAddr,
) -> Result<ConnHandle, ConnectionError> {
let conn_id = self.next_conn_id;
if conn_id > CONN_ID_MASK as u32 {
return Err(ConnectionError::IdSpaceExhausted("conn_id"));
}
let bgid = self.next_bgid;
let Some(next_bgid) = self.next_bgid.checked_add(1) else {
return Err(ConnectionError::IdSpaceExhausted("bgid"));
};
cfg.conn_id = conn_id;
cfg.bgid = bgid;
let mut conn = ConnectionState::new(cfg, addr)?;
conn.submit_connect(&mut self.proactor)?;
debug_assert_eq!(self.conns.len(), conn_id as usize);
self.conns.push(Some(conn));
self.active_count += 1;
self.next_conn_id = conn_id + 1;
self.next_bgid = next_bgid;
Ok(ConnHandle(conn_id))
}
fn drop_slot(&mut self, conn_id: u32) {
let Some(slot) = self.conns.get_mut(conn_id as usize) else {
return;
};
if let Some(mut dead) = slot.take() {
if let Some(mut ring) = dead.buf_ring.take() {
let _ = ring.unregister(&mut self.proactor);
}
self.active_count = self.active_count.saturating_sub(1);
}
}
fn drive_conn_until_open(&mut self, conn_id: u32) -> Result<(), ConnectionError> {
loop {
self.pump_impl(1, |_, _| { })?;
let conn = self
.conns
.get_mut(conn_id as usize)
.and_then(Option::as_mut)
.expect("just-added conn must exist");
conn.sync_ws_open_state();
match conn.state() {
State::Open => return Ok(()),
State::Closed => return Err(ConnectionError::PeerClosed),
_ => {}
}
}
}
pub fn send_text(&mut self, h: ConnHandle, payload: &[u8]) -> Result<(), ConnectionError> {
let conn = self.conn_mut(h)?;
conn.assert_open()?;
conn.ws.send_text(payload)?;
Ok(())
}
pub fn send_binary(&mut self, h: ConnHandle, payload: &[u8]) -> Result<(), ConnectionError> {
let conn = self.conn_mut(h)?;
conn.assert_open()?;
conn.ws.send_binary(payload)?;
Ok(())
}
pub fn send_ping(&mut self, h: ConnHandle, payload: &[u8]) -> Result<(), ConnectionError> {
let conn = self.conn_mut(h)?;
conn.assert_open()?;
conn.ws.send_ping(payload)?;
Ok(())
}
pub fn send_pong(&mut self, h: ConnHandle, payload: &[u8]) -> Result<(), ConnectionError> {
let conn = self.conn_mut(h)?;
conn.assert_open()?;
conn.ws.send_pong(payload)?;
Ok(())
}
pub fn initiate_close(
&mut self,
h: ConnHandle,
code: u16,
reason: &str,
) -> Result<(), ConnectionError> {
let conn = self.conn_mut(h)?;
if matches!(conn.state(), State::Closed | State::Closing) {
return Ok(());
}
conn.ws.send_close(code, reason)?;
if matches!(conn.state(), State::Open) {
conn.state = State::Closing;
}
Ok(())
}
pub fn pump<F>(&mut self, sink: F) -> Result<(), ConnectionError>
where
F: FnMut(ConnHandle, WsEvent<'_>),
{
self.pump_impl(1, sink)
}
pub fn pump_nowait<F>(&mut self, sink: F) -> Result<(), ConnectionError>
where
F: FnMut(ConnHandle, WsEvent<'_>),
{
self.pump_impl(0, sink)
}
pub fn pump_spin<F>(&mut self, spin_iters: usize, sink: F) -> Result<bool, ConnectionError>
where
F: FnMut(ConnHandle, WsEvent<'_>),
{
self.pump_spin_impl(spin_iters, sink)
}
pub fn pump_data<F>(&mut self, sink: F) -> Result<(), ConnectionError>
where
F: for<'a> FnMut(ConnHandle, WsDataEvent<'a>),
{
self.pump_data_impl(1, sink)
}
pub fn pump_data_nowait<F>(&mut self, sink: F) -> Result<(), ConnectionError>
where
F: for<'a> FnMut(ConnHandle, WsDataEvent<'a>),
{
self.pump_data_impl(0, sink)
}
pub fn pump_data_marked<F>(&mut self, sink: F) -> Result<(), ConnectionError>
where
F: for<'a> FnMut(ConnHandle, WsMarkedDataEvent<'a>),
{
self.pump_data_marked_impl(1, sink)
}
pub fn pump_data_marked_nowait<F>(&mut self, sink: F) -> Result<(), ConnectionError>
where
F: for<'a> FnMut(ConnHandle, WsMarkedDataEvent<'a>),
{
self.pump_data_marked_impl(0, sink)
}
pub fn pump_data_spin<F>(&mut self, spin_iters: usize, sink: F) -> Result<bool, ConnectionError>
where
F: for<'a> FnMut(ConnHandle, WsDataEvent<'a>),
{
self.pump_data_spin_impl(spin_iters, sink)
}
pub fn pump_data_spin_marked<F>(
&mut self,
spin_iters: usize,
sink: F,
) -> Result<bool, ConnectionError>
where
F: for<'a> FnMut(ConnHandle, WsMarkedDataEvent<'a>),
{
self.pump_data_spin_marked_impl(spin_iters, sink)
}
fn pump_data_impl<F>(&mut self, wait_nr: usize, mut sink: F) -> Result<(), ConnectionError>
where
F: for<'a> FnMut(ConnHandle, WsDataEvent<'a>),
{
let Self {
proactor,
conns,
completions_buf,
active_count,
..
} = self;
let mut first_err: Option<ConnectionError> = None;
for slot in conns.iter_mut() {
let Some(conn) = slot.as_mut() else { continue };
if let Err(e) = conn.try_submit_send(proactor) {
fail_conn(conn, e, &mut first_err);
continue;
}
if let Err(e) = conn.try_rearm_multishot(proactor) {
fail_conn(conn, e, &mut first_err);
}
}
proactor.submit()?;
proactor.wait_for_cqe(wait_nr)?;
completions_buf.clear();
proactor.drain_completions(|c| completions_buf.push(c));
dispatch_conn_completions_data(conns, proactor, completions_buf, &mut sink, &mut first_err);
sync_active_count(conns, active_count);
first_err.map_or(Ok(()), Err)
}
fn pump_data_marked_impl<F>(
&mut self,
wait_nr: usize,
mut sink: F,
) -> Result<(), ConnectionError>
where
F: for<'a> FnMut(ConnHandle, WsMarkedDataEvent<'a>),
{
let Self {
proactor,
conns,
completions_buf,
active_count,
..
} = self;
let mut first_err: Option<ConnectionError> = None;
for slot in conns.iter_mut() {
let Some(conn) = slot.as_mut() else { continue };
if let Err(e) = conn.try_submit_send(proactor) {
fail_conn(conn, e, &mut first_err);
continue;
}
if let Err(e) = conn.try_rearm_multishot(proactor) {
fail_conn(conn, e, &mut first_err);
}
}
proactor.submit()?;
proactor.wait_for_cqe(wait_nr)?;
completions_buf.clear();
proactor.drain_completions(|c| completions_buf.push(c));
for &c in completions_buf.iter() {
let conn_id =
u32::try_from(c.user_data.token() & CONN_ID_MASK).expect("28-bit mask fits u32");
if let Some(conn) = conns.get_mut(conn_id as usize).and_then(Option::as_mut) {
let handle = ConnHandle(conn.conn_id());
match conn.handle_completion_data_marked(proactor, c, |ev| sink(handle, ev)) {
Ok(_) => {
conn.sync_ws_open_state();
conn.sync_ws_close_state();
}
Err(e) => fail_conn(conn, e, &mut first_err),
}
}
}
sync_active_count(conns, active_count);
first_err.map_or(Ok(()), Err)
}
fn pump_data_spin_impl<F>(
&mut self,
spin_iters: usize,
mut sink: F,
) -> Result<bool, ConnectionError>
where
F: for<'a> FnMut(ConnHandle, WsDataEvent<'a>),
{
let post_progress_spin_iters = self.post_progress_spin_iters;
let Self {
proactor,
conns,
completions_buf,
active_count,
..
} = self;
let mut first_err: Option<ConnectionError> = None;
let mut progressed = false;
submit_conn_ops(conns, proactor, &mut first_err);
proactor.submit()?;
for iter in 0..=spin_iters {
let cqes = drain_conn_completions_data(
conns,
proactor,
completions_buf,
&mut sink,
&mut first_err,
);
if cqes > 0 {
progressed = true;
for _ in 0..post_progress_spin_iters {
std::hint::spin_loop();
let _ = drain_conn_completions_data(
conns,
proactor,
completions_buf,
&mut sink,
&mut first_err,
);
if first_err.is_some() {
break;
}
}
}
if progressed || first_err.is_some() {
break;
}
if iter < spin_iters {
std::hint::spin_loop();
}
}
sync_active_count(conns, active_count);
match first_err {
Some(e) => Err(e),
None => Ok(progressed),
}
}
fn pump_data_spin_marked_impl<F>(
&mut self,
spin_iters: usize,
mut sink: F,
) -> Result<bool, ConnectionError>
where
F: for<'a> FnMut(ConnHandle, WsMarkedDataEvent<'a>),
{
let Self {
proactor,
conns,
completions_buf,
active_count,
..
} = self;
let mut first_err: Option<ConnectionError> = None;
let mut progressed = false;
submit_conn_ops(conns, proactor, &mut first_err);
proactor.submit()?;
for iter in 0..=spin_iters {
let cqes = drain_conn_completions_data_marked(
conns,
proactor,
completions_buf,
&mut sink,
&mut first_err,
);
progressed |= cqes > 0;
if progressed || first_err.is_some() {
break;
}
if iter < spin_iters {
std::hint::spin_loop();
}
}
sync_active_count(conns, active_count);
match first_err {
Some(e) => Err(e),
None => Ok(progressed),
}
}
fn pump_impl<F>(&mut self, wait_nr: usize, mut sink: F) -> Result<(), ConnectionError>
where
F: FnMut(ConnHandle, WsEvent<'_>),
{
let Self {
proactor,
conns,
completions_buf,
active_count,
..
} = self;
let mut first_err: Option<ConnectionError> = None;
for slot in conns.iter_mut() {
let Some(conn) = slot.as_mut() else { continue };
if let Err(e) = conn.try_submit_send(proactor) {
fail_conn(conn, e, &mut first_err);
continue;
}
if let Err(e) = conn.try_rearm_multishot(proactor) {
fail_conn(conn, e, &mut first_err);
}
}
proactor.submit()?;
proactor.wait_for_cqe(wait_nr)?;
completions_buf.clear();
proactor.drain_completions(|c| completions_buf.push(c));
for &c in completions_buf.iter() {
let conn_id =
u32::try_from(c.user_data.token() & CONN_ID_MASK).expect("28-bit mask fits u32");
if let Some(conn) = conns.get_mut(conn_id as usize).and_then(Option::as_mut)
&& let Err(e) = conn.handle_completion(proactor, c)
{
fail_conn(conn, e, &mut first_err);
}
}
for slot in conns.iter_mut() {
let Some(conn) = slot.as_mut() else { continue };
let handle = ConnHandle(conn.conn_id());
while let Some(res) = conn.ws.poll_event() {
match res {
Ok(ev) => sink(handle, ev),
Err(e) => {
fail_conn(conn, ConnectionError::Ws(e), &mut first_err);
break;
}
}
}
conn.sync_ws_open_state();
conn.sync_ws_close_state();
conn.clear_ws_ingress_dirty();
}
sync_active_count(conns, active_count);
first_err.map_or(Ok(()), Err)
}
fn pump_spin_impl<F>(&mut self, spin_iters: usize, mut sink: F) -> Result<bool, ConnectionError>
where
F: FnMut(ConnHandle, WsEvent<'_>),
{
let Self {
proactor,
conns,
completions_buf,
active_count,
..
} = self;
let mut first_err: Option<ConnectionError> = None;
let mut progressed = false;
submit_conn_ops(conns, proactor, &mut first_err);
proactor.submit()?;
for iter in 0..=spin_iters {
let cqes = drain_conn_completions(conns, proactor, completions_buf, &mut first_err);
progressed |= cqes > 0;
for slot in conns.iter_mut() {
let Some(conn) = slot.as_mut() else { continue };
let handle = ConnHandle(conn.conn_id());
while let Some(res) = conn.ws.poll_event() {
progressed = true;
match res {
Ok(ev) => sink(handle, ev),
Err(e) => {
fail_conn(conn, ConnectionError::Ws(e), &mut first_err);
break;
}
}
}
conn.sync_ws_open_state();
conn.sync_ws_close_state();
conn.clear_ws_ingress_dirty();
}
if progressed || first_err.is_some() {
break;
}
if iter < spin_iters {
std::hint::spin_loop();
}
}
sync_active_count(conns, active_count);
match first_err {
Some(e) => Err(e),
None => Ok(progressed),
}
}
pub fn state(&self, h: ConnHandle) -> Option<State> {
self.conns
.get(h.0 as usize)
.and_then(Option::as_ref)
.map(ConnectionState::state)
}
#[must_use]
pub fn conn_count(&self) -> usize {
self.active_count as usize
}
#[must_use]
pub fn ingress_stats(&self, h: ConnHandle) -> Option<IngressStats> {
self.conns
.get(h.0 as usize)
.and_then(Option::as_ref)
.map(ConnectionState::ingress_stats)
}
#[must_use]
pub fn prometheus_metrics(&self) -> String {
let mut out = String::new();
self.write_prometheus_metrics(&mut out)
.expect("writing Prometheus metrics to String cannot fail");
out
}
pub fn write_prometheus_metrics<W: fmt::Write>(&self, out: &mut W) -> fmt::Result {
LatencyHistograms::write_prometheus_help(out)?;
write_ingress_prometheus_help(out)?;
for conn in self.conns.iter().flatten() {
conn.write_prometheus_metrics(out)?;
}
Ok(())
}
#[must_use]
pub fn prometheus_metrics_and_reset_interval(&mut self) -> String {
let mut out = String::new();
self.write_prometheus_metrics_and_reset_interval(&mut out)
.expect("writing Prometheus metrics to String cannot fail");
out
}
pub fn write_prometheus_metrics_and_reset_interval<W: fmt::Write>(
&mut self,
out: &mut W,
) -> fmt::Result {
LatencyHistograms::write_prometheus_help(out)?;
write_ingress_prometheus_help(out)?;
for conn in self.conns.iter_mut().flatten() {
conn.write_prometheus_metrics_and_reset_interval(out)?;
}
Ok(())
}
fn conn_mut(&mut self, h: ConnHandle) -> Result<&mut ConnectionState, ConnectionError> {
self.conns
.get_mut(h.0 as usize)
.and_then(Option::as_mut)
.ok_or(ConnectionError::InvalidState(State::Closed))
}
}
fn resolve_addr(cfg: &ConnectionConfig) -> Result<SocketAddr, ConnectionError> {
(cfg.host.as_str(), cfg.port)
.to_socket_addrs()?
.next()
.ok_or_else(|| ConnectionError::DnsEmpty(cfg.host.clone()))
}
fn submit_conn_ops(
conns: &mut [Option<ConnectionState>],
proactor: &mut Proactor,
first_err: &mut Option<ConnectionError>,
) {
for slot in conns.iter_mut() {
let Some(conn) = slot.as_mut() else { continue };
if let Err(e) = conn.try_submit_send(proactor) {
fail_conn(conn, e, first_err);
continue;
}
if let Err(e) = conn.try_rearm_multishot(proactor) {
fail_conn(conn, e, first_err);
}
}
}
#[inline]
fn completion_conn_id(c: Completion) -> u32 {
u32::try_from(c.user_data.token() & CONN_ID_MASK).expect("28-bit mask fits u32")
}
fn drain_conn_completions(
conns: &mut [Option<ConnectionState>],
proactor: &mut Proactor,
completions_buf: &mut Vec<Completion>,
first_err: &mut Option<ConnectionError>,
) -> usize {
completions_buf.clear();
let count = proactor.drain_completions(|c| completions_buf.push(c));
for &c in completions_buf.iter() {
let conn_id =
u32::try_from(c.user_data.token() & CONN_ID_MASK).expect("28-bit mask fits u32");
if let Some(conn) = conns.get_mut(conn_id as usize).and_then(Option::as_mut)
&& let Err(e) = conn.handle_completion(proactor, c)
{
fail_conn(conn, e, first_err);
}
}
count
}
fn drain_conn_completions_data<F>(
conns: &mut [Option<ConnectionState>],
proactor: &mut Proactor,
completions_buf: &mut Vec<Completion>,
sink: &mut F,
first_err: &mut Option<ConnectionError>,
) -> usize
where
F: for<'a> FnMut(ConnHandle, WsDataEvent<'a>),
{
completions_buf.clear();
let count = proactor.drain_completions(|c| completions_buf.push(c));
dispatch_conn_completions_data(conns, proactor, completions_buf, sink, first_err);
count
}
fn dispatch_conn_completions_data<F>(
conns: &mut [Option<ConnectionState>],
proactor: &mut Proactor,
completions_buf: &[Completion],
sink: &mut F,
first_err: &mut Option<ConnectionError>,
) where
F: for<'a> FnMut(ConnHandle, WsDataEvent<'a>),
{
let mut i = 0_usize;
while i < completions_buf.len() {
let c = completions_buf[i];
let conn_id = completion_conn_id(c);
let Some(conn) = conns.get_mut(conn_id as usize).and_then(Option::as_mut) else {
i += 1;
continue;
};
let handle = ConnHandle(conn.conn_id());
if conn.can_handle_plain_recv_data_batch(c) {
let mut end = i + 1;
while end < completions_buf.len() {
let next = completions_buf[end];
if completion_conn_id(next) != conn_id
|| !conn.can_handle_plain_recv_data_batch(next)
{
break;
}
end += 1;
}
if end > i + 1 {
match conn.handle_plain_recv_data_batch(&completions_buf[i..end], &mut |ev| {
sink(handle, ev);
}) {
Ok(_) => {
conn.sync_ws_open_state();
conn.sync_ws_close_state();
}
Err(e) => fail_conn(conn, e, first_err),
}
i = end;
continue;
}
}
match conn.handle_completion_data(proactor, c, |ev| sink(handle, ev)) {
Ok(_) => {
conn.sync_ws_open_state();
conn.sync_ws_close_state();
}
Err(e) => fail_conn(conn, e, first_err),
}
i += 1;
}
}
fn drain_conn_completions_data_marked<F>(
conns: &mut [Option<ConnectionState>],
proactor: &mut Proactor,
completions_buf: &mut Vec<Completion>,
sink: &mut F,
first_err: &mut Option<ConnectionError>,
) -> usize
where
F: for<'a> FnMut(ConnHandle, WsMarkedDataEvent<'a>),
{
completions_buf.clear();
let count = proactor.drain_completions(|c| completions_buf.push(c));
for &c in completions_buf.iter() {
let conn_id =
u32::try_from(c.user_data.token() & CONN_ID_MASK).expect("28-bit mask fits u32");
if let Some(conn) = conns.get_mut(conn_id as usize).and_then(Option::as_mut) {
let handle = ConnHandle(conn.conn_id());
match conn.handle_completion_data_marked(proactor, c, |ev| sink(handle, ev)) {
Ok(_) => {
conn.sync_ws_open_state();
conn.sync_ws_close_state();
}
Err(e) => fail_conn(conn, e, first_err),
}
}
}
count
}
fn write_ingress_prometheus_help<W: fmt::Write>(out: &mut W) -> fmt::Result {
writeln!(
out,
"# HELP talaris_ingress_recv_data_cqes_total Positive-length recv data CQEs handled by a connection."
)?;
writeln!(out, "# TYPE talaris_ingress_recv_data_cqes_total counter")?;
writeln!(
out,
"# HELP talaris_ingress_recv_bytes_total Ciphertext bytes carried by recv data CQEs."
)?;
writeln!(out, "# TYPE talaris_ingress_recv_bytes_total counter")?;
writeln!(
out,
"# HELP talaris_ingress_recv_multishot_rearms_total Recv multishot SQEs submitted or rearmed."
)?;
writeln!(
out,
"# TYPE talaris_ingress_recv_multishot_rearms_total counter"
)?;
writeln!(
out,
"# HELP talaris_ingress_recv_ring_exhaustions_total Recv multishot terminations caused by provided-buffer ring exhaustion."
)?;
writeln!(
out,
"# TYPE talaris_ingress_recv_ring_exhaustions_total counter"
)?;
writeln!(
out,
"# HELP talaris_ingress_plain_recv_batches_total Consecutive plain TCP recv CQE runs handled by the data-pump batch path."
)?;
writeln!(
out,
"# TYPE talaris_ingress_plain_recv_batches_total counter"
)?;
writeln!(
out,
"# HELP talaris_ingress_plain_recv_batch_cqes_total Total recv CQEs included in plain TCP data-pump batch runs."
)?;
writeln!(
out,
"# TYPE talaris_ingress_plain_recv_batch_cqes_total counter"
)?;
writeln!(
out,
"# HELP talaris_ingress_plain_recv_copied_batches_total Plain TCP data-pump batch runs parsed through the reusable copy scratch buffer."
)?;
writeln!(
out,
"# TYPE talaris_ingress_plain_recv_copied_batches_total counter"
)?;
writeln!(
out,
"# HELP talaris_ingress_plain_recv_copied_bytes_total Bytes copied into the reusable plain TCP data-pump batch scratch buffer."
)?;
writeln!(
out,
"# TYPE talaris_ingress_plain_recv_copied_bytes_total counter"
)?;
writeln!(
out,
"# HELP talaris_ingress_plaintext_chunks_total Plaintext chunks fed into the WebSocket parser."
)?;
writeln!(out, "# TYPE talaris_ingress_plaintext_chunks_total counter")?;
writeln!(
out,
"# HELP talaris_ingress_plaintext_bytes_total Plaintext bytes fed into the WebSocket parser."
)?;
writeln!(out, "# TYPE talaris_ingress_plaintext_bytes_total counter")?;
writeln!(
out,
"# HELP talaris_ingress_ws_data_drains_total Data-pump CQEs that fed plaintext into WebSocket receive processing."
)?;
writeln!(out, "# TYPE talaris_ingress_ws_data_drains_total counter")?;
writeln!(
out,
"# HELP talaris_ingress_ws_data_drain_skips_total Data-pump CQEs that skipped WebSocket draining because no plaintext arrived."
)?;
writeln!(
out,
"# TYPE talaris_ingress_ws_data_drain_skips_total counter"
)?;
writeln!(
out,
"# HELP talaris_ingress_ws_data_events_total Text/Binary data messages emitted to the user's data sink."
)?;
writeln!(out, "# TYPE talaris_ingress_ws_data_events_total counter")?;
writeln!(
out,
"# HELP talaris_ingress_ws_text_events_total Text messages emitted to the user's data sink."
)?;
writeln!(out, "# TYPE talaris_ingress_ws_text_events_total counter")?;
writeln!(
out,
"# HELP talaris_ingress_ws_binary_events_total Binary messages emitted to the user's data sink."
)?;
writeln!(out, "# TYPE talaris_ingress_ws_binary_events_total counter")
}
fn fail_conn(
conn: &mut ConnectionState,
err: ConnectionError,
first_err: &mut Option<ConnectionError>,
) {
tracing::warn!(conn_id = conn.conn_id(), error = %err, "pool conn failed");
let ws_close_in_progress = matches!(conn.ws.state(), WsConnState::Closing);
conn.state = if ws_close_in_progress {
State::Closing
} else {
State::Closed
};
if first_err.is_none() {
*first_err = Some(err);
}
}
fn sync_active_count(conns: &[Option<ConnectionState>], active_count: &mut u32) {
let count = conns
.iter()
.filter(|slot| {
slot.as_ref()
.is_some_and(|conn| !matches!(conn.state(), State::Closed))
})
.count();
*active_count = u32::try_from(count).unwrap_or(u32::MAX);
}
impl Drop for Pool {
fn drop(&mut self) {
for slot in self.conns.iter_mut() {
if let Some(conn) = slot.as_mut()
&& let Some(mut ring) = conn.buf_ring.take()
{
let _ = ring.unregister(&mut self.proactor);
}
}
}
}
#[cfg(all(test, target_os = "linux"))]
#[allow(
clippy::unwrap_used,
clippy::expect_used,
clippy::indexing_slicing,
clippy::panic
)]
mod tests {
use super::*;
use crate::connection::{ConnectionConfig, State};
use crate::test_helpers::run_echo_server;
use crate::ws::DataEvent as WsDataEvent;
use std::net::{Ipv4Addr, SocketAddrV4, TcpListener};
use std::sync::mpsc;
use std::thread;
#[test]
fn pool_single_conn_plain_ws_echo_roundtrip() {
let listener = TcpListener::bind(SocketAddrV4::new(Ipv4Addr::LOCALHOST, 0)).unwrap();
let local_addr = listener.local_addr().unwrap();
let (_shutdown_tx, shutdown_rx) = mpsc::channel::<()>();
let server = thread::spawn(move || run_echo_server(listener, shutdown_rx));
let cfg = ConnectionConfig::new("localhost", local_addr.port(), "/echo").with_tls(false);
let mut pool = Pool::new(PoolConfig::new(cfg.proactor)).expect("pool");
let handle = pool.connect_blocking_to(cfg, local_addr).expect("connect");
assert_eq!(pool.state(handle), Some(State::Open));
pool.send_text(handle, b"hello").unwrap();
let mut got_text: Option<String> = None;
for _ in 0..50 {
pool.pump_data(|h, ev| {
assert_eq!(h, handle);
if let WsDataEvent::Text(s) = ev {
got_text = Some(s.to_owned());
}
})
.unwrap();
if got_text.is_some() {
break;
}
}
assert_eq!(got_text.as_deref(), Some("hello"));
pool.initiate_close(handle, 1000, "bye").unwrap();
for _ in 0..50 {
if matches!(pool.state(handle), Some(State::Closed | State::Closing)) {
let _ = pool.pump_nowait(|_, _| {});
}
if matches!(pool.state(handle), Some(State::Closed)) {
break;
}
let _ = pool.pump(|_, _| {});
}
server.join().unwrap();
}
#[test]
#[ignore = "需要外网 + test.deribit.com 可达;手动 --ignored 跑"]
fn tls_smoke_deribit_testnet() {
let cfg = ConnectionConfig::new("test.deribit.com", 443, "/ws/api/v2");
let mut pool = Pool::new(PoolConfig::new(cfg.proactor)).expect("pool");
let handle = pool
.connect_blocking(cfg)
.expect("tls handshake + ws upgrade");
assert_eq!(pool.state(handle), Some(State::Open));
eprintln!("TLS+WS handshake OK, sending public/test ...");
pool.send_text(
handle,
br#"{"jsonrpc":"2.0","id":1,"method":"public/test","params":{}}"#,
)
.unwrap();
let mut got = false;
for _ in 0..100 {
pool.pump(|_h, ev| {
if let WsEvent::Text(s) = ev {
eprintln!("got text: {s}");
got = true;
}
})
.unwrap();
if got {
break;
}
}
assert!(got, "no response from test.deribit.com");
pool.initiate_close(handle, 1000, "bye").unwrap();
for _ in 0..20 {
let _ = pool.pump_nowait(|_, _| {});
if matches!(pool.state(handle), Some(State::Closed)) {
break;
}
}
}
#[test]
fn pool_two_conns_no_cross_talk() {
let listener_a = TcpListener::bind(SocketAddrV4::new(Ipv4Addr::LOCALHOST, 0)).unwrap();
let addr_a = listener_a.local_addr().unwrap();
let listener_b = TcpListener::bind(SocketAddrV4::new(Ipv4Addr::LOCALHOST, 0)).unwrap();
let addr_b = listener_b.local_addr().unwrap();
let (_tx_a, rx_a) = mpsc::channel::<()>();
let (_tx_b, rx_b) = mpsc::channel::<()>();
let server_a = thread::spawn(move || run_echo_server(listener_a, rx_a));
let server_b = thread::spawn(move || run_echo_server(listener_b, rx_b));
let mut pool = Pool::new(PoolConfig::default()).expect("pool");
let cfg_a = ConnectionConfig::new("localhost", addr_a.port(), "/a").with_tls(false);
let cfg_b = ConnectionConfig::new("localhost", addr_b.port(), "/b").with_tls(false);
let h_a = pool.connect_blocking_to(cfg_a, addr_a).expect("connect a");
let h_b = pool.connect_blocking_to(cfg_b, addr_b).expect("connect b");
assert_eq!(pool.conn_count(), 2);
assert_ne!(h_a, h_b);
assert_eq!(pool.state(h_a), Some(State::Open));
assert_eq!(pool.state(h_b), Some(State::Open));
assert!(h_b.as_u32() > h_a.as_u32());
pool.send_text(h_a, b"alpha").unwrap();
pool.send_text(h_b, b"bravo").unwrap();
let mut a_text: Option<String> = None;
let mut b_text: Option<String> = None;
let mut wrong_route = false;
for _ in 0..200 {
pool.pump(|h, ev| {
if let WsEvent::Text(s) = ev {
if h == h_a {
if s != "alpha" {
wrong_route = true;
}
a_text = Some(s.to_owned());
} else if h == h_b {
if s != "bravo" {
wrong_route = true;
}
b_text = Some(s.to_owned());
} else {
wrong_route = true;
}
}
})
.unwrap();
if a_text.is_some() && b_text.is_some() {
break;
}
}
assert!(
!wrong_route,
"CQE 路由错位:handle 收到了不属于它的 payload"
);
assert_eq!(a_text.as_deref(), Some("alpha"));
assert_eq!(b_text.as_deref(), Some("bravo"));
pool.initiate_close(h_a, 1000, "bye").unwrap();
pool.initiate_close(h_b, 1000, "bye").unwrap();
for _ in 0..50 {
let _ = pool.pump_nowait(|_, _| {});
let done_a = matches!(pool.state(h_a), Some(State::Closed));
let done_b = matches!(pool.state(h_b), Some(State::Closed));
if done_a && done_b {
break;
}
let _ = pool.pump(|_, _| {});
}
server_a.join().unwrap();
server_b.join().unwrap();
}
}