feat: publish Dever core with project scaffolding and versioned AI skills
This commit is contained in:
1 parent
70fa73cdc2
commit
ac458ffa7f
960 files changed
+221319
-1163
No files matched your search
@@ -0,0 +1,125 @@
|
||||
use std::time::Duration;
|
||||
|
||||
use hyper::Uri;
|
||||
|
||||
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
|
||||
pub enum Protocol {
|
||||
Tcp,
|
||||
WebSocket,
|
||||
Sse,
|
||||
}
|
||||
|
||||
impl Protocol {
|
||||
fn parse(value: &str) -> Result<Self, String> {
|
||||
match value {
|
||||
"tcp" => Ok(Self::Tcp),
|
||||
"ws" => Ok(Self::WebSocket),
|
||||
"sse" => Ok(Self::Sse),
|
||||
_ => Err("PROTOCOL must be tcp, ws, or sse".into()),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn name(self) -> &'static str {
|
||||
match self {
|
||||
Self::Tcp => "tcp",
|
||||
Self::WebSocket => "ws",
|
||||
Self::Sse => "sse",
|
||||
}
|
||||
}
|
||||
|
||||
fn scheme(self) -> &'static str {
|
||||
match self {
|
||||
Self::Tcp => "tcp",
|
||||
Self::WebSocket => "ws",
|
||||
Self::Sse => "http",
|
||||
}
|
||||
}
|
||||
|
||||
fn path(self) -> &'static str {
|
||||
match self {
|
||||
Self::Tcp => "/",
|
||||
Self::WebSocket => "/ws",
|
||||
Self::Sse => "/events",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct Target {
|
||||
pub port: u16,
|
||||
pub path: String,
|
||||
pub url: String,
|
||||
}
|
||||
|
||||
impl Target {
|
||||
fn parse(protocol: Protocol, value: &str) -> Result<Self, String> {
|
||||
let uri = value
|
||||
.parse::<Uri>()
|
||||
.map_err(|error| format!("URL is invalid: {error}"))?;
|
||||
if uri.scheme_str() != Some(protocol.scheme()) {
|
||||
return Err(format!(
|
||||
"{} URL must use the {} scheme",
|
||||
protocol.name(),
|
||||
protocol.scheme()
|
||||
));
|
||||
}
|
||||
if uri.host() != Some("127.0.0.1") {
|
||||
return Err("live benchmark URL must use 127.0.0.1".into());
|
||||
}
|
||||
let port = uri.port_u16().ok_or("URL must include an explicit port")?;
|
||||
if uri.path() != protocol.path() || uri.query().is_some() {
|
||||
return Err(format!(
|
||||
"{} URL must use the {} path without a query",
|
||||
protocol.name(),
|
||||
protocol.path()
|
||||
));
|
||||
}
|
||||
Ok(Self {
|
||||
port,
|
||||
path: uri.path().to_owned(),
|
||||
url: value.to_owned(),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct Arguments {
|
||||
pub protocol: Protocol,
|
||||
pub target: Target,
|
||||
pub connections: usize,
|
||||
pub cycles: usize,
|
||||
pub hold: Duration,
|
||||
pub timeout: Duration,
|
||||
}
|
||||
|
||||
impl Arguments {
|
||||
pub fn parse(values: Vec<String>) -> Result<Self, String> {
|
||||
let [protocol, url, connections, cycles, hold_ms, timeout_ms] = values.as_slice() else {
|
||||
return Err(crate::usage());
|
||||
};
|
||||
let protocol = Protocol::parse(protocol)?;
|
||||
let target = Target::parse(protocol, url)?;
|
||||
let connections = number(connections, "CONNECTIONS", 1, 4_096)? as usize;
|
||||
let cycles = number(cycles, "CYCLES", 1, 100)? as usize;
|
||||
let hold_ms = number(hold_ms, "HOLD_MS", 300, 60_000)?;
|
||||
let timeout_ms = number(timeout_ms, "TIMEOUT_MS", 1, 300_000)?;
|
||||
Ok(Self {
|
||||
protocol,
|
||||
target,
|
||||
connections,
|
||||
cycles,
|
||||
hold: Duration::from_millis(hold_ms),
|
||||
timeout: Duration::from_millis(timeout_ms),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
fn number(value: &str, name: &str, minimum: u64, maximum: u64) -> Result<u64, String> {
|
||||
let value = value
|
||||
.parse::<u64>()
|
||||
.map_err(|_| format!("{name} must be an integer"))?;
|
||||
if !(minimum..=maximum).contains(&value) {
|
||||
return Err(format!("{name} must be between {minimum} and {maximum}"));
|
||||
}
|
||||
Ok(value)
|
||||
}
|
||||
@@ -0,0 +1,325 @@
|
||||
use std::time::Duration;
|
||||
|
||||
use bytes::Bytes;
|
||||
use futures_util::{SinkExt, StreamExt};
|
||||
use http_body_util::{BodyExt, Empty};
|
||||
use hyper::body::Incoming;
|
||||
use hyper::client::conn::http1;
|
||||
use hyper::header::{CONTENT_TYPE, HOST};
|
||||
use hyper::{Request, StatusCode};
|
||||
use hyper_util::rt::TokioIo;
|
||||
use tokio::net::TcpStream;
|
||||
use tokio::task::JoinHandle;
|
||||
use tokio::time::{Instant, timeout, timeout_at};
|
||||
use tokio_tungstenite::tungstenite::Message;
|
||||
use tokio_tungstenite::{WebSocketStream, client_async};
|
||||
|
||||
use crate::config::{Protocol, Target};
|
||||
|
||||
const INTERACTION_INTERVAL: Duration = Duration::from_millis(100);
|
||||
const MAX_SSE_BUFFER: usize = 8_192;
|
||||
const SSE_EVENT: &[u8] = b"event: live\nid: 1\ndata: ready\n\n";
|
||||
const SSE_HEARTBEAT: &[u8] = b":\n\n";
|
||||
|
||||
pub struct Session {
|
||||
connection: Connection,
|
||||
cycle: usize,
|
||||
index: usize,
|
||||
timeout: Duration,
|
||||
messages: u64,
|
||||
}
|
||||
|
||||
enum Connection {
|
||||
Tcp(TcpStream),
|
||||
WebSocket(Box<WebSocketStream<TcpStream>>),
|
||||
Sse(SseConnection),
|
||||
}
|
||||
|
||||
struct SseConnection {
|
||||
body: Incoming,
|
||||
driver: ConnectionDriver,
|
||||
pending: Vec<u8>,
|
||||
}
|
||||
|
||||
struct ConnectionDriver(Option<JoinHandle<Result<(), hyper::Error>>>);
|
||||
|
||||
impl ConnectionDriver {
|
||||
fn new(driver: JoinHandle<Result<(), hyper::Error>>) -> Self {
|
||||
Self(Some(driver))
|
||||
}
|
||||
|
||||
async fn close(&mut self) {
|
||||
if let Some(driver) = self.0.take() {
|
||||
driver.abort();
|
||||
let _ = driver.await;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for ConnectionDriver {
|
||||
fn drop(&mut self) {
|
||||
if let Some(driver) = &self.0 {
|
||||
driver.abort();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Session {
|
||||
pub async fn connect(
|
||||
protocol: Protocol,
|
||||
target: &Target,
|
||||
cycle: usize,
|
||||
index: usize,
|
||||
operation_timeout: Duration,
|
||||
) -> Result<Self, String> {
|
||||
let stream = timeout(
|
||||
operation_timeout,
|
||||
TcpStream::connect(("127.0.0.1", target.port)),
|
||||
)
|
||||
.await
|
||||
.map_err(|_| "connection timed out".to_owned())?
|
||||
.map_err(|error| error.to_string())?;
|
||||
stream
|
||||
.set_nodelay(true)
|
||||
.map_err(|error| error.to_string())?;
|
||||
let connection = match protocol {
|
||||
Protocol::Tcp => Connection::Tcp(stream),
|
||||
Protocol::WebSocket => {
|
||||
let (socket, response) =
|
||||
timeout(operation_timeout, client_async(target.url.as_str(), stream))
|
||||
.await
|
||||
.map_err(|_| "WebSocket handshake timed out".to_owned())?
|
||||
.map_err(|error| error.to_string())?;
|
||||
if response.status() != StatusCode::SWITCHING_PROTOCOLS {
|
||||
return Err(format!(
|
||||
"WebSocket handshake returned {}",
|
||||
response.status()
|
||||
));
|
||||
}
|
||||
Connection::WebSocket(Box::new(socket))
|
||||
}
|
||||
Protocol::Sse => {
|
||||
Connection::Sse(SseConnection::open(stream, target, operation_timeout).await?)
|
||||
}
|
||||
};
|
||||
let mut session = Self {
|
||||
connection,
|
||||
cycle,
|
||||
index,
|
||||
timeout: operation_timeout,
|
||||
messages: 0,
|
||||
};
|
||||
session.exchange(0).await?;
|
||||
Ok(session)
|
||||
}
|
||||
|
||||
pub async fn hold(&mut self, duration: Duration) -> Result<(), String> {
|
||||
let deadline = Instant::now() + duration;
|
||||
let mut sequence = 1_u64;
|
||||
while Instant::now() < deadline {
|
||||
if let Connection::Sse(connection) = &mut self.connection {
|
||||
let heartbeats = connection.read_heartbeats(deadline).await?;
|
||||
if heartbeats == 0 {
|
||||
return Err("SSE hold ended without a heartbeat".into());
|
||||
}
|
||||
self.messages += heartbeats;
|
||||
break;
|
||||
}
|
||||
self.exchange(sequence).await?;
|
||||
sequence += 1;
|
||||
let next = (Instant::now() + INTERACTION_INTERVAL).min(deadline);
|
||||
tokio::time::sleep_until(next).await;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn close(mut self) -> Result<u64, String> {
|
||||
match &mut self.connection {
|
||||
Connection::Tcp(_) => {}
|
||||
Connection::WebSocket(socket) => timeout(self.timeout, socket.as_mut().close(None))
|
||||
.await
|
||||
.map_err(|_| "WebSocket close timed out".to_owned())?
|
||||
.map_err(|error| error.to_string())?,
|
||||
Connection::Sse(connection) => connection.close().await,
|
||||
}
|
||||
Ok(self.messages)
|
||||
}
|
||||
|
||||
async fn exchange(&mut self, sequence: u64) -> Result<(), String> {
|
||||
match &mut self.connection {
|
||||
Connection::Tcp(stream) => {
|
||||
let payload = payload(self.cycle, self.index, sequence);
|
||||
timeout(self.timeout, write_all(stream, &payload))
|
||||
.await
|
||||
.map_err(|_| "TCP write timed out".to_owned())??;
|
||||
let mut received = vec![0; payload.len()];
|
||||
timeout(self.timeout, read_exact(stream, &mut received))
|
||||
.await
|
||||
.map_err(|_| "TCP read timed out".to_owned())??;
|
||||
if received != payload {
|
||||
return Err("TCP echo body mismatch".into());
|
||||
}
|
||||
}
|
||||
Connection::WebSocket(socket) => {
|
||||
let payload = String::from_utf8(payload(self.cycle, self.index, sequence))
|
||||
.expect("benchmark payload is ASCII");
|
||||
timeout(
|
||||
self.timeout,
|
||||
socket.send(Message::Text(payload.clone().into())),
|
||||
)
|
||||
.await
|
||||
.map_err(|_| "WebSocket write timed out".to_owned())?
|
||||
.map_err(|error| error.to_string())?;
|
||||
let message = timeout(self.timeout, socket.next())
|
||||
.await
|
||||
.map_err(|_| "WebSocket read timed out".to_owned())?
|
||||
.ok_or("WebSocket ended before echo")?
|
||||
.map_err(|error| error.to_string())?;
|
||||
if !matches!(message, Message::Text(value) if value.as_str() == payload) {
|
||||
return Err("WebSocket echo body mismatch".into());
|
||||
}
|
||||
}
|
||||
Connection::Sse(connection) => connection.read_event(self.timeout).await?,
|
||||
}
|
||||
self.messages += 1;
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
impl SseConnection {
|
||||
async fn open(
|
||||
stream: TcpStream,
|
||||
target: &Target,
|
||||
operation_timeout: Duration,
|
||||
) -> Result<Self, String> {
|
||||
let (mut sender, connection) =
|
||||
timeout(operation_timeout, http1::handshake(TokioIo::new(stream)))
|
||||
.await
|
||||
.map_err(|_| "SSE HTTP handshake timed out".to_owned())?
|
||||
.map_err(|error| error.to_string())?;
|
||||
let driver = ConnectionDriver::new(tokio::spawn(connection));
|
||||
let request = Request::builder()
|
||||
.method("GET")
|
||||
.uri(&target.path)
|
||||
.header(HOST, format!("127.0.0.1:{}", target.port))
|
||||
.body(Empty::<Bytes>::new())
|
||||
.map_err(|error| error.to_string())?;
|
||||
let response = timeout(operation_timeout, sender.send_request(request))
|
||||
.await
|
||||
.map_err(|_| "SSE request timed out".to_owned())?
|
||||
.map_err(|error| error.to_string())?;
|
||||
if response.status() != StatusCode::OK {
|
||||
return Err(format!("SSE request returned {}", response.status()));
|
||||
}
|
||||
let content_type = response
|
||||
.headers()
|
||||
.get(CONTENT_TYPE)
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.unwrap_or_default();
|
||||
if content_type.split(';').next().map(str::trim) != Some("text/event-stream") {
|
||||
return Err("SSE response has the wrong Content-Type".into());
|
||||
}
|
||||
Ok(Self {
|
||||
body: response.into_body(),
|
||||
driver,
|
||||
pending: Vec::new(),
|
||||
})
|
||||
}
|
||||
|
||||
async fn read_event(&mut self, operation_timeout: Duration) -> Result<(), String> {
|
||||
while self.pending.len() < SSE_EVENT.len() {
|
||||
let frame = timeout(operation_timeout, self.body.frame())
|
||||
.await
|
||||
.map_err(|_| "SSE event timed out".to_owned())?
|
||||
.ok_or("SSE body ended before the initial event")?
|
||||
.map_err(|error| error.to_string())?;
|
||||
if let Ok(bytes) = frame.into_data() {
|
||||
self.extend(bytes.as_ref())?;
|
||||
}
|
||||
if self.pending.len() < SSE_EVENT.len() && !SSE_EVENT.starts_with(&self.pending) {
|
||||
return Err("SSE initial event body mismatch".into());
|
||||
}
|
||||
}
|
||||
if !self.pending.starts_with(SSE_EVENT) {
|
||||
return Err("SSE initial event body mismatch".into());
|
||||
}
|
||||
self.pending.drain(..SSE_EVENT.len());
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn read_heartbeats(&mut self, deadline: Instant) -> Result<u64, String> {
|
||||
let mut messages = consume_heartbeats(&mut self.pending)?;
|
||||
loop {
|
||||
let frame = match timeout_at(deadline, self.body.frame()).await {
|
||||
Err(_) => break,
|
||||
Ok(None) => return Err("SSE body ended during hold".into()),
|
||||
Ok(Some(Err(error))) => return Err(error.to_string()),
|
||||
Ok(Some(Ok(frame))) => frame,
|
||||
};
|
||||
if let Ok(bytes) = frame.into_data() {
|
||||
self.extend(bytes.as_ref())?;
|
||||
messages += consume_heartbeats(&mut self.pending)?;
|
||||
}
|
||||
}
|
||||
Ok(messages)
|
||||
}
|
||||
|
||||
fn extend(&mut self, bytes: &[u8]) -> Result<(), String> {
|
||||
if self.pending.len().saturating_add(bytes.len()) > MAX_SSE_BUFFER {
|
||||
return Err("SSE validation buffer exceeded its limit".into());
|
||||
}
|
||||
self.pending.extend_from_slice(bytes);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn close(&mut self) {
|
||||
self.driver.close().await;
|
||||
}
|
||||
}
|
||||
|
||||
fn consume_heartbeats(buffer: &mut Vec<u8>) -> Result<u64, String> {
|
||||
let mut messages = 0;
|
||||
while buffer.len() >= SSE_HEARTBEAT.len() {
|
||||
if !buffer.starts_with(SSE_HEARTBEAT) {
|
||||
return Err("SSE hold data was not a heartbeat".into());
|
||||
}
|
||||
buffer.drain(..SSE_HEARTBEAT.len());
|
||||
messages += 1;
|
||||
}
|
||||
if !SSE_HEARTBEAT.starts_with(buffer) {
|
||||
return Err("SSE heartbeat body mismatch".into());
|
||||
}
|
||||
Ok(messages)
|
||||
}
|
||||
|
||||
fn payload(cycle: usize, index: usize, sequence: u64) -> Vec<u8> {
|
||||
format!("live:{cycle}:{index}:{sequence}\n").into_bytes()
|
||||
}
|
||||
|
||||
async fn write_all(stream: &TcpStream, bytes: &[u8]) -> Result<(), String> {
|
||||
let mut offset = 0;
|
||||
while offset < bytes.len() {
|
||||
stream.writable().await.map_err(|error| error.to_string())?;
|
||||
match stream.try_write(&bytes[offset..]) {
|
||||
Ok(0) => return Err("TCP write returned zero bytes".into()),
|
||||
Ok(length) => offset += length,
|
||||
Err(error) if error.kind() == std::io::ErrorKind::WouldBlock => {}
|
||||
Err(error) => return Err(error.to_string()),
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn read_exact(stream: &TcpStream, bytes: &mut [u8]) -> Result<(), String> {
|
||||
let mut offset = 0;
|
||||
while offset < bytes.len() {
|
||||
stream.readable().await.map_err(|error| error.to_string())?;
|
||||
match stream.try_read(&mut bytes[offset..]) {
|
||||
Ok(0) => return Err("TCP connection ended before the echo completed".into()),
|
||||
Ok(length) => offset += length,
|
||||
Err(error) if error.kind() == std::io::ErrorKind::WouldBlock => {}
|
||||
Err(error) => return Err(error.to_string()),
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
@@ -0,0 +1,284 @@
|
||||
use std::io::{self, Write};
|
||||
use std::time::Duration;
|
||||
|
||||
use tokio::runtime::Builder;
|
||||
use tokio::sync::{mpsc, watch};
|
||||
use tokio::task::JoinSet;
|
||||
use tokio::time::{Instant, timeout_at};
|
||||
|
||||
use crate::config::Arguments;
|
||||
use crate::protocol::Session;
|
||||
|
||||
const SETTLE_TIME: Duration = Duration::from_millis(300);
|
||||
|
||||
#[derive(Default)]
|
||||
struct Counters {
|
||||
attempted_connections: u64,
|
||||
connected: u64,
|
||||
disconnected: u64,
|
||||
messages: u64,
|
||||
errors: u64,
|
||||
}
|
||||
|
||||
pub fn run(arguments: Arguments) -> Result<(), String> {
|
||||
let workers = crate::settings::Settings::load()?.workers()?;
|
||||
let runtime = Builder::new_multi_thread()
|
||||
.worker_threads(workers)
|
||||
.max_blocking_threads(1)
|
||||
.enable_io()
|
||||
.enable_time()
|
||||
.build()
|
||||
.map_err(|error| format!("cannot start live benchmark runtime: {error}"))?;
|
||||
let started = std::time::Instant::now();
|
||||
let result = runtime.block_on(execute(&arguments));
|
||||
let (counters, failure) = match result {
|
||||
Ok(counters) => (counters, None),
|
||||
Err((counters, error)) => (counters, Some(error)),
|
||||
};
|
||||
println!(
|
||||
"{{\"protocol\":\"{}\",\"connections\":{},\"cycles\":{},\"attempted_connections\":{},\"connected\":{},\"disconnected\":{},\"messages\":{},\"errors\":{},\"elapsed_ms\":{}}}",
|
||||
arguments.protocol.name(),
|
||||
arguments.connections,
|
||||
arguments.cycles,
|
||||
counters.attempted_connections,
|
||||
counters.connected,
|
||||
counters.disconnected,
|
||||
counters.messages,
|
||||
counters.errors,
|
||||
started.elapsed().as_millis()
|
||||
);
|
||||
io::stdout().flush().map_err(|error| error.to_string())?;
|
||||
match failure {
|
||||
Some(error) => Err(error),
|
||||
None if counters.errors != 0 => Err("live benchmark recorded connection errors".into()),
|
||||
None => Ok(()),
|
||||
}
|
||||
}
|
||||
|
||||
async fn execute(arguments: &Arguments) -> Result<Counters, (Counters, String)> {
|
||||
let mut counters = Counters::default();
|
||||
for cycle in 0..arguments.cycles {
|
||||
if let Err(error) = execute_cycle(arguments, cycle, &mut counters).await {
|
||||
counters.errors += 1;
|
||||
return Err((counters, error));
|
||||
}
|
||||
}
|
||||
Ok(counters)
|
||||
}
|
||||
|
||||
async fn execute_cycle(
|
||||
arguments: &Arguments,
|
||||
cycle: usize,
|
||||
counters: &mut Counters,
|
||||
) -> Result<(), String> {
|
||||
phase(cycle, "connecting")?;
|
||||
counters.attempted_connections += arguments.connections as u64;
|
||||
|
||||
let (ready_sender, mut ready_receiver) = mpsc::channel(arguments.connections);
|
||||
let (held_sender, mut held_receiver) = mpsc::channel(arguments.connections);
|
||||
let (failure_sender, mut failure_receiver) = mpsc::channel(arguments.connections);
|
||||
let (hold_sender, hold_receiver) = watch::channel(false);
|
||||
let (close_sender, close_receiver) = watch::channel(false);
|
||||
let mut sessions = JoinSet::new();
|
||||
for index in 0..arguments.connections {
|
||||
sessions.spawn(report_session_failure(
|
||||
session_task(
|
||||
arguments.clone(),
|
||||
cycle,
|
||||
index,
|
||||
ready_sender.clone(),
|
||||
held_sender.clone(),
|
||||
hold_receiver.clone(),
|
||||
close_receiver.clone(),
|
||||
),
|
||||
failure_sender.clone(),
|
||||
));
|
||||
}
|
||||
drop(ready_sender);
|
||||
drop(held_sender);
|
||||
drop(failure_sender);
|
||||
|
||||
if let Err(error) = wait_for_phase(
|
||||
&mut ready_receiver,
|
||||
arguments.connections,
|
||||
Instant::now() + arguments.timeout,
|
||||
"connections",
|
||||
&mut failure_receiver,
|
||||
)
|
||||
.await
|
||||
{
|
||||
abort_sessions(&mut sessions).await;
|
||||
return Err(error);
|
||||
}
|
||||
counters.connected += arguments.connections as u64;
|
||||
if let Err(error) = phase(cycle, "connected") {
|
||||
abort_sessions(&mut sessions).await;
|
||||
return Err(error);
|
||||
}
|
||||
if hold_sender.send(true).is_err() {
|
||||
abort_sessions(&mut sessions).await;
|
||||
return Err("connection tasks ended before hold started".into());
|
||||
}
|
||||
|
||||
let hold_deadline = Instant::now() + arguments.hold + arguments.timeout;
|
||||
if let Err(error) = wait_for_phase(
|
||||
&mut held_receiver,
|
||||
arguments.connections,
|
||||
hold_deadline,
|
||||
"held connections",
|
||||
&mut failure_receiver,
|
||||
)
|
||||
.await
|
||||
{
|
||||
let _ = close_sender.send(true);
|
||||
abort_sessions(&mut sessions).await;
|
||||
return Err(error);
|
||||
}
|
||||
if let Err(error) = phase(cycle, "closing") {
|
||||
abort_sessions(&mut sessions).await;
|
||||
return Err(error);
|
||||
}
|
||||
if close_sender.send(true).is_err() {
|
||||
abort_sessions(&mut sessions).await;
|
||||
return Err("connection tasks ended before close started".into());
|
||||
}
|
||||
|
||||
let close_deadline = Instant::now() + arguments.timeout;
|
||||
while !sessions.is_empty() {
|
||||
let result = match timeout_at(close_deadline, sessions.join_next()).await {
|
||||
Err(_) => {
|
||||
abort_sessions(&mut sessions).await;
|
||||
return Err("connection shutdown timed out".into());
|
||||
}
|
||||
Ok(None) => {
|
||||
abort_sessions(&mut sessions).await;
|
||||
return Err("connection task set ended unexpectedly".into());
|
||||
}
|
||||
Ok(Some(Err(error))) => {
|
||||
abort_sessions(&mut sessions).await;
|
||||
return Err(format!("connection task failed: {error}"));
|
||||
}
|
||||
Ok(Some(Ok(result))) => result,
|
||||
};
|
||||
match result {
|
||||
Ok(messages) => {
|
||||
counters.disconnected += 1;
|
||||
counters.messages += messages;
|
||||
}
|
||||
Err(error) => {
|
||||
abort_sessions(&mut sessions).await;
|
||||
return Err(error);
|
||||
}
|
||||
}
|
||||
}
|
||||
phase(cycle, "disconnected")?;
|
||||
tokio::time::sleep(SETTLE_TIME).await;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn session_task(
|
||||
arguments: Arguments,
|
||||
cycle: usize,
|
||||
index: usize,
|
||||
ready: mpsc::Sender<()>,
|
||||
held: mpsc::Sender<()>,
|
||||
mut hold: watch::Receiver<bool>,
|
||||
mut close: watch::Receiver<bool>,
|
||||
) -> Result<u64, String> {
|
||||
let mut session = Session::connect(
|
||||
arguments.protocol,
|
||||
&arguments.target,
|
||||
cycle,
|
||||
index,
|
||||
arguments.timeout,
|
||||
)
|
||||
.await
|
||||
.map_err(|error| connection_error(cycle, index, error))?;
|
||||
ready.send(()).await.map_err(|_| "ready receiver closed")?;
|
||||
drop(ready);
|
||||
wait_for_signal(&mut hold, arguments.timeout, "hold").await?;
|
||||
session
|
||||
.hold(arguments.hold)
|
||||
.await
|
||||
.map_err(|error| connection_error(cycle, index, error))?;
|
||||
held.send(()).await.map_err(|_| "held receiver closed")?;
|
||||
drop(held);
|
||||
wait_for_signal(&mut close, arguments.timeout, "close").await?;
|
||||
session
|
||||
.close()
|
||||
.await
|
||||
.map_err(|error| connection_error(cycle, index, error))
|
||||
}
|
||||
|
||||
async fn report_session_failure(
|
||||
session: impl std::future::Future<Output = Result<u64, String>>,
|
||||
failure: mpsc::Sender<String>,
|
||||
) -> Result<u64, String> {
|
||||
let result = session.await;
|
||||
if let Err(error) = &result {
|
||||
let _ = failure.send(error.clone()).await;
|
||||
}
|
||||
result
|
||||
}
|
||||
|
||||
async fn wait_for_phase(
|
||||
receiver: &mut mpsc::Receiver<()>,
|
||||
expected: usize,
|
||||
deadline: Instant,
|
||||
name: &str,
|
||||
failures: &mut mpsc::Receiver<String>,
|
||||
) -> Result<(), String> {
|
||||
for completed in 0..expected {
|
||||
tokio::select! {
|
||||
biased;
|
||||
failure = failures.recv() => {
|
||||
return Err(failure.unwrap_or_else(|| format!(
|
||||
"only {completed} of {expected} {name} reached the phase"
|
||||
)));
|
||||
}
|
||||
signal = timeout_at(deadline, receiver.recv()) => match signal {
|
||||
Ok(Some(())) => {}
|
||||
Ok(None) => {
|
||||
return Err(format!(
|
||||
"only {completed} of {expected} {name} reached the phase"
|
||||
));
|
||||
}
|
||||
Err(_) => return Err(format!("timed out waiting for {name}")),
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn wait_for_signal(
|
||||
receiver: &mut watch::Receiver<bool>,
|
||||
operation_timeout: Duration,
|
||||
name: &str,
|
||||
) -> Result<(), String> {
|
||||
if *receiver.borrow() {
|
||||
return Ok(());
|
||||
}
|
||||
tokio::time::timeout(operation_timeout, receiver.changed())
|
||||
.await
|
||||
.map_err(|_| format!("timed out waiting for {name} signal"))?
|
||||
.map_err(|_| format!("{name} signal closed"))?;
|
||||
if !*receiver.borrow() {
|
||||
return Err(format!("{name} signal did not start the phase"));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn abort_sessions(sessions: &mut JoinSet<Result<u64, String>>) {
|
||||
sessions.abort_all();
|
||||
while sessions.join_next().await.is_some() {}
|
||||
}
|
||||
|
||||
fn phase(cycle: usize, name: &str) -> Result<(), String> {
|
||||
let mut output = io::stdout().lock();
|
||||
writeln!(output, "PHASE|{cycle}|{name}").map_err(|error| error.to_string())?;
|
||||
output.flush().map_err(|error| error.to_string())
|
||||
}
|
||||
|
||||
fn connection_error(cycle: usize, index: usize, error: String) -> String {
|
||||
format!("cycle {cycle} connection {index}: {error}")
|
||||
}
|
||||
Reference in new issue
Block a user