use std::collections::{hash_map, HashMap, HashSet};
use std::io::prelude::*;
use std::net::TcpStream;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::{Arc, Mutex, MutexGuard, OnceLock, RwLock, RwLockReadGuard, RwLockWriteGuard};
use std::thread::{self, JoinHandle};
use std::time::Duration;
use crossbeam::channel::{self, Receiver, Sender};
use log::{debug, error, info, trace, warn};
use crate::accounts::types::AccountId;
use crate::client::handshake_state::HandshakeState;
use crate::client::ids::{OrderId, RequestId, WireId};
use crate::connection::sync::Connection;
use super::common::{lock, log_orphan, read_lock, report_unroutable_frame, validate_frame_length, write_lock, Lease, LeaseRef};
use super::order_taps::{NewTap, OrderTaps};
use super::rate_limiter::RateLimiter;
use super::raw_capture::RawFrameTap;
use super::routing::{
classify_error, determine_routing, is_foreign, order_routing_strategy, order_update_notice, DecodedError, ErrorDisposition, OrderRoutingStrategy,
RoutingDecision,
};
use super::{
Admit, BoundState, BufferBound, InternalSubscription, MessageBus, Response, RoutedItem, SharedCounts, SharedTicket, Signal, SubscriptionBuilder,
};
use crate::messages::{
is_cancel_frame, shared_channel_configuration, transport_reconnect_notice, IncomingMessages, Notice, OutgoingMessages, ResponseMessage,
};
use crate::subscriptions::notice_stream::sync_impl::NoticeStream;
use crate::Error;
const TWS_READ_TIMEOUT: Duration = Duration::from_secs(1);
const MID_FRAME_STALL: Duration = Duration::from_secs(30);
const MID_FRAME_TIMEOUT_LIMIT: u32 = (MID_FRAME_STALL.as_millis() / TWS_READ_TIMEOUT.as_millis()) as u32;
const _: () = assert!(MID_FRAME_TIMEOUT_LIMIT > 0, "MID_FRAME_STALL must exceed TWS_READ_TIMEOUT");
const STARTUP_STALL: Duration = Duration::from_secs(30);
pub(crate) const STARTUP_TIMEOUT_LIMIT: u32 = (STARTUP_STALL.as_millis() / TWS_READ_TIMEOUT.as_millis()) as u32;
const _: () = assert!(STARTUP_TIMEOUT_LIMIT > 0, "STARTUP_STALL must exceed TWS_READ_TIMEOUT");
const BACKLOG_WATERMARK: usize = 10_000;
fn backlog_watermark_crossed(depth: usize) -> bool {
depth > 0 && depth.is_multiple_of(BACKLOG_WATERMARK)
}
pub(super) fn warn_if_backlogged(label: std::fmt::Arguments<'_>, depth: usize) {
if backlog_watermark_crossed(depth) {
warn!("{label} at {depth} messages and growing — consumer is stalling");
}
}
#[derive(Debug)]
struct SharedSubscriber {
request: OutgoingMessages,
responses: &'static [IncomingMessages],
sender: Sender<RoutedItem>,
lease: LeaseRef,
}
impl SharedSubscriber {
fn receives(&self, message_type: IncomingMessages) -> bool {
self.responses.contains(&message_type)
}
}
#[derive(Debug)]
struct SharedChannels {
response_types: HashSet<IncomingMessages>,
subscribers: Mutex<Vec<SharedSubscriber>>,
counts: Mutex<SharedCounts>,
}
impl SharedChannels {
pub fn new() -> Self {
Self {
response_types: shared_channel_configuration::CHANNEL_MAPPINGS
.iter()
.flat_map(|mapping| mapping.responses.iter().copied())
.collect(),
subscribers: Mutex::new(Vec::new()),
counts: Mutex::new(SharedCounts::default()),
}
}
fn subscribers(&self) -> MutexGuard<'_, Vec<SharedSubscriber>> {
lock(&self.subscribers)
}
fn add(&self, request: OutgoingMessages, sender: Sender<RoutedItem>, lease: LeaseRef) {
let responses = shared_channel_configuration::response_types(request)
.unwrap_or_else(|| panic!("unsupported request message {request:?}. check mapping in messages::shared_channel_configuration"));
self.subscribers().push(SharedSubscriber {
request,
responses,
sender,
lease,
});
}
fn remove(&self, lease: &LeaseRef) {
let mut subscribers = self.subscribers();
let before = subscribers.len();
subscribers.retain(|subscriber| !subscriber.lease.is(lease));
debug!("cleanup shared subscription: removed={}", before != subscribers.len());
}
fn is_shared_response(&self, message_type: IncomingMessages) -> bool {
self.response_types.contains(&message_type)
}
fn subscribe(
&self,
message_type: OutgoingMessages,
account: Option<&AccountId>,
write: impl FnOnce() -> Result<(), Error>,
) -> Result<SharedTicket, Error> {
let mut counts = lock(&self.counts);
counts.check_account_updates(account)?;
write()?;
Ok(counts.subscribe(message_type, account))
}
fn unsubscribe(&self, ticket: SharedTicket, write: impl FnOnce() -> Result<(), Error>) -> Result<(), Error> {
let mut counts = lock(&self.counts);
if counts.unsubscribe(ticket) {
write()
} else {
Ok(())
}
}
fn reset_counts(&self) {
lock(&self.counts).reset();
}
fn send_to<F, I>(&self, filter: F, item: I) -> usize
where
F: Fn(&SharedSubscriber) -> bool,
I: Fn() -> RoutedItem,
{
let mut selected = 0;
self.subscribers().retain(|subscriber| {
if !filter(subscriber) {
return true;
}
selected += 1;
match subscriber.sender.send(item()) {
Ok(()) => {
warn_if_backlogged(format_args!("shared channel for {:?}", subscriber.request), subscriber.sender.len());
true
}
Err(_) => {
debug!("shared subscription {:?} dropped: removed", subscriber.request);
false
}
}
});
selected
}
fn send_message(&self, message_type: IncomingMessages, message: &ResponseMessage) -> bool {
let delivered = self.send_to(|subscriber| subscriber.receives(message_type), || message.clone().into()) > 0;
if !delivered {
trace!("no shared subscription for {message_type:?}: frame dropped");
}
delivered
}
fn fail_all<F>(&self, message_fn: F)
where
F: Fn() -> RoutedItem,
{
let mut subscribers = self.subscribers();
for subscriber in subscribers.iter() {
let _ = subscriber.sender.send(message_fn());
}
subscribers.clear();
}
fn fail_one_shot_channels<F>(&self, error_fn: F)
where
F: Fn() -> RoutedItem,
{
let one_shot = shared_channel_configuration::exclusive_one_shot_response_types();
self.send_to(|subscriber| subscriber.responses.iter().any(|r| one_shot.contains(r)), error_fn);
}
}
#[derive(Debug)]
pub(crate) struct NoticeBroadcaster {
senders: Mutex<Option<Vec<Sender<Notice>>>>,
}
impl super::common::NoticeSink for NoticeBroadcaster {
fn deliver(&self, notice: Notice) {
self.broadcast(notice);
}
}
impl NoticeBroadcaster {
pub(crate) fn new() -> Self {
Self {
senders: Mutex::new(Some(Vec::new())),
}
}
pub(crate) fn subscribe(&self) -> Receiver<Notice> {
let (sender, receiver) = channel::unbounded();
if let Some(senders) = lock(&self.senders).as_mut() {
senders.push(sender);
}
receiver
}
pub(crate) fn broadcast(&self, notice: Notice) {
if let Some(senders) = lock(&self.senders).as_mut() {
senders.retain(|s| {
let sent = s.send(notice.clone()).is_ok();
if sent {
warn_if_backlogged(format_args!("notice stream"), s.len());
}
sent
});
}
}
pub(crate) fn close(&self) {
*lock(&self.senders) = None;
}
}
#[derive(Debug)]
pub struct TcpMessageBus<S: Stream> {
connection: Connection<S>,
handles: Mutex<Vec<JoinHandle<()>>>,
requests: SenderHash<RequestId, RoutedItem>,
orders: SenderHash<OrderId, RoutedItem>,
executions: SenderHash<String, RoutedItem>,
shared_channels: SharedChannels,
signals_send: Sender<Signal>,
signals_recv: Receiver<Signal>,
shutdown_send: Sender<()>,
shutdown_recv: Receiver<()>,
shutdown: Arc<ShutdownSignal>,
handshake_state: OnceLock<HandshakeState>,
session_generation: AtomicU64,
rate_limiter: OnceLock<RateLimiter>,
order_taps: OrderTaps<Sender<RoutedItem>>,
connection_state: ConnectionSignal,
}
impl<S: Stream> TcpMessageBus<S> {
pub fn new(connection: Connection<S>) -> TcpMessageBus<S> {
let order_taps = OrderTaps::new();
let (signals_send, signals_recv) = channel::unbounded();
let (shutdown_send, shutdown_recv) = channel::bounded(1);
let shutdown = connection.shutdown_signal();
TcpMessageBus {
connection,
handles: Mutex::new(Vec::default()),
requests: SenderHash::new(),
orders: SenderHash::new(),
executions: SenderHash::new(),
shared_channels: SharedChannels::new(),
signals_send,
signals_recv,
shutdown_send,
shutdown_recv,
shutdown,
handshake_state: OnceLock::new(),
session_generation: AtomicU64::new(0),
rate_limiter: OnceLock::new(),
order_taps,
connection_state: ConnectionSignal::default(),
}
}
pub(crate) fn set_handshake_state(&self, state: HandshakeState) {
self.handshake_state.set(state).expect("handshake state installed twice");
}
pub(crate) fn set_rate_limiter(&self, limiter: RateLimiter) {
self.rate_limiter.set(limiter).expect("rate limiter installed twice");
}
fn is_shutting_down(&self) -> bool {
self.shutdown.is_requested()
}
fn request_shutdown(&self) {
debug!("shutdown requested");
self.requests.fail_all(|| Error::Shutdown.into());
self.orders.fail_all(|| Error::Shutdown.into());
self.shared_channels.fail_all(|| Error::Shutdown.into());
self.executions.clear();
self.connection.notice_broadcaster.close();
self.connection_state.shutdown();
self.shutdown.request();
self.order_taps.close();
let _ = self.shutdown_send.try_send(());
self.connection.socket.close();
}
fn ensure_connected(&self) -> Result<(), Error> {
self.connection_state.check()?;
if self.is_shutting_down() {
return Err(Error::Shutdown);
}
Ok(())
}
fn ended_reason(&self, refusal: Error) -> Error {
self.connection_state.check_ended().err().unwrap_or(refusal)
}
fn write_message(&self, message: &[u8]) -> Result<(), Error> {
self.throttle(message);
self.write_now(message)
}
fn throttle(&self, message: &[u8]) {
if let Some(limiter) = self.rate_limiter.get() {
limiter.acquire_blocking(is_cancel_frame(message));
}
}
fn write_now(&self, message: &[u8]) -> Result<(), Error> {
self.ensure_connected()?;
self.connection.write_message(message)
}
fn reset(&self) {
debug!("reset message bus");
self.session_generation.fetch_add(1, Ordering::AcqRel);
self.requests.fail_all(|| Error::ConnectionReset.into());
self.orders.fail_all(|| Error::ConnectionReset.into());
self.shared_channels.fail_all(|| Error::ConnectionReset.into());
self.shared_channels.reset_counts();
self.executions.clear();
self.order_taps.reset();
}
fn clean_request(&self, request_id: RequestId, lease: &LeaseRef) {
let removed = self.requests.remove_if_same(request_id, lease);
let aliases = self.executions.remove_all_same(lease);
debug!(
"cleanup request_id {request_id}: removed={removed}, aliases={aliases}, requests.len()={}",
self.requests.len()
);
}
fn clean_order(&self, order_id: OrderId, lease: &LeaseRef) {
let removed = self.orders.remove_if_same(order_id, lease);
let aliases = self.executions.remove_all_same(lease);
debug!(
"cleanup order_id {order_id}: removed={removed}, aliases={aliases}, orders.len()={}",
self.orders.len()
);
}
fn clear_order_update_stream(&self, lease: &LeaseRef) {
let removed = self.order_taps.release_updates(lease);
debug!("cleanup order_update_stream: removed={removed}");
}
fn read_message(&self) -> Response {
self.connection.read_message()
}
pub(crate) fn dispatch(&self) -> Result<(), Error> {
match self.read_message() {
Ok(message) => {
if message.is_shutdown() {
self.request_shutdown();
Err(Error::Shutdown)
} else {
self.dispatch_message(message);
Ok(())
}
}
Err(ref err) if err.is_read_timeout() => {
if self.is_shutting_down() {
debug!("dispatcher thread exiting");
return Err(Error::Shutdown);
}
Ok(())
}
Err(ref err) if err.is_connection_lost() => {
if self.is_shutting_down() {
debug!("dispatcher thread exiting");
return Err(Error::Shutdown);
}
error!("error reading next message (will attempt reconnect): {err:?}");
self.connection_state.set_disconnected();
self.reset();
match self.connection.reconnect() {
Ok(()) => {}
Err(Error::Shutdown) => {
debug!("shutdown requested during reconnect; dispatcher thread exiting");
return Err(Error::Shutdown);
}
Err(reconnect_err) => {
error!("failed to reconnect to TWS/Gateway: {reconnect_err:?}");
self.connection_state.fail();
self.request_shutdown();
return Err(Error::ConnectionFailed);
}
}
if self.is_shutting_down() {
debug!("shutdown requested during reconnect; dispatcher thread exiting");
return Err(Error::Shutdown);
}
if let Some(state) = self.handshake_state.get() {
state.refresh(&self.connection.connection_metadata());
}
info!("successfully reconnected to TWS/Gateway");
self.connection_state.set_connected();
self.connection.notice_broadcaster.broadcast(transport_reconnect_notice());
Ok(())
}
Err(Error::Shutdown) => {
debug!("socket closed; dispatcher thread exiting");
self.request_shutdown();
Err(Error::Shutdown)
}
Err(err) => {
error!("error reading next message (shutting down): {err:?}");
self.connection_state.fail();
self.request_shutdown();
Err(err)
}
}
}
fn start_dispatcher_thread(self: &Arc<Self>) -> JoinHandle<()> {
let message_bus = Arc::clone(self);
thread::spawn(move || {
loop {
match message_bus.dispatch() {
Ok(_) => {}
Err(Error::Shutdown | Error::ConnectionFailed) => break,
Err(e) => {
error!("Dispatcher encountered an error: {e:?}");
break;
}
}
}
debug!("Dispatcher thread finished.");
})
}
fn dispatch_message(&self, message: ResponseMessage) {
match determine_routing(&message) {
RoutingDecision::Error(payload) => self.route_error_message(payload),
RoutingDecision::ByOrderId(_) => {
self.process_orders(message);
}
RoutingDecision::ByRequestId(id) => {
self.process_response_with_id(WireId::classify(id), message, false);
}
_ => {
self.process_response(message, false);
}
}
}
fn route_error_message(&self, payload: DecodedError) {
let sent_to_update_stream = order_update_notice(&payload).is_some_and(|notice| self.send_order_update_item(RoutedItem::Notice(notice)));
match classify_error(payload) {
ErrorDisposition::NoticeOnly(notice) => {
notice.log();
self.connection.notice_broadcaster.broadcast(notice);
}
ErrorDisposition::NoticeAndFailOneShots(notice, error) => {
notice.log();
self.connection.notice_broadcaster.broadcast(notice);
self.shared_channels.fail_one_shot_channels(|| RoutedItem::Error(error.clone()));
}
ErrorDisposition::Route(id, item) => {
self.order_taps.publish_rejection(id, &item);
self.deliver(id, item, sent_to_update_stream);
}
}
}
fn process_response(&self, message: ResponseMessage, routed: bool) {
let id = message.request_id().and_then(WireId::classify);
self.process_response_with_id(id, message, routed);
}
fn process_response_with_id(&self, id: Option<WireId>, message: ResponseMessage, routed: bool) {
match id {
Some(WireId::Request(request_id)) if self.requests.contains(&request_id) => {
let _ = self.requests.deliver(&request_id, message.into());
return;
}
Some(WireId::Order(order_id)) if self.orders.contains(&order_id) => {
let _ = self.orders.deliver(&order_id, message.into());
return;
}
_ => {}
}
if self.shared_channels.is_shared_response(message.message_type()) {
self.shared_channels.send_message(message.message_type(), &message);
} else if !routed {
report_unroutable_frame(&message, &*self.connection.notice_broadcaster);
}
}
fn deliver(&self, id: WireId, item: RoutedItem, sent_to_update_stream: bool) {
let unrouted = match id {
WireId::Request(request_id) => self.requests.deliver(&request_id, item),
WireId::Order(order_id) => self.orders.deliver(&order_id, item),
};
if let Err(item) = unrouted {
if !sent_to_update_stream {
log_orphan(id, &item);
}
}
}
fn process_orders(&self, message: ResponseMessage) {
let foreign = is_foreign(&message, self.connection.client_id);
let strategy = order_routing_strategy(message.message_type(), foreign);
let message_order_id = message.order_id().map(OrderId::from);
let message_request_id = message.request_id().and_then(RequestId::from_raw);
if !foreign {
self.order_taps.publish_status(&message);
}
match strategy {
OrderRoutingStrategy::OrderUpdateOnly => {
self.send_order_update(&message);
}
OrderRoutingStrategy::ExecutionData => {
let sent_to_update_stream = self.send_order_update(&message);
let execution_id = message.execution_id();
if let Err(item) = self.deliver_to_order_or_request(message_order_id, message_request_id, message.into(), execution_id.as_ref()) {
if !sent_to_update_stream {
warn!("could not route message {item:?}");
}
}
}
OrderRoutingStrategy::ExecutionDataEnd => {
if let Err(item) = self.deliver_to_order_or_request(message_order_id, message_request_id, message.into(), None) {
warn!("could not route message {item:?}");
}
}
OrderRoutingStrategy::OrderAndShared => {
let sent_to_update_stream = self.send_order_update(&message);
let shared = self.shared_channels.send_message(message.message_type(), &message);
let unrouted = match message_order_id {
Some(order_id) => self.orders.deliver(&order_id, message.into()).err(),
None => Some(message.into()),
};
if let Some(item) = unrouted {
if !shared && !sent_to_update_stream {
warn!("could not route message {item:?}");
}
}
}
OrderRoutingStrategy::ByExecutionId => {
let sent_to_update_stream = self.send_order_update(&message);
let unrouted = match message.execution_id() {
Some(execution_id) => self.executions.deliver(&execution_id, message.into()),
None => Err(message.into()),
};
if let Err(item) = unrouted {
if !sent_to_update_stream {
warn!("could not route commission report {item:?}");
}
}
}
OrderRoutingStrategy::SharedOnly => {
self.shared_channels.send_message(message.message_type(), &message);
}
OrderRoutingStrategy::ByRequestId => {
let execution_id = message.execution_id();
if let Err(item) = self.deliver_to_order_or_request(None, message_request_id, message.into(), execution_id.as_ref()) {
debug!("another client's execution answers no request: {item:?}");
}
}
OrderRoutingStrategy::Ignore => {}
OrderRoutingStrategy::ByOrderId => {
warn!("unhandled order message type: {message:?}");
}
}
}
fn open_request(&self, request_id: RequestId, message: &[u8], bound: Option<BufferBound>) -> Result<InternalSubscription, Error> {
let session = self.session_generation();
self.ensure_connected()?;
let (sender, receiver) = channel::unbounded();
let lease = Lease::new();
let lease_ref = lease.downgrade();
match bound {
Some(bound) => self.requests.insert_bounded(request_id, sender.clone(), lease_ref.clone(), bound),
None => self.requests.insert(request_id, sender.clone(), lease_ref.clone()),
};
if let Err(e) = self.write_message(message) {
self.requests.remove_if_same(request_id, &lease_ref);
return Err(e);
}
let subscription = SubscriptionBuilder::new()
.receiver(receiver)
.sender(sender)
.signaler(self.signals_send.clone())
.lease(lease)
.request_id(request_id)
.session(session)
.build();
Ok(subscription)
}
fn deliver_to_order_or_request(
&self,
order_id: Option<OrderId>,
request_id: Option<RequestId>,
item: RoutedItem,
execution_id: Option<&String>,
) -> Result<(), RoutedItem> {
match order_id {
Some(id) => self.orders.deliver_aliased(&id, item, execution_id, &self.executions),
None => Err(item),
}
.or_else(|item| match request_id {
Some(id) => self.requests.deliver_aliased(&id, item, execution_id, &self.executions),
None => Err(item),
})
}
fn send_order_update(&self, message: &ResponseMessage) -> bool {
self.send_order_update_item(message.clone().into())
}
fn send_order_update_item(&self, item: RoutedItem) -> bool {
self.order_taps.send_update(item)
}
fn start_cleanup_thread(self: &Arc<Self>) -> JoinHandle<()> {
let message_bus = Arc::clone(self);
thread::spawn(move || {
let signal_recv = message_bus.signals_recv.clone();
let shutdown_recv = message_bus.shutdown_recv.clone();
loop {
crossbeam::select! {
recv(signal_recv) -> signal => match signal {
Ok(Signal::Request(request_id, lease)) => message_bus.clean_request(request_id, &lease),
Ok(Signal::Order(order_id, lease)) => message_bus.clean_order(order_id, &lease),
Ok(Signal::OrderUpdateStream(lease)) => message_bus.clear_order_update_stream(&lease),
Ok(Signal::OrderStatus(order_id, lease)) => message_bus.order_taps.release_status(order_id, &lease),
Ok(Signal::Shared(lease)) => message_bus.shared_channels.remove(&lease),
Err(_) => {
debug!("cleanup signal channel closed");
return;
}
},
recv(shutdown_recv) -> _ => {
debug!("cleanup thread exiting");
return;
}
}
}
})
}
pub(crate) fn process_messages(self: &Arc<Self>) -> Result<(), Error> {
let handle = self.start_dispatcher_thread();
self.add_join_handle(handle);
let handle = self.start_cleanup_thread();
self.add_join_handle(handle);
Ok(())
}
fn add_join_handle(&self, handle: JoinHandle<()>) {
let mut handles = lock(&self.handles);
handles.push(handle);
}
pub fn join(&self) {
let mut handles = lock(&self.handles);
for handle in handles.drain(..) {
if let Err(e) = handle.join() {
warn!("could not join thread: {e:?}");
}
}
}
fn send_shared(&self, message_type: OutgoingMessages, account: Option<&AccountId>, message: &[u8]) -> Result<InternalSubscription, Error> {
let session = self.session_generation();
self.ensure_connected()?;
let (sender, receiver) = channel::unbounded();
let lease = Lease::new();
let lease_ref = lease.downgrade();
self.shared_channels.add(message_type, sender.clone(), lease_ref.clone());
self.throttle(message);
let ticket = match self.shared_channels.subscribe(message_type, account, || self.write_now(message)) {
Ok(ticket) => ticket,
Err(e) => {
self.shared_channels.remove(&lease_ref);
return Err(e);
}
};
let subscription = SubscriptionBuilder::new()
.receiver(receiver)
.sender(sender)
.signaler(self.signals_send.clone())
.lease(lease)
.shared(ticket)
.session(session)
.build();
Ok(subscription)
}
}
impl<S: Stream> MessageBus for TcpMessageBus<S> {
fn send_request_capped(&self, request_id: RequestId, message: &[u8], bound: Option<BufferBound>) -> Result<InternalSubscription, Error> {
self.open_request(request_id, message, bound)
}
fn send_order_request(&self, order_id: OrderId, message: &[u8]) -> Result<InternalSubscription, Error> {
let session = self.session_generation();
self.ensure_connected()?;
let (sender, receiver) = channel::unbounded();
let lease = Lease::new();
let lease_ref = lease.downgrade();
self.orders.insert(order_id, sender.clone(), lease_ref.clone());
debug!("Registered order subscription for order_id={}", order_id);
if let Err(e) = self.write_message(message) {
self.orders.remove_if_same(order_id, &lease_ref);
return Err(e);
}
let subscription = SubscriptionBuilder::new()
.receiver(receiver)
.sender(sender)
.signaler(self.signals_send.clone())
.lease(lease)
.order_id(order_id)
.session(session)
.build();
Ok(subscription)
}
fn send_message(&self, message: &[u8]) -> Result<(), Error> {
self.write_message(message)?;
Ok(())
}
fn send_market_data_type(&self, message: &[u8]) -> Result<(), Error> {
self.connection.remember_market_data_type(message.to_vec());
self.write_message(message)
}
fn create_order_status_subscription(&self, order_id: OrderId) -> Result<InternalSubscription, Error> {
let NewTap { sender, receiver, lease } = self.order_taps.subscribe_status(order_id, 0).map_err(|e| self.ended_reason(e))?;
Ok(SubscriptionBuilder::new()
.receiver(receiver)
.sender(sender)
.signaler(self.signals_send.clone())
.lease(lease)
.order_status(order_id)
.session(self.session_generation())
.build())
}
fn create_order_update_subscription(&self) -> Result<InternalSubscription, Error> {
let NewTap { sender, receiver, lease } = self.order_taps.subscribe_updates(0).map_err(|e| self.ended_reason(e))?;
Ok(SubscriptionBuilder::new()
.receiver(receiver)
.sender(sender)
.signaler(self.signals_send.clone())
.lease(lease)
.session(self.session_generation())
.build())
}
fn send_shared_request(&self, message_type: OutgoingMessages, message: &[u8]) -> Result<InternalSubscription, Error> {
self.send_shared(message_type, None, message)
}
fn send_account_updates_request(&self, account: &AccountId, message: &[u8]) -> Result<InternalSubscription, Error> {
self.send_shared(OutgoingMessages::RequestAccountData, Some(account), message)
}
fn cancel_shared_subscription(&self, ticket: SharedTicket, message: Option<&[u8]>) -> Result<(), Error> {
if let Some(message) = message {
self.throttle(message);
}
self.shared_channels
.unsubscribe(ticket, || message.map_or(Ok(()), |message| self.write_now(message)))
}
fn notice_subscribe(&self) -> NoticeStream {
NoticeStream::new(self.connection.notice_broadcaster.subscribe())
}
fn ensure_shutdown(&self) {
self.request_shutdown();
self.join();
}
fn wait_connected(&self) -> Result<(), Error> {
self.connection_state.wait_connected()
}
fn is_connected(&self) -> bool {
self.connection_state.is_connected() && !self.is_shutting_down()
}
fn session_generation(&self) -> u64 {
self.session_generation.load(Ordering::Acquire)
}
}
#[derive(Debug)]
struct Entry<V> {
sender: Sender<V>,
lease: LeaseRef,
bound: Option<BoundState>,
}
impl<V> Entry<V> {
fn new(sender: Sender<V>, lease: LeaseRef) -> Self {
Self { sender, lease, bound: None }
}
fn closed(&self) -> bool {
self.bound.as_ref().is_some_and(BoundState::closed)
}
fn alias(&self) -> Self {
Self::new(self.sender.clone(), self.lease.clone())
}
}
impl Entry<RoutedItem> {
fn deliver(&self, id: &impl std::fmt::Debug, item: RoutedItem) {
let item = match &self.bound {
Some(bound) => match bound.admit(&item, self.sender.len()) {
Admit::Deliver => item,
Admit::Overflow => bound.overflow_error(),
Admit::Discard => {
trace!("discarding item for closed route {id:?}");
return;
}
},
None => item,
};
if self.sender.send(item).is_err() {
trace!("discarding item for dropped subscription {id:?}");
} else {
warn_if_backlogged(format_args!("subscription queue for {id:?}"), self.sender.len());
}
}
}
#[derive(Debug)]
struct SenderHash<K, V> {
senders: RwLock<HashMap<K, Entry<V>>>,
}
impl<K: std::hash::Hash + Eq + std::fmt::Debug, V: std::fmt::Debug> SenderHash<K, V> {
pub fn new() -> Self {
Self {
senders: RwLock::new(HashMap::new()),
}
}
fn read(&self) -> RwLockReadGuard<'_, HashMap<K, Entry<V>>> {
read_lock(&self.senders)
}
fn write(&self) -> RwLockWriteGuard<'_, HashMap<K, Entry<V>>> {
write_lock(&self.senders)
}
#[cfg(test)]
pub fn copy_sender(&self, id: K) -> Option<Sender<V>> {
self.with_entry(&id, |entry| entry.sender.clone())
}
#[cfg(test)]
pub fn lease(&self, id: K) -> Option<LeaseRef> {
self.with_entry(&id, |entry| entry.lease.clone())
}
#[cfg(test)]
fn with_entry<R>(&self, id: &K, f: impl FnOnce(&Entry<V>) -> R) -> Option<R> {
let senders = self.read();
senders.get(id).map(f)
}
pub fn insert(&self, id: K, sender: Sender<V>, lease: LeaseRef) {
self.insert_entry(id, Entry::new(sender, lease));
}
pub fn insert_bounded(&self, id: K, sender: Sender<V>, lease: LeaseRef, bound: BufferBound) {
let entry = Entry {
bound: Some(BoundState::new(bound)),
..Entry::new(sender, lease)
};
self.insert_entry(id, entry);
}
fn insert_entry(&self, id: K, entry: Entry<V>) {
self.write().insert(id, entry);
}
pub fn remove_if_same(&self, id: K, lease: &LeaseRef) -> bool {
match self.write().entry(id) {
hash_map::Entry::Occupied(registered) if registered.get().lease.is(lease) => {
registered.remove();
true
}
_ => false,
}
}
pub fn remove_all_same(&self, lease: &LeaseRef) -> usize {
let mut senders = self.write();
let before = senders.len();
senders.retain(|_, registered| !registered.lease.is(lease));
before - senders.len()
}
pub fn contains(&self, id: &K) -> bool {
let senders = self.read();
senders.contains_key(id)
}
pub fn len(&self) -> usize {
let senders = self.read();
senders.len()
}
pub fn clear(&self) {
let mut senders = self.write();
senders.clear();
}
}
impl<K: std::hash::Hash + Eq + std::fmt::Debug> SenderHash<K, RoutedItem> {
pub fn deliver(&self, id: &K, item: RoutedItem) -> Result<(), RoutedItem> {
self.deliver_then(id, item, |_| {})
}
pub fn deliver_aliased<A: std::hash::Hash + Eq + Clone + std::fmt::Debug>(
&self,
id: &K,
item: RoutedItem,
alias: Option<&A>,
aliases: &SenderHash<A, RoutedItem>,
) -> Result<(), RoutedItem> {
self.deliver_then(id, item, |entry| {
if let Some(alias) = alias {
aliases.insert_entry(alias.clone(), entry.alias());
}
})
}
fn deliver_then(&self, id: &K, item: RoutedItem, then: impl FnOnce(&Entry<RoutedItem>)) -> Result<(), RoutedItem> {
let senders = self.read();
let Some(entry) = senders.get(id) else {
return Err(item);
};
entry.deliver(id, item);
then(entry);
Ok(())
}
pub fn fail_all<F>(&self, message_fn: F)
where
F: Fn() -> RoutedItem,
{
let mut senders = self.write();
for entry in senders.values().filter(|entry| !entry.closed()) {
let _ = entry.sender.send(message_fn());
}
senders.clear();
}
}
#[derive(Debug)]
struct CloseHandle {
stream: TcpStream,
closed: bool,
}
#[derive(Debug)]
pub(crate) struct TcpSocket {
reader: Mutex<TcpStream>,
writer: Mutex<TcpStream>,
close_handle: Mutex<CloseHandle>,
connection_url: String,
tcp_no_delay: bool,
tap: RawFrameTap,
}
impl TcpSocket {
pub fn connect(address: &str, tcp_no_delay: bool) -> Result<Self, Error> {
let stream = TcpStream::connect(address)?;
Self::new(stream, address, tcp_no_delay)
}
pub fn new(stream: TcpStream, connection_url: &str, tcp_no_delay: bool) -> Result<Self, Error> {
let writer = stream.try_clone()?;
let close_handle = CloseHandle {
stream: stream.try_clone()?,
closed: false,
};
stream.set_read_timeout(Some(TWS_READ_TIMEOUT))?;
stream.set_nodelay(tcp_no_delay)?;
Ok(Self {
reader: Mutex::new(stream),
writer: Mutex::new(writer),
close_handle: Mutex::new(close_handle),
connection_url: connection_url.to_string(),
tcp_no_delay,
tap: RawFrameTap::from_env(),
})
}
fn closed_or(&self, err: Error) -> Error {
if lock(&self.close_handle).closed {
Error::Shutdown
} else {
err
}
}
}
impl Reconnect for TcpSocket {
fn reconnect(&self) -> Result<(), Error> {
let stream = TcpStream::connect(&self.connection_url)?;
stream.set_read_timeout(Some(TWS_READ_TIMEOUT))?;
stream.set_nodelay(self.tcp_no_delay)?;
let mut reader = self.reader.lock()?;
let mut writer = self.writer.lock()?;
let mut handle = self.close_handle.lock()?;
if handle.closed {
return Err(Error::Shutdown);
}
*reader = stream.try_clone()?;
*writer = stream.try_clone()?;
handle.stream = stream;
self.tap.start_new_segment();
Ok(())
}
fn sleep(&self, duration: std::time::Duration, shutdown: &ShutdownSignal) {
shutdown.wait_timeout(duration)
}
}
pub(crate) trait Reconnect {
fn reconnect(&self) -> Result<(), Error>;
fn sleep(&self, duration: std::time::Duration, shutdown: &ShutdownSignal);
}
pub(crate) trait Stream: Io + Reconnect + Sync + Send + 'static + std::fmt::Debug {}
impl Stream for TcpSocket {}
struct FrameRead<'a, R> {
reader: &'a mut R,
started: bool,
timeouts: u32,
}
impl<'a, R: Read> FrameRead<'a, R> {
fn new(reader: &'a mut R) -> Self {
Self {
reader,
started: false,
timeouts: 0,
}
}
fn fill(&mut self, buf: &mut [u8]) -> Result<(), Error> {
let mut filled = 0;
while filled < buf.len() {
match self.reader.read(&mut buf[filled..]) {
Ok(0) => return Err(std::io::Error::new(std::io::ErrorKind::UnexpectedEof, "failed to fill whole buffer").into()),
Ok(n) => {
filled += n;
self.started = true;
self.timeouts = 0;
}
Err(e) if e.kind() == std::io::ErrorKind::Interrupted => {}
Err(e) if matches!(e.kind(), std::io::ErrorKind::WouldBlock | std::io::ErrorKind::TimedOut) => {
if !self.started {
return Err(e.into());
}
self.timeouts += 1;
if self.timeouts > MID_FRAME_TIMEOUT_LIMIT {
return Err(Error::InvalidFrame(format!("frame stalled mid-read for {} read timeouts", self.timeouts)));
}
}
Err(e) => return Err(e.into()),
}
}
Ok(())
}
}
fn read_header(frame: &mut FrameRead<'_, impl Read>, tap: &RawFrameTap) -> Result<usize, Error> {
let mut buffer = [0_u8; 4];
frame.fill(&mut buffer)?;
tap.record_length_prefix(&buffer);
validate_frame_length(u32::from_be_bytes(buffer) as usize)
}
pub(crate) fn read_message(reader: &mut impl Read, tap: &RawFrameTap) -> Result<Vec<u8>, Error> {
let mut frame = FrameRead::new(reader);
let message_size = read_header(&mut frame, tap)?;
let mut data = vec![0_u8; message_size];
frame.fill(&mut data)?;
tap.record_body(&data);
Ok(data)
}
impl Io for TcpSocket {
fn read_message(&self) -> Result<Vec<u8>, Error> {
let mut reader = self.reader.lock()?;
read_message(&mut *reader, &self.tap).map_err(|err| self.closed_or(err))
}
fn write_all(&self, buf: &[u8]) -> Result<(), Error> {
let mut writer = self.writer.lock()?;
writer.write_all(buf).map_err(|err| self.closed_or(err.into()))
}
fn close(&self) {
let mut handle = lock(&self.close_handle);
if std::mem::replace(&mut handle.closed, true) {
return;
}
if let Err(e) = handle.stream.shutdown(std::net::Shutdown::Both) {
debug!("socket shutdown returned: {e:?}");
}
}
}
pub(crate) trait Io {
fn read_message(&self) -> Result<Vec<u8>, Error>;
fn write_all(&self, buf: &[u8]) -> Result<(), Error>;
fn close(&self);
}
mod connection_signal;
mod shutdown;
pub(crate) use connection_signal::ConnectionSignal;
pub(crate) use shutdown::ShutdownSignal;
#[cfg(test)]
mod memory;
#[cfg(test)]
pub(crate) use memory::MemoryStream;
#[cfg(test)]
pub(crate) mod test_listener;
#[cfg(test)]
#[path = "sync_tests.rs"]
mod tests;