pezkuwichain/pezkuwi-telemetry
1
0
Fork 0
mirror of https://github.com/pezkuwichain/pezkuwi-telemetry.git synced 2026-04-29 12:37:59 +00:00
Code Issues Packages Projects Releases Wiki Activity

cargo fmt

Browse Source
This commit is contained in:
James Wilson
2021-07-27 12:38:05 +01:00
parent a44f39cfaf
commit 2ca5f6a1fb
23 changed files with 707 additions and 511 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/test_utils/src/server/channels.rs
+25 -17
View File
@@ -1,4 +1,7 @@
use std::{ops::{Deref, DerefMut}, time::Duration};
use std::{
ops::{Deref, DerefMut},
time::Duration,
};
use crate::feed_message_de::FeedMessage;
use common::ws_client;
@@ -51,10 +54,7 @@ impl ShardSender {
self.unbounded_send(ws_client::SentMessage::Binary(bytes))
}
/// Send JSON as a textual websocket message
pub fn send_json_text(
&mut self,
json: serde_json::Value,
) -> Result<(), ws_client::SendError> {
pub 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.unbounded_send(ws_client::SentMessage::Text(s))
}
@@ -169,7 +169,6 @@ impl FeedSender {
}
}
/// Wrap a `ws_client::Receiver` with convenient utility methods for feed connections
pub struct FeedReceiver(ws_client::Receiver);
@@ -208,23 +207,26 @@ impl FeedReceiver {
/// Prefer [`FeedReceiver::recv_feed_messages`]; tests should generally be
/// 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_timeout(&mut self, timeout: Duration) -> Result<Vec<FeedMessage>, anyhow::Error> {
pub async fn recv_feed_messages_once_timeout(
&mut self,
timeout: Duration,
) -> Result<Vec<FeedMessage>, anyhow::Error> {
let msg = match tokio::time::timeout(timeout, self.0.next()).await {
// Timeout elapsed; no messages back:
Err(_) => return Ok(Vec::new()),
// Something back; Complain if error no stream closed:
Ok(res) => res.ok_or_else(|| anyhow::anyhow!("Stream closed: no more messages"))??
Ok(res) => res.ok_or_else(|| anyhow::anyhow!("Stream closed: no more messages"))??,
};
match msg {
ws_client::RecvMessage::Binary(data) => {
let messages = FeedMessage::from_bytes(&data)?;
Ok(messages)
},
}
ws_client::RecvMessage::Text(text) => {
let messages = FeedMessage::from_bytes(text.as_bytes())?;
Ok(messages)
},
}
}
}
@@ -232,7 +234,8 @@ impl FeedReceiver {
/// See `recv_feed_messages_once_timeout`.
pub async fn recv_feed_messages_once(&mut self) -> Result<Vec<FeedMessage>, anyhow::Error> {
// Default to a timeout of 30 seconds, meaning that the test will eventually end,
self.recv_feed_messages_once_timeout(Duration::from_secs(30)).await
self.recv_feed_messages_once_timeout(Duration::from_secs(30))
.await
}
/// Wait for feed messages to be sent back, building up a list of output messages until
@@ -241,12 +244,16 @@ impl FeedReceiver {
/// If no new messages are received within the timeout given, bail with whatever we have so far.
/// This differs from `recv_feed_messages` and `recv_feed_messages_once`, which will block indefinitely
/// waiting for something to arrive
pub async fn recv_feed_messages_timeout(&mut self, timeout: Duration) -> Result<Vec<FeedMessage>, anyhow::Error> {
pub async fn recv_feed_messages_timeout(
&mut self,
timeout: Duration,
) -> Result<Vec<FeedMessage>, anyhow::Error> {
// Block as long as needed for messages to start coming in:
let mut feed_messages = match tokio::time::timeout(timeout, self.recv_feed_messages_once()).await {
Ok(msgs) => msgs?,
Err(_) => return Ok(Vec::new()),
};
let mut feed_messages =
match tokio::time::timeout(timeout, self.recv_feed_messages_once()).await {
Ok(msgs) => msgs?,
Err(_) => return Ok(Vec::new()),
};
// Then, loop a little to make sure we catch any additional messages that are sent soon after:
loop {
@@ -271,6 +278,7 @@ impl FeedReceiver {
/// See `recv_feed_messages_timeout`.
pub async fn recv_feed_messages(&mut self) -> Result<Vec<FeedMessage>, anyhow::Error> {
// Default to a timeout of 30 seconds, meaning that the test will eventually end,
self.recv_feed_messages_timeout(Duration::from_secs(30)).await
self.recv_feed_messages_timeout(Duration::from_secs(30))
.await
}
}
backend/test_utils/src/server/server.rs
+76 -53
View File
@@ -36,7 +36,7 @@ pub enum StartOpts {
/// Where is the process that we can subscribe to the `/feed` of?
/// Eg: `127.0.0.1:3000`
feed_host: String,
}
},
}
/// This represents a telemetry server. It can be in different modes
@@ -47,7 +47,7 @@ pub enum Server {
/// A virtual shard that we can hand out.
virtual_shard: ShardProcess,
/// Core process that we can connect to.
core: CoreProcess
core: CoreProcess,
},
ShardAndCoreMode {
/// Command to run to start a new shard.
@@ -67,10 +67,9 @@ pub enum Server {
shards: DenseMap<ProcessId, ShardProcess>,
/// Core process that we can connect to.
core: CoreProcess,
}
},
}
#[derive(thiserror::Error, Debug)]
pub enum Error {
#[error("Can't establsih connection: {0}")]
@@ -88,23 +87,23 @@ pub enum Error {
)]
CannotAddShard,
#[error("The URI provided was invalid: {0}")]
InvalidUri(#[from] http::uri::InvalidUri)
InvalidUri(#[from] http::uri::InvalidUri),
}
impl Server {
pub fn get_core(&self) -> &CoreProcess {
match self {
Server::SingleProcessMode { core, .. } => core,
Server::ShardAndCoreMode { core, ..} => core,
Server::ConnectToExistingMode { core, .. } => core
Server::ShardAndCoreMode { core, .. } => core,
Server::ConnectToExistingMode { core, .. } => core,
}
}
pub fn get_shard(&self, id: ProcessId) -> Option<&ShardProcess> {
match self {
Server::SingleProcessMode { virtual_shard, .. } => Some(virtual_shard),
Server::ShardAndCoreMode { shards, ..} => shards.get(id),
Server::ConnectToExistingMode { shards, .. } => shards.get(id)
Server::ShardAndCoreMode { shards, .. } => shards.get(id),
Server::ConnectToExistingMode { shards, .. } => shards.get(id),
}
}
@@ -112,8 +111,8 @@ impl Server {
let shard = match self {
// Can't remove the pretend shard:
Server::SingleProcessMode { .. } => return false,
Server::ShardAndCoreMode { shards, ..} => shards.remove(id),
Server::ConnectToExistingMode { shards, .. } => shards.remove(id)
Server::ShardAndCoreMode { shards, .. } => shards.remove(id),
Server::ConnectToExistingMode { shards, .. } => shards.remove(id),
};
let shard = match shard {
@@ -137,12 +136,9 @@ impl Server {
// Run all kill futs simultaneously.
let handle = tokio::spawn(async move {
let (core, shards) = match self {
Server::SingleProcessMode { core, .. }
=> (core, DenseMap::new()),
Server::ShardAndCoreMode { core, shards, ..}
=> (core, shards),
Server::ConnectToExistingMode { core, shards, .. }
=> (core, shards)
Server::SingleProcessMode { core, .. } => (core, DenseMap::new()),
Server::ShardAndCoreMode { core, shards, .. } => (core, shards),
Server::ConnectToExistingMode { core, shards, .. } => (core, shards),
};
let shard_kill_futs = shards.into_iter().map(|(_, s)| s.kill());
@@ -157,15 +153,18 @@ impl Server {
pub async fn add_shard(&mut self) -> Result<ProcessId, Error> {
match self {
// Always get back the same "virtual" shard; we're always just talking to the core anyway.
Server::SingleProcessMode { virtual_shard, .. } => {
Ok(virtual_shard.id)
},
Server::SingleProcessMode { virtual_shard, .. } => Ok(virtual_shard.id),
// We're connecting to an existing process. Find the next host we've been told about
// round-robin style and use that as our new virtual shard.
Server::ConnectToExistingMode { submit_hosts, next_submit_host_idx, shards, .. } => {
Server::ConnectToExistingMode {
submit_hosts,
next_submit_host_idx,
shards,
..
} => {
let host = match submit_hosts.get(*next_submit_host_idx % submit_hosts.len()) {
Some(host) => host,
None => return Err(Error::CannotAddShard)
None => return Err(Error::CannotAddShard),
};
*next_submit_host_idx += 1;
@@ -177,9 +176,13 @@ impl Server {
});
Ok(pid)
},
}
// Start a new process and return that.
Server::ShardAndCoreMode { shard_command, shards, core } => {
Server::ShardAndCoreMode {
shard_command,
shards,
core,
} => {
// Where is the URI we'll want to submit things to?
let core_shard_submit_uri = format!("http://{}/shard_submit", core.host);
@@ -224,7 +227,7 @@ impl Server {
});
Ok(pid)
},
}
}
}
@@ -240,31 +243,35 @@ impl Server {
id: ProcessId(0),
host: virtual_shard_host,
handle: None,
_channel_type: PhantomData
}
}
},
StartOpts::ShardAndCore { core_command, shard_command } => {
let core_process = Server::start_core(core_command).await?;
Server::ShardAndCoreMode {
core: core_process,
shard_command,
shards: DenseMap::new()
}
},
StartOpts::ConnectToExisting { feed_host, submit_hosts } => {
Server::ConnectToExistingMode {
submit_hosts,
next_submit_host_idx: 0,
shards: DenseMap::new(),
core: Process {
id: ProcessId(0),
host: feed_host,
handle: None,
_channel_type: PhantomData,
},
}
}
StartOpts::ShardAndCore {
core_command,
shard_command,
} => {
let core_process = Server::start_core(core_command).await?;
Server::ShardAndCoreMode {
core: core_process,
shard_command,
shards: DenseMap::new(),
}
}
StartOpts::ConnectToExisting {
feed_host,
submit_hosts,
} => Server::ConnectToExistingMode {
submit_hosts,
next_submit_host_idx: 0,
shards: DenseMap::new(),
core: Process {
id: ProcessId(0),
host: feed_host,
handle: None,
_channel_type: PhantomData,
},
},
};
Ok(server)
@@ -340,7 +347,9 @@ impl<Channel> Process<Channel> {
}
/// Establish a raw WebSocket connection (not cancel-safe)
async fn connect_to_uri_raw(uri: &http::Uri) -> Result<(ws_client::RawSender, ws_client::RawReceiver), Error> {
async fn connect_to_uri_raw(
uri: &http::Uri,
) -> Result<(ws_client::RawSender, ws_client::RawReceiver), Error> {
ws_client::connect(uri)
.await
.map(|c| c.into_raw())
@@ -371,19 +380,26 @@ impl<Send: From<ws_client::Sender>, Recv: From<ws_client::Receiver>> Process<(Se
impl ShardProcess {
/// Establish a raw connection to the process
pub async fn connect_node_raw(&self) -> Result<(ws_client::RawSender, ws_client::RawReceiver), Error> {
pub async fn connect_node_raw(
&self,
) -> Result<(ws_client::RawSender, ws_client::RawReceiver), Error> {
let uri = format!("http://{}/submit", self.host).parse()?;
connect_to_uri_raw(&uri).await
}
/// Establish a connection to the process
pub async fn connect_node(&self) -> Result<(channels::ShardSender, channels::ShardReceiver), Error> {
pub async fn connect_node(
&self,
) -> Result<(channels::ShardSender, channels::ShardReceiver), Error> {
let uri = format!("http://{}/submit", self.host).parse()?;
Process::connect_to_uri(&uri).await
}
/// Establish multiple connections to the process
pub async fn connect_multiple_nodes(&self, num_connections: usize) -> Result<Vec<(channels::ShardSender, channels::ShardReceiver)>, Error> {
pub async fn connect_multiple_nodes(
&self,
num_connections: usize,
) -> Result<Vec<(channels::ShardSender, channels::ShardReceiver)>, Error> {
let uri = format!("http://{}/submit", self.host).parse()?;
Process::connect_multiple_to_uri(&uri, num_connections).await
}
@@ -391,19 +407,26 @@ impl ShardProcess {
impl CoreProcess {
/// Establish a raw connection to the process
pub async fn connect_feed_raw(&self) -> Result<(ws_client::RawSender, ws_client::RawReceiver), Error> {
pub async fn connect_feed_raw(
&self,
) -> Result<(ws_client::RawSender, ws_client::RawReceiver), Error> {
let uri = format!("http://{}/feed", self.host).parse()?;
connect_to_uri_raw(&uri).await
}
/// Establish a connection to the process
pub async fn connect_feed(&self) -> Result<(channels::FeedSender, channels::FeedReceiver), Error> {
pub async fn connect_feed(
&self,
) -> Result<(channels::FeedSender, channels::FeedReceiver), Error> {
let uri = format!("http://{}/feed", self.host).parse()?;
Process::connect_to_uri(&uri).await
}
/// Establish multiple connections to the process
pub async fn connect_multiple_feeds(&self, num_connections: usize) -> Result<Vec<(channels::FeedSender, channels::FeedReceiver)>, Error> {
pub async fn connect_multiple_feeds(
&self,
num_connections: usize,
) -> Result<Vec<(channels::FeedSender, channels::FeedReceiver)>, Error> {
let uri = format!("http://{}/feed", self.host).parse()?;
Process::connect_multiple_to_uri(&uri, num_connections).await
}
backend/test_utils/src/server/utils.rs
+3 -7
View File
@@ -1,5 +1,5 @@
use common::ws_client;
use anyhow::{anyhow, Context};
use common::ws_client;
use tokio::io::BufReader;
use tokio::io::{AsyncBufReadExt, AsyncRead, AsyncWrite};
use tokio::time::Duration;
@@ -45,13 +45,9 @@ pub async fn wait_for_line_containing<R: AsyncRead + Unpin, F: Fn(&str) -> bool>
let line = match line {
// timeout expired; couldn't get port:
Err(_) => {
return Err(anyhow!("Timeout elapsed waiting for text match"))
}
Err(_) => return Err(anyhow!("Timeout elapsed waiting for text match")),
// Something went wrong reading line; bail:
Ok(Err(e)) => {
return Err(anyhow!("Could not read line from stdout: {}", e))
},
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!(