use std::{net::SocketAddr, sync::Arc};
use anyhow::{Context, Result};
use async_trait::async_trait;
use serde_json::Value;
use tokio::{
io::{AsyncWriteExt, BufReader},
net::{
tcp::{OwnedReadHalf, OwnedWriteHalf},
TcpStream,
},
sync::{mpsc::Sender, Mutex},
};
use crate::{
config::Config,
doslimit::{ConnectionLimits, GlobalLimits},
query::Query,
rpc::{daemon::connection::Connection, rpcstats::RpcStats},
signal::NetworkNotifier,
util::Channel,
};
use super::{
communication::Communcation,
parse_requests::{parse_requests, tcp_line_validator},
server::ConnectionId,
Message,
};
use crate::query::scripthash_subscriptions::ScripthashSubscriptions;
pub struct TcpCommuncation {
reader: Option<OwnedReadHalf>,
writer: OwnedWriteHalf,
}
impl TcpCommuncation {
pub fn new(stream: TcpStream) -> Self {
let (reader, writer) = stream.into_split();
Self {
reader: Some(reader),
writer,
}
}
}
#[async_trait]
impl Communcation for TcpCommuncation {
async fn send_values(&mut self, values: &[Value]) -> Result<()> {
for value in values {
let line = value.to_string() + "\n";
if let Err(e) = self.writer.write_all(line.as_bytes()).await {
let truncated: String = line.chars().take(80).collect();
return Err(e).context(format!("failed to send {}", truncated));
}
}
Ok(())
}
fn start_request_receiver(&mut self, rpc_queue_sender: Sender<Message>, idle_timeout: u64) {
let reader = self.reader.take();
crate::thread::spawn_task("peer_reader", async move {
let bufreader = BufReader::new(reader.unwrap());
if let Err(e) = parse_requests(
bufreader,
rpc_queue_sender,
idle_timeout,
tcp_line_validator,
)
.await
{
trace!("peer_reader thread: {}", e);
}
Ok(())
});
}
async fn shutdown(&mut self) {
let _ = self.writer.shutdown().await;
}
}
#[allow(clippy::too_many_arguments)]
pub async fn start_tcp_server(
conn_acceptor: Channel<Option<(TcpStream, SocketAddr)>>,
global_limits: Arc<GlobalLimits>,
config: Arc<Config>,
query: Arc<Query>,
stats: Arc<RpcStats>,
connection_limits: ConnectionLimits,
network_notifier: Arc<NetworkNotifier>,
rpc_queue_senders: Arc<Mutex<std::collections::HashMap<ConnectionId, Sender<Message>>>>,
subscription_manager: Arc<ScripthashSubscriptions>,
) -> Result<()> {
while let Some(conn) = conn_acceptor.receiver().await.recv().await {
let (mut stream, addr) = match conn {
Some(c) => c,
None => break,
};
let global_limits = global_limits.clone();
let mut connections = match global_limits.inc_connection(&addr.ip()).await {
Err(e) => {
trace!("[{}] dropping tcp peer - {}", addr, e);
let _ = stream.shutdown().await;
drop(stream);
continue;
}
Ok(n) => n,
};
let query_cpy = Arc::clone(&query);
let stats_cpy = Arc::clone(&stats);
let network_notifier_cpy = Arc::clone(&network_notifier);
let config = Arc::clone(&config);
let rpc_queue = Channel::bounded(config.rpc_buffer_size);
let conn_id = ConnectionId::new();
let subscription_manager_cpy = Arc::clone(&subscription_manager);
let rpc_queue_senders_cpy = Arc::clone(&rpc_queue_senders);
rpc_queue_senders
.lock()
.await
.insert(conn_id, rpc_queue.sender());
crate::thread::spawn_task("tcp_peer", async move {
if connections != (0, 0) {
debug!(
"[{} tcp] connected peer ({:?} out of {:?} connection slots used)",
addr,
connections,
global_limits.connection_limits(),
);
}
let communication = TcpCommuncation::new(stream);
let conn = Connection::new(
config,
query_cpy,
addr,
stats_cpy,
connection_limits,
Box::new(communication),
network_notifier_cpy,
conn_id,
subscription_manager_cpy.clone(),
);
let conn_id_for_cleanup = conn_id;
conn.run(rpc_queue).await;
subscription_manager_cpy
.remove_connection(conn_id_for_cleanup)
.await;
rpc_queue_senders_cpy
.lock()
.await
.remove(&conn_id_for_cleanup);
match global_limits.dec_connection(&addr.ip()).await {
Ok(n) => connections = n,
Err(e) => error!("{}", e),
};
debug!(
"[{} tcp] disconnected peer ({:?} out of {:?} connection slots used)",
addr,
connections,
global_limits.connection_limits(),
);
Ok(())
});
}
info!("TCP RPC connections no longer accepted");
Ok(())
}