pezkuwichain/pezkuwi-telemetry
1
0
Fork 0
mirror of https://github.com/pezkuwichain/pezkuwi-telemetry.git synced 2026-06-12 08:51:11 +00:00
Code Issues Packages Projects Releases Wiki Activity

cargo fmt

Browse Source
This commit is contained in:
James Wilson
2021-07-13 16:51:24 +01:00
parent 9ac5ea7624
commit c1208b9e81
13 changed files with 466 additions and 306 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/test_utils/src/feed_message_de.rs
+104 -47
View File
@@ -1,4 +1,6 @@
use common::node_types::{BlockDetails, BlockHash, BlockNumber, NodeLocation, NodeStats, Timestamp};
use common::node_types::{
BlockDetails, BlockHash, BlockNumber, NodeLocation, NodeStats, Timestamp,
};
use serde_json::value::RawValue;
#[derive(Debug, PartialEq)]
@@ -77,7 +79,7 @@ pub enum FeedMessage {
address: String,
block_number: BlockNumber,
block_hash: BlockHash,
voter: Option<String>
voter: Option<String>,
},
AfgReceivedPrecommit {
address: String,
@@ -85,7 +87,8 @@ pub enum FeedMessage {
block_hash: BlockHash,
voter: Option<String>,
},
AfgAuthoritySet { // Not used currently; not sure what "address" params are:
AfgAuthoritySet {
// Not used currently; not sure what "address" params are:
a1: String,
a2: String,
a3: String,
@@ -102,8 +105,8 @@ pub enum FeedMessage {
/// A "special" case when we don't know how to decode an action:
UnknownValue {
action: u8,
value: String
}
value: String,
},
}
#[derive(Debug, PartialEq)]
@@ -138,22 +141,30 @@ impl FeedMessage {
0 => {
let version = serde_json::from_str(raw_val.get())?;
FeedMessage::Version(version)
},
}
// BestBlock
1 => {
let (block_number, timestamp, avg_block_time) = serde_json::from_str(raw_val.get())?;
FeedMessage::BestBlock { block_number, timestamp, avg_block_time }
},
let (block_number, timestamp, avg_block_time) =
serde_json::from_str(raw_val.get())?;
FeedMessage::BestBlock {
block_number,
timestamp,
avg_block_time,
}
}
// BestFinalized
2 => {
let (block_number, block_hash) = serde_json::from_str(raw_val.get())?;
FeedMessage::BestFinalized { block_number, block_hash }
FeedMessage::BestFinalized {
block_number,
block_hash,
}
}
// AddNode
3 => {
let (
node_id,
( name, implementation, version, validator, network_id ),
(name, implementation, version, validator, network_id),
stats,
io,
hardware,
@@ -163,113 +174,153 @@ impl FeedMessage {
) = serde_json::from_str(raw_val.get())?;
// Give these two types but don't use the results:
let (_,_): (&RawValue, &RawValue) = (io, hardware);
let (_, _): (&RawValue, &RawValue) = (io, hardware);
FeedMessage::AddedNode {
node_id,
node: NodeDetails { name, implementation, version, validator, network_id },
node: NodeDetails {
name,
implementation,
version,
validator,
network_id,
},
stats,
block_details,
location,
startup_time,
}
},
}
// RemoveNode
4 => {
let node_id = serde_json::from_str(raw_val.get())?;
FeedMessage::RemovedNode { node_id }
},
}
// LocatedNode
5 => {
let (node_id, lat, long, city) = serde_json::from_str(raw_val.get())?;
FeedMessage::LocatedNode { node_id, lat, long, city }
},
FeedMessage::LocatedNode {
node_id,
lat,
long,
city,
}
}
// ImportedBlock
6 => {
let (node_id, block_details) = serde_json::from_str(raw_val.get())?;
FeedMessage::ImportedBlock { node_id, block_details }
},
FeedMessage::ImportedBlock {
node_id,
block_details,
}
}
// FinalizedBlock
7 => {
let (node_id, block_number, block_hash) = serde_json::from_str(raw_val.get())?;
FeedMessage::FinalizedBlock { node_id, block_number, block_hash }
},
FeedMessage::FinalizedBlock {
node_id,
block_number,
block_hash,
}
}
// NodeStatsUpdate
8 => {
let (node_id, stats) = serde_json::from_str(raw_val.get())?;
FeedMessage::NodeStatsUpdate { node_id, stats }
},
}
// Hardware
9 => {
let (node_id, _hardware): (_, &RawValue) = serde_json::from_str(raw_val.get())?;
FeedMessage::Hardware { node_id }
},
}
// TimeSync
10 => {
let time = serde_json::from_str(raw_val.get())?;
FeedMessage::TimeSync { time }
},
}
// AddedChain
11 => {
let (name, node_count) = serde_json::from_str(raw_val.get())?;
FeedMessage::AddedChain { name, node_count }
},
}
// RemovedChain
12 => {
let name = serde_json::from_str(raw_val.get())?;
FeedMessage::RemovedChain { name }
},
}
// SubscribedTo
13 => {
let name = serde_json::from_str(raw_val.get())?;
FeedMessage::SubscribedTo { name }
},
}
// UnsubscribedFrom
14 => {
let name = serde_json::from_str(raw_val.get())?;
FeedMessage::UnsubscribedFrom { name }
},
}
// Pong
15 => {
let msg = serde_json::from_str(raw_val.get())?;
FeedMessage::Pong { msg }
},
}
// AfgFinalized
16 => {
let (address, block_number, block_hash) = serde_json::from_str(raw_val.get())?;
FeedMessage::AfgFinalized { address, block_number, block_hash }
},
FeedMessage::AfgFinalized {
address,
block_number,
block_hash,
}
}
// AfgReceivedPrevote
17 => {
let (address, block_number, block_hash, voter) = serde_json::from_str(raw_val.get())?;
FeedMessage::AfgReceivedPrevote { address, block_number, block_hash, voter }
},
let (address, block_number, block_hash, voter) =
serde_json::from_str(raw_val.get())?;
FeedMessage::AfgReceivedPrevote {
address,
block_number,
block_hash,
voter,
}
}
// AfgReceivedPrecommit
18 => {
let (address, block_number, block_hash, voter) = serde_json::from_str(raw_val.get())?;
FeedMessage::AfgReceivedPrecommit { address, block_number, block_hash, voter }
},
let (address, block_number, block_hash, voter) =
serde_json::from_str(raw_val.get())?;
FeedMessage::AfgReceivedPrecommit {
address,
block_number,
block_hash,
voter,
}
}
// AfgAuthoritySet
19 => {
let (a1, a2, a3, block_number, block_hash) = serde_json::from_str(raw_val.get())?;
FeedMessage::AfgAuthoritySet { a1, a2, a3, block_number, block_hash }
},
FeedMessage::AfgAuthoritySet {
a1,
a2,
a3,
block_number,
block_hash,
}
}
// StaleNode
20 => {
let node_id = serde_json::from_str(raw_val.get())?;
FeedMessage::StaleNode { node_id }
},
}
// NodeIOUpdate
21 => {
// ignore NodeIO for now:
let (node_id, _node_io): (_, &RawValue) = serde_json::from_str(raw_val.get())?;
FeedMessage::NodeIOUpdate { node_id }
},
}
// A catchall for messages we don't know/care about yet:
_ => {
let value = raw_val.to_string();
FeedMessage::UnknownValue { action, value }
},
}
};
Ok(feed_message)
@@ -288,7 +339,9 @@ mod test {
assert_eq!(
FeedMessage::from_bytes(msg.as_bytes()).unwrap(),
vec![FeedMessage::RemovedChain { name: "".to_owned() }]
vec![FeedMessage::RemovedChain {
name: "".to_owned()
}]
);
}
@@ -300,10 +353,14 @@ mod test {
assert_eq!(
FeedMessage::from_bytes(msg.as_bytes()).unwrap(),
vec![
FeedMessage::RemovedChain { name: "".to_owned() },
FeedMessage::AddedChain { name: "Local Testnet".to_owned(), node_count: 1 },
FeedMessage::RemovedChain {
name: "".to_owned()
},
FeedMessage::AddedChain {
name: "Local Testnet".to_owned(),
node_count: 1
},
]
);
}
}
}
backend/test_utils/src/server/cargo_run_commands.rs
+1 -1
View File
@@ -33,4 +33,4 @@ fn try_find_workspace_dir() -> Result<PathBuf, std::io::Error> {
let mut dir = std::env::current_dir()?;
while !dir.ends_with("backend") && dir.pop() {}
Ok(dir)
}
}
backend/test_utils/src/server/channels.rs
+83 -30
View File
@@ -1,45 +1,65 @@
use std::time::Duration;
use crate::feed_message_de::FeedMessage;
use crate::ws_client;
use futures::{Sink, SinkExt, Stream, StreamExt};
use crate::feed_message_de::FeedMessage;
/// Wrap a `ws_client::Sender` with convenient utility methods for shard connections
pub struct ShardSender(ws_client::Sender);
impl From<ws_client::Sender> for ShardSender {
fn from(c: ws_client::Sender) -> Self { ShardSender(c) }
fn from(c: ws_client::Sender) -> Self {
ShardSender(c)
}
}
impl ShardSender {
/// Close this connection
pub async fn close(&mut self) -> Result<(),ws_client::SendError> {
pub async fn close(&mut self) -> Result<(), ws_client::SendError> {
self.0.close().await
}
}
impl Sink<ws_client::Message> for ShardSender {
type Error = ws_client::SendError;
fn poll_ready(mut self: std::pin::Pin<&mut Self>, cx: &mut std::task::Context<'_>) -> std::task::Poll<Result<(), Self::Error>> {
fn poll_ready(
mut self: std::pin::Pin<&mut Self>,
cx: &mut std::task::Context<'_>,
) -> std::task::Poll<Result<(), Self::Error>> {
self.0.poll_ready_unpin(cx)
}
fn start_send(mut self: std::pin::Pin<&mut Self>, item: ws_client::Message) -> Result<(), Self::Error> {
fn start_send(
mut self: std::pin::Pin<&mut Self>,
item: ws_client::Message,
) -> Result<(), Self::Error> {
self.0.start_send_unpin(item)
}
fn poll_flush(mut self: std::pin::Pin<&mut Self>, cx: &mut std::task::Context<'_>) -> std::task::Poll<Result<(), Self::Error>> {
fn poll_flush(
mut self: std::pin::Pin<&mut Self>,
cx: &mut std::task::Context<'_>,
) -> std::task::Poll<Result<(), Self::Error>> {
self.0.poll_flush_unpin(cx)
}
fn poll_close(mut self: std::pin::Pin<&mut Self>, cx: &mut std::task::Context<'_>) -> std::task::Poll<Result<(), Self::Error>> {
fn poll_close(
mut self: std::pin::Pin<&mut Self>,
cx: &mut std::task::Context<'_>,
) -> std::task::Poll<Result<(), Self::Error>> {
self.0.poll_close_unpin(cx)
}
}
impl ShardSender {
pub async fn send_json_binary(&mut self, json: serde_json::Value) -> Result<(), ws_client::SendError> {
pub async fn send_json_binary(
&mut self,
json: serde_json::Value,
) -> Result<(), ws_client::SendError> {
let bytes = serde_json::to_vec(&json).expect("valid bytes");
self.send(ws_client::Message::Binary(bytes)).await
}
pub async fn send_json_text(&mut self, json: serde_json::Value) -> Result<(), ws_client::SendError> {
pub async fn send_json_text(
&mut self,
json: serde_json::Value,
) -> Result<(), ws_client::SendError> {
let s = serde_json::to_string(&json).expect("valid string");
self.send(ws_client::Message::Text(s)).await
}
@@ -49,43 +69,70 @@ impl ShardSender {
pub struct ShardReceiver(ws_client::Receiver);
impl From<ws_client::Receiver> for ShardReceiver {
fn from(c: ws_client::Receiver) -> Self { ShardReceiver(c) }
fn from(c: ws_client::Receiver) -> Self {
ShardReceiver(c)
}
}
impl Stream for ShardReceiver {
type Item = Result<ws_client::Message, ws_client::RecvError>;
fn poll_next(mut self: std::pin::Pin<&mut Self>, cx: &mut std::task::Context<'_>) -> std::task::Poll<Option<Self::Item>> {
fn poll_next(
mut self: std::pin::Pin<&mut Self>,
cx: &mut std::task::Context<'_>,
) -> std::task::Poll<Option<Self::Item>> {
self.0.poll_next_unpin(cx)
}
}
/// Wrap a `ws_client::Sender` with convenient utility methods for feed connections
pub struct FeedSender(ws_client::Sender);
impl From<ws_client::Sender> for FeedSender {
fn from(c: ws_client::Sender) -> Self { FeedSender(c) }
fn from(c: ws_client::Sender) -> Self {
FeedSender(c)
}
}
impl Sink<ws_client::Message> for FeedSender {
type Error = ws_client::SendError;
fn poll_ready(mut self: std::pin::Pin<&mut Self>, cx: &mut std::task::Context<'_>) -> std::task::Poll<Result<(), Self::Error>> {
fn poll_ready(
mut self: std::pin::Pin<&mut Self>,
cx: &mut std::task::Context<'_>,
) -> std::task::Poll<Result<(), Self::Error>> {
self.0.poll_ready_unpin(cx)
}
fn start_send(mut self: std::pin::Pin<&mut Self>, item: ws_client::Message) -> Result<(), Self::Error> {
fn start_send(
mut self: std::pin::Pin<&mut Self>,
item: ws_client::Message,
) -> Result<(), Self::Error> {
self.0.start_send_unpin(item)
}
fn poll_flush(mut self: std::pin::Pin<&mut Self>, cx: &mut std::task::Context<'_>) -> std::task::Poll<Result<(), Self::Error>> {
fn poll_flush(
mut self: std::pin::Pin<&mut Self>,
cx: &mut std::task::Context<'_>,
) -> std::task::Poll<Result<(), Self::Error>> {
self.0.poll_flush_unpin(cx)
}
fn poll_close(mut self: std::pin::Pin<&mut Self>, cx: &mut std::task::Context<'_>) -> std::task::Poll<Result<(), Self::Error>> {
fn poll_close(
mut self: std::pin::Pin<&mut Self>,
cx: &mut std::task::Context<'_>,
) -> std::task::Poll<Result<(), Self::Error>> {
self.0.poll_close_unpin(cx)
}
}
impl FeedSender {
pub async fn send_command<S: AsRef<str>>(&mut self, command: S, param: S) -> Result<(), ws_client::SendError> {
self.send(ws_client::Message::Text(format!("{}:{}", command.as_ref(), param.as_ref()))).await
pub async fn send_command<S: AsRef<str>>(
&mut self,
command: S,
param: S,
) -> Result<(), ws_client::SendError> {
self.send(ws_client::Message::Text(format!(
"{}:{}",
command.as_ref(),
param.as_ref()
)))
.await
}
}
@@ -93,12 +140,17 @@ impl FeedSender {
pub struct FeedReceiver(ws_client::Receiver);
impl From<ws_client::Receiver> for FeedReceiver {
fn from(c: ws_client::Receiver) -> Self { FeedReceiver(c) }
fn from(c: ws_client::Receiver) -> Self {
FeedReceiver(c)
}
}
impl Stream for FeedReceiver {
type Item = Result<ws_client::Message, ws_client::RecvError>;
fn poll_next(mut self: std::pin::Pin<&mut Self>, cx: &mut std::task::Context<'_>) -> std::task::Poll<Option<Self::Item>> {
fn poll_next(
mut self: std::pin::Pin<&mut Self>,
cx: &mut std::task::Context<'_>,
) -> std::task::Poll<Option<Self::Item>> {
self.0.poll_next_unpin(cx).map_err(|e| e.into())
}
}
@@ -111,7 +163,8 @@ impl FeedReceiver {
/// robust in assuming that messages may not all be delivered at once (unless we are
/// specifically testing which messages are buffered together).
pub async fn recv_feed_messages_once(&mut self) -> Result<Vec<FeedMessage>, anyhow::Error> {
let msg = self.0
let msg = self
.0
.next()
.await
.ok_or_else(|| anyhow::anyhow!("Stream closed: no more messages"))??;
@@ -120,7 +173,7 @@ impl FeedReceiver {
ws_client::Message::Binary(data) => {
let messages = FeedMessage::from_bytes(&data)?;
Ok(messages)
},
}
ws_client::Message::Text(text) => {
let messages = FeedMessage::from_bytes(text.as_bytes())?;
Ok(messages)
@@ -135,20 +188,20 @@ impl FeedReceiver {
let mut feed_messages = self.recv_feed_messages_once().await?;
// Then, loop a little to make sure we catch any additional messages that are sent soon after:
loop {
match tokio::time::timeout(Duration::from_millis(250), self.recv_feed_messages_once()).await {
match tokio::time::timeout(Duration::from_millis(250), self.recv_feed_messages_once())
.await
{
// Timeout elapsed; return the messages we have so far
Err(_) => {
break Ok(feed_messages);
},
}
// Append messages that come back to our vec
Ok(Ok(mut msgs)) => {
feed_messages.append(&mut msgs);
},
// Error came back receiving messages; return it
Ok(Err(e)) => {
break Err(e)
}
// Error came back receiving messages; return it
Ok(Err(e)) => break Err(e),
}
}
}
}
}
backend/test_utils/src/server/mod.rs
+2 -2
View File
@@ -1,6 +1,6 @@
mod utils;
mod server;
mod utils;
pub mod cargo_run_commands;
pub mod channels;
pub use server::*;
pub use server::*;
backend/test_utils/src/server/server.rs
+41 -38
View File
@@ -1,9 +1,9 @@
use super::{channels, utils};
use crate::ws_client;
use common::{id_type, DenseMap};
use std::ffi::OsString;
use std::marker::PhantomData;
use crate::ws_client;
use tokio::process::{ self, Command as TokioCommand };
use super::{ channels, utils };
use common::{ id_type, DenseMap };
use tokio::process::{self, Command as TokioCommand};
id_type! {
/// The ID of a running process. Cannot be constructed externally.
@@ -16,7 +16,7 @@ pub struct StartOpts {
pub shard_command: Command,
/// Command to run to start a telemetry core process.
/// The `--listen` and `--log` arguments will be appended within and shouldn't be provided.
pub core_command: Command
pub core_command: Command,
}
pub struct ConnectToExistingOpts {
@@ -38,7 +38,9 @@ pub enum Error {
ErrorObtainingPort(anyhow::Error),
#[error("Whoops; attempt to kill a process we didn't start (and so have no handle to)")]
CannotKillNoHandle,
#[error("Whoops; attempt to add a shard to a server we didn't start (and so have no handle to)")]
#[error(
"Whoops; attempt to add a shard to a server we didn't start (and so have no handle to)"
)]
CannotAddShardNoHandle,
}
@@ -68,20 +70,21 @@ impl Server {
}
pub fn iter_shards(&self) -> impl Iterator<Item = &ShardProcess> {
self.shards.iter().map(|(_,v)| v)
self.shards.iter().map(|(_, v)| v)
}
pub async fn kill_shard(&mut self, id: ProcessId) -> bool {
let shard = match self.shards.remove(id) {
Some(shard) => shard,
None => return false
None => return false,
};
// With this, killing will complete even if the promise returned is cancelled
// (it should regardless, but just to play it safe..)
let _ = tokio::spawn(async move {
let _ = shard.kill().await;
}).await;
})
.await;
true
}
@@ -91,14 +94,9 @@ impl Server {
// Spawn so we don't need to await cleanup if we don't care.
// Run all kill futs simultaneously.
let handle = tokio::spawn(async move {
let shard_kill_futs = self.shards
.into_iter()
.map(|(_,s)| s.kill());
let shard_kill_futs = self.shards.into_iter().map(|(_, s)| s.kill());
let _ = tokio::join!(
futures::future::join_all(shard_kill_futs),
self.core.kill()
);
let _ = tokio::join!(futures::future::join_all(shard_kill_futs), self.core.kill());
});
// You can wait for cleanup but aren't obliged to:
@@ -109,10 +107,11 @@ impl Server {
pub async fn add_shard(&mut self) -> Result<ProcessId, Error> {
let core_uri = match &self.core_shard_submit_uri {
Some(uri) => uri,
None => return Err(Error::CannotAddShardNoHandle)
None => return Err(Error::CannotAddShardNoHandle),
};
let mut shard_cmd: TokioCommand = self.shard_command
let mut shard_cmd: TokioCommand = self
.shard_command
.clone()
.ok_or_else(|| Error::CannotAddShardNoHandle)?
.into();
@@ -141,8 +140,9 @@ impl Server {
let _ = utils::wait_for_line_containing(
&mut child_stdout,
"Connected to telemetry core",
std::time::Duration::from_secs(5)
).await;
std::time::Duration::from_secs(5),
)
.await;
// Since we're piping stdout from the child process, we need somewhere for it to go
// else the process will get stuck when it tries to produce output:
@@ -156,7 +156,7 @@ impl Server {
id,
handle: Some(shard_process),
uri: shard_uri,
_channel_type: PhantomData
_channel_type: PhantomData,
});
Ok(pid)
@@ -164,7 +164,6 @@ impl Server {
/// Start a telemetry_core process. From here, we can add/remove shards as needed.
pub async fn start(opts: StartOpts) -> Result<Server, Error> {
let mut core_cmd: TokioCommand = opts.core_command.into();
let mut child = core_cmd
@@ -194,16 +193,18 @@ impl Server {
Ok(Server {
shard_command: Some(opts.shard_command),
core_shard_submit_uri: Some(format!("http://127.0.0.1:{}/shard_submit", core_port)
.parse()
.expect("valid shard_submit URI")),
core_shard_submit_uri: Some(
format!("http://127.0.0.1:{}/shard_submit", core_port)
.parse()
.expect("valid shard_submit URI"),
),
shards: DenseMap::new(),
core: Process {
id: ProcessId(0),
handle: Some(child),
uri: feed_uri,
_channel_type: PhantomData,
}
},
})
}
@@ -229,12 +230,11 @@ impl Server {
uri: opts.feed_uri,
handle: None,
_channel_type: PhantomData,
}
},
}
}
}
/// This represents a running process that we can connect to, which
/// may be either a `telemetry_shard` or `telemetry_core`.
pub struct Process<Channel> {
@@ -245,7 +245,7 @@ pub struct Process<Channel> {
/// The URI that we can use to connect to the process socket.
uri: http::Uri,
/// The kind of the process (lets us add methods specific to shard/core).
_channel_type: PhantomData<Channel>
_channel_type: PhantomData<Channel>,
}
/// A shard process with shard-specific methods.
@@ -254,7 +254,7 @@ pub type ShardProcess = Process<(channels::ShardSender, channels::ShardReceiver)
/// A core process with core-specific methods.
pub type CoreProcess = Process<(channels::FeedSender, channels::FeedReceiver)>;
impl <Channel> Process<Channel> {
impl<Channel> Process<Channel> {
/// Get the ID of this process
pub fn id(&self) -> ProcessId {
self.id
@@ -265,25 +265,28 @@ impl <Channel> Process<Channel> {
async fn kill(self) -> Result<(), Error> {
match self.handle {
Some(mut handle) => Ok(handle.kill().await?),
None => Err(Error::CannotKillNoHandle)
None => Err(Error::CannotKillNoHandle),
}
}
}
impl <Send: From<ws_client::Sender>, Recv: From<ws_client::Receiver>> Process<(Send, Recv)> {
impl<Send: From<ws_client::Sender>, Recv: From<ws_client::Receiver>> Process<(Send, Recv)> {
/// Establish a connection to the process
pub async fn connect(&self) -> Result<(Send, Recv), Error> {
ws_client::connect(&self.uri)
.await
.map(|(s,r)| (s.into(), r.into()))
.map(|(s, r)| (s.into(), r.into()))
.map_err(|e| e.into())
}
/// Establish multiple connections to the process
pub async fn connect_multiple(&self, num_connections: usize) -> Result<Vec<(Send, Recv)>, Error> {
pub async fn connect_multiple(
&self,
num_connections: usize,
) -> Result<Vec<(Send, Recv)>, Error> {
utils::connect_multiple_to_uri(&self.uri, num_connections)
.await
.map(|v| v.into_iter().map(|(s,r)| (s.into(), r.into())).collect())
.map(|v| v.into_iter().map(|(s, r)| (s.into(), r.into())).collect())
.map_err(|e| e.into())
}
}
@@ -294,14 +297,14 @@ impl <Send: From<ws_client::Sender>, Recv: From<ws_client::Receiver>> Process<(S
#[derive(Clone, Debug)]
pub struct Command {
command: OsString,
args: Vec<OsString>
args: Vec<OsString>,
}
impl Command {
pub fn new<S: Into<OsString>>(command: S) -> Command {
Command {
command: command.into(),
args: Vec::new()
args: Vec::new(),
}
}
@@ -317,4 +320,4 @@ impl Into<TokioCommand> for Command {
cmd.args(self.args);
cmd
}
}
}
backend/test_utils/src/server/utils.rs
+28 -16
View File
@@ -1,8 +1,8 @@
use crate::ws_client;
use anyhow::{anyhow, Context};
use tokio::io::BufReader;
use tokio::io::{ AsyncRead, AsyncWrite, AsyncBufReadExt };
use tokio::io::{AsyncBufReadExt, AsyncRead, AsyncWrite};
use tokio::time::Duration;
use anyhow::{ anyhow, Context };
/// Reads from the stdout of the shard/core process to extract the port that was assigned to it,
/// with the side benefit that we'll wait for it to start listening before returning. We do this
@@ -23,25 +23,35 @@ pub async fn get_port<R: AsyncRead + Unpin>(reader: R) -> Result<u16, anyhow::Er
/// Wait for a line of output containing the text given. Also provide a timeout,
/// such that if we don't see a new line of output within the timeout we bail out
/// and return an error.
pub async fn wait_for_line_containing<R: AsyncRead + Unpin>(reader: R, text: &str, max_wait_between_lines: Duration) -> Result<String, anyhow::Error> {
pub async fn wait_for_line_containing<R: AsyncRead + Unpin>(
reader: R,
text: &str,
max_wait_between_lines: Duration,
) -> Result<String, anyhow::Error> {
let reader = BufReader::new(reader);
let mut reader_lines = reader.lines();
loop {
let line = tokio::time::timeout(
max_wait_between_lines,
reader_lines.next_line()
).await;
let line = tokio::time::timeout(max_wait_between_lines, reader_lines.next_line()).await;
let line = match line {
// timeout expired; couldn't get port:
Err(_) => return Err(anyhow!("Timeout expired waiting for output containing: {}", text)),
Err(_) => {
return Err(anyhow!(
"Timeout expired waiting for output containing: {}",
text
))
}
// Something went wrong reading line; bail:
Ok(Err(e)) => return Err(anyhow!("Could not read line from stdout: {}", e)),
// No more output; process ended? bail:
Ok(Ok(None)) => return Err(anyhow!("No more output from stdout; has the process ended?")),
Ok(Ok(None)) => {
return Err(anyhow!(
"No more output from stdout; has the process ended?"
))
}
// All OK, and a line is given back; phew!
Ok(Ok(Some(line))) => line
Ok(Ok(Some(line))) => line,
};
if line.contains(text) {
@@ -51,10 +61,12 @@ pub async fn wait_for_line_containing<R: AsyncRead + Unpin>(reader: R, text: &st
}
/// Establish multiple connections to a URI and return them all.
pub async fn connect_multiple_to_uri(uri: &http::Uri, num_connections: usize) -> Result<Vec<(ws_client::Sender, ws_client::Receiver)>, ws_client::ConnectError> {
let connect_futs = (0..num_connections)
.map(|_| ws_client::connect(uri));
let sockets: Result<Vec<_>,_> = futures::future::join_all(connect_futs)
pub async fn connect_multiple_to_uri(
uri: &http::Uri,
num_connections: usize,
) -> Result<Vec<(ws_client::Sender, ws_client::Receiver)>, ws_client::ConnectError> {
let connect_futs = (0..num_connections).map(|_| ws_client::connect(uri));
let sockets: Result<Vec<_>, _> = futures::future::join_all(connect_futs)
.await
.into_iter()
.collect();
@@ -66,9 +78,9 @@ pub async fn connect_multiple_to_uri(uri: &http::Uri, num_connections: usize) ->
pub fn drain<R, W>(mut reader: R, mut writer: W)
where
R: AsyncRead + Unpin + Send + 'static,
W: AsyncWrite + Unpin + Send + 'static
W: AsyncWrite + Unpin + Send + 'static,
{
tokio::spawn(async move {
let _ = tokio::io::copy(&mut reader, &mut writer).await;
});
}
}
backend/test_utils/src/ws_client.rs
+60 -42
View File
@@ -1,17 +1,17 @@
use futures::channel::{ mpsc };
use soketto::handshake::{Client, ServerResponse};
use tokio_util::compat::{ TokioAsyncReadCompatExt };
use tokio::net::TcpStream;
use futures::channel::mpsc;
use futures::{Sink, SinkExt, Stream, StreamExt};
use soketto::handshake::{Client, ServerResponse};
use tokio::net::TcpStream;
use tokio_util::compat::TokioAsyncReadCompatExt;
/// Send messages into the connection
#[derive(Clone)]
pub struct Sender {
inner: mpsc::UnboundedSender<SentMessage>
inner: mpsc::UnboundedSender<SentMessage>,
}
impl Sender {
pub async fn close(&mut self) -> Result<(),SendError> {
pub async fn close(&mut self) -> Result<(), SendError> {
self.inner.send(SentMessage::Close).await?;
Ok(())
}
@@ -20,28 +20,39 @@ impl Sender {
#[derive(thiserror::Error, Debug, Clone)]
pub enum SendError {
#[error("Failed to send message: {0}")]
ChannelError(#[from] mpsc::SendError)
ChannelError(#[from] mpsc::SendError),
}
impl Sink<Message> for Sender {
type Error = SendError;
fn poll_ready(mut self: std::pin::Pin<&mut Self>, cx: &mut std::task::Context<'_>) -> std::task::Poll<Result<(), Self::Error>> {
fn poll_ready(
mut self: std::pin::Pin<&mut Self>,
cx: &mut std::task::Context<'_>,
) -> std::task::Poll<Result<(), Self::Error>> {
self.inner.poll_ready_unpin(cx).map_err(|e| e.into())
}
fn start_send(mut self: std::pin::Pin<&mut Self>, item: Message) -> Result<(), Self::Error> {
self.inner.start_send_unpin(SentMessage::Message(item)).map_err(|e| e.into())
self.inner
.start_send_unpin(SentMessage::Message(item))
.map_err(|e| e.into())
}
fn poll_flush(mut self: std::pin::Pin<&mut Self>, cx: &mut std::task::Context<'_>) -> std::task::Poll<Result<(), Self::Error>> {
fn poll_flush(
mut self: std::pin::Pin<&mut Self>,
cx: &mut std::task::Context<'_>,
) -> std::task::Poll<Result<(), Self::Error>> {
self.inner.poll_flush_unpin(cx).map_err(|e| e.into())
}
fn poll_close(mut self: std::pin::Pin<&mut Self>, cx: &mut std::task::Context<'_>) -> std::task::Poll<Result<(), Self::Error>> {
fn poll_close(
mut self: std::pin::Pin<&mut Self>,
cx: &mut std::task::Context<'_>,
) -> std::task::Poll<Result<(), Self::Error>> {
self.inner.poll_close_unpin(cx).map_err(|e| e.into())
}
}
/// Receive messages out of a connection
pub struct Receiver {
inner: mpsc::UnboundedReceiver<Result<Message,RecvError>>
inner: mpsc::UnboundedReceiver<Result<Message, RecvError>>,
}
#[derive(thiserror::Error, Debug)]
@@ -49,12 +60,15 @@ pub enum RecvError {
#[error("Text message contains invalid UTF8: {0}")]
InvalidUtf8(#[from] std::string::FromUtf8Error),
#[error("Stream finished")]
StreamFinished
StreamFinished,
}
impl Stream for Receiver {
type Item = Result<Message, RecvError>;
fn poll_next(mut self: std::pin::Pin<&mut Self>, cx: &mut std::task::Context<'_>) -> std::task::Poll<Option<Self::Item>> {
fn poll_next(
mut self: std::pin::Pin<&mut Self>,
cx: &mut std::task::Context<'_>,
) -> std::task::Poll<Option<Self::Item>> {
self.inner.poll_next_unpin(cx).map_err(|e| e.into())
}
}
@@ -62,13 +76,13 @@ impl Stream for Receiver {
/// A message type that can be sent or received from the connection
pub enum Message {
Text(String),
Binary(Vec<u8>)
Binary(Vec<u8>),
}
/// Sent messages can be anything publically visible, or a close message.
enum SentMessage {
Message(Message),
Close
Close,
}
#[derive(thiserror::Error, Debug)]
@@ -98,12 +112,10 @@ pub async fn connect(uri: &http::Uri) -> Result<(Sender, Receiver), ConnectError
// Establish a WS connection:
let mut client = Client::new(socket.compat(), host, &path);
let (mut ws_to_connection, mut ws_from_connection) = match client.handshake().await? {
ServerResponse::Accepted { .. } => {
client.into_builder().finish()
},
ServerResponse::Accepted { .. } => client.into_builder().finish(),
ServerResponse::Redirect { status_code, .. } => {
return Err(ConnectError::ConnectionFailedRedirect { status_code })
},
}
ServerResponse::Rejected { status_code } => {
return Err(ConnectError::ConnectionFailedRejected { status_code })
}
@@ -124,23 +136,20 @@ pub async fn connect(uri: &http::Uri) -> Result<(Sender, Receiver), ConnectError
Err(e) => {
// Couldn't receive data may mean all senders are gone, so log
// the error and shut this down:
log::error!("Shutting down websocket connection: Failed to receive data: {}", e);
log::error!(
"Shutting down websocket connection: Failed to receive data: {}",
e
);
break;
},
Ok(data) => {
data
}
Ok(data) => data,
};
let msg = match message_data {
soketto::Data::Text(_) => {
Ok(Message::Binary(data))
},
soketto::Data::Binary(_) => {
String::from_utf8(data)
.map(|s| Message::Text(s))
.map_err(|e| e.into())
},
soketto::Data::Text(_) => Ok(Message::Binary(data)),
soketto::Data::Binary(_) => String::from_utf8(data)
.map(|s| Message::Text(s))
.map_err(|e| e.into()),
};
data = Vec::with_capacity(128);
@@ -148,7 +157,10 @@ pub async fn connect(uri: &http::Uri) -> Result<(Sender, Receiver), ConnectError
if let Err(e) = tx_to_external.send(msg).await {
// Failure to send likely means that the recv has been dropped,
// so let's drop this loop too.
log::error!("Shutting down websocket connection: Failed to send data out: {}", e);
log::error!(
"Shutting down websocket connection: Failed to send data out: {}",
e
);
break;
}
}
@@ -161,16 +173,22 @@ pub async fn connect(uri: &http::Uri) -> Result<(Sender, Receiver), ConnectError
match msg {
SentMessage::Message(Message::Text(s)) => {
if let Err(e) = ws_to_connection.send_text_owned(s).await {
log::error!("Shutting down websocket connection: Failed to send text data: {}", e);
log::error!(
"Shutting down websocket connection: Failed to send text data: {}",
e
);
break;
}
},
}
SentMessage::Message(Message::Binary(bytes)) => {
if let Err(e) = ws_to_connection.send_binary_mut(bytes).await {
log::error!("Shutting down websocket connection: Failed to send binary data: {}", e);
log::error!(
"Shutting down websocket connection: Failed to send binary data: {}",
e
);
break;
}
},
}
SentMessage::Close => {
if let Err(e) = ws_to_connection.close().await {
log::error!("Error attempting to close connection: {}", e);
@@ -180,14 +198,14 @@ pub async fn connect(uri: &http::Uri) -> Result<(Sender, Receiver), ConnectError
}
if let Err(e) = ws_to_connection.flush().await {
log::error!("Shutting down websocket connection: Failed to flush data: {}", e);
log::error!(
"Shutting down websocket connection: Failed to flush data: {}",
e
);
break;
}
}
});
Ok((
Sender { inner: tx_to_ws },
Receiver { inner: rx_from_ws }
))
Ok((Sender { inner: tx_to_ws }, Receiver { inner: rx_from_ws }))
}