cargo fmt

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
+13 -13
View File
@@ -22,7 +22,6 @@ pub struct NodeDetails {
pub startup_time: Option<Box<str>>, pub startup_time: Option<Box<str>>,
} }
/// ///
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] #[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub struct NodeStats { pub struct NodeStats {
@@ -42,7 +41,7 @@ impl Serialize for NodeStats {
} }
} }
impl <'de> Deserialize<'de> for NodeStats { impl<'de> Deserialize<'de> for NodeStats {
fn deserialize<D>(deserializer: D) -> Result<Self, D::Error> fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
where where
D: serde::Deserializer<'de>, D: serde::Deserializer<'de>,
@@ -52,7 +51,6 @@ impl <'de> Deserialize<'de> for NodeStats {
} }
} }
/// ///
#[derive(Default)] #[derive(Default)]
pub struct NodeIO { pub struct NodeIO {
@@ -71,7 +69,6 @@ impl Serialize for NodeIO {
} }
} }
/// ///
#[derive(Deserialize, Serialize, Debug, Clone, Copy, PartialEq)] #[derive(Deserialize, Serialize, Debug, Clone, Copy, PartialEq)]
pub struct Block { pub struct Block {
@@ -88,7 +85,6 @@ impl Block {
} }
} }
/// ///
#[derive(Default)] #[derive(Default)]
pub struct NodeHardware { pub struct NodeHardware {
@@ -114,7 +110,6 @@ impl Serialize for NodeHardware {
} }
} }
/// ///
#[derive(Debug, Clone, PartialEq)] #[derive(Debug, Clone, PartialEq)]
pub struct NodeLocation { pub struct NodeLocation {
@@ -136,17 +131,20 @@ impl Serialize for NodeLocation {
} }
} }
impl <'de> Deserialize<'de> for NodeLocation { impl<'de> Deserialize<'de> for NodeLocation {
fn deserialize<D>(deserializer: D) -> Result<Self, D::Error> fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
where where
D: serde::Deserializer<'de>, D: serde::Deserializer<'de>,
{ {
let (latitude, longitude, city) = <(f32, f32, Box<str>)>::deserialize(deserializer)?; let (latitude, longitude, city) = <(f32, f32, Box<str>)>::deserialize(deserializer)?;
Ok(NodeLocation { latitude, longitude, city }) Ok(NodeLocation {
latitude,
longitude,
city,
})
} }
} }
/// ///
#[derive(Debug, Clone, Copy, PartialEq)] #[derive(Debug, Clone, Copy, PartialEq)]
pub struct BlockDetails { pub struct BlockDetails {
@@ -182,18 +180,20 @@ impl Serialize for BlockDetails {
} }
} }
impl <'de> Deserialize<'de> for BlockDetails { impl<'de> Deserialize<'de> for BlockDetails {
fn deserialize<D>(deserializer: D) -> Result<Self, D::Error> fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
where where
D: serde::Deserializer<'de>, D: serde::Deserializer<'de>,
{ {
let tup = <(u64, BlockHash, u64, u64, Option<u64>)>::deserialize(deserializer)?; let tup = <(u64, BlockHash, u64, u64, Option<u64>)>::deserialize(deserializer)?;
Ok(BlockDetails { Ok(BlockDetails {
block: Block { height: tup.0, hash: tup.1 }, block: Block {
height: tup.0,
hash: tup.1,
},
block_time: tup.2, block_time: tup.2,
block_timestamp: tup.3, block_timestamp: tup.3,
propagation_time: tup.4 propagation_time: tup.4,
}) })
} }
} }
@@ -340,7 +340,8 @@ impl InnerLoop {
.collect(); .collect();
// ... and remove them: // ... and remove them:
self.remove_nodes_and_broadcast_result(node_ids_to_remove).await; self.remove_nodes_and_broadcast_result(node_ids_to_remove)
.await;
} }
} }
} }
@@ -557,11 +558,7 @@ impl InnerLoop {
} }
/// Send a message to all chain feeds. /// Send a message to all chain feeds.
fn broadcast_to_chain_feeds( fn broadcast_to_chain_feeds(&mut self, genesis_hash: &BlockHash, message: ToFeedWebsocket) {
&mut self,
genesis_hash: &BlockHash,
message: ToFeedWebsocket,
) {
if let Some(feeds) = self.chain_to_feed_conn_ids.get(genesis_hash) { if let Some(feeds) = self.chain_to_feed_conn_ids.get(genesis_hash) {
for &feed_id in feeds { for &feed_id in feeds {
if let Some(chan) = self.feed_channels.get_mut(&feed_id) { if let Some(chan) = self.feed_channels.get_mut(&feed_id) {
+1 -1
View File
@@ -231,4 +231,4 @@ impl FeedMessageWrite for AddedNode<'_> {
&node.startup_time(), &node.startup_time(),
)); ));
} }
} }
+9 -8
View File
@@ -94,14 +94,15 @@ async fn start_server(opts: Opts) -> anyhow::Result<()> {
// We can decide how many messages can be buffered to be sent, but not specifically how // We can decide how many messages can be buffered to be sent, but not specifically how
// large those messages are cumulatively allowed to be: // large those messages are cumulatively allowed to be:
ws.max_send_queue(1_000 ).on_upgrade(move |websocket| async move { ws.max_send_queue(1_000)
let (mut tx_to_aggregator, websocket) = .on_upgrade(move |websocket| async move {
handle_feed_websocket_connection(websocket, tx_to_aggregator).await; let (mut tx_to_aggregator, websocket) =
log::info!("Closing /feed connection from {:?}", addr); handle_feed_websocket_connection(websocket, tx_to_aggregator).await;
// Tell the aggregator that this connection has closed, so it can tidy up. log::info!("Closing /feed connection from {:?}", addr);
let _ = tx_to_aggregator.send(FromFeedWebsocket::Disconnected).await; // Tell the aggregator that this connection has closed, so it can tidy up.
let _ = websocket.close().await; let _ = tx_to_aggregator.send(FromFeedWebsocket::Disconnected).await;
}) let _ = websocket.close().await;
})
}); });
// Merge the routes and start our server: // Merge the routes and start our server:
+119 -101
View File
@@ -1,17 +1,19 @@
use test_utils::{ use common::node_types::BlockHash;
feed_message_de::{ FeedMessage, NodeDetails },
server::{ self, Server },
assert_contains_matches
};
use serde_json::json; use serde_json::json;
use std::time::Duration; use std::time::Duration;
use common::node_types::{ BlockHash }; use test_utils::{
assert_contains_matches,
feed_message_de::{FeedMessage, NodeDetails},
server::{self, Server},
};
async fn cargo_run_server() -> Server { async fn cargo_run_server() -> Server {
Server::start(server::StartOpts { Server::start(server::StartOpts {
shard_command: server::cargo_run_commands::telemetry_shard().expect("valid shard command"), shard_command: server::cargo_run_commands::telemetry_shard().expect("valid shard command"),
core_command: server::cargo_run_commands::telemetry_core().expect("valid core command") core_command: server::cargo_run_commands::telemetry_core().expect("valid core command"),
}).await.unwrap() })
.await
.unwrap()
} }
#[tokio::test] #[tokio::test]
@@ -23,7 +25,11 @@ async fn feed_sent_version_on_connect() {
// Expect a version response of 31: // Expect a version response of 31:
let feed_messages = feed_rx.recv_feed_messages().await.unwrap(); let feed_messages = feed_rx.recv_feed_messages().await.unwrap();
assert_eq!(feed_messages, vec![FeedMessage::Version(31)], "expecting version"); assert_eq!(
feed_messages,
vec![FeedMessage::Version(31)],
"expecting version"
);
// Tidy up: // Tidy up:
server.shutdown().await; server.shutdown().await;
@@ -41,7 +47,12 @@ async fn feed_ping_responded_to_with_pong() {
// Expect a pong response: // Expect a pong response:
let feed_messages = feed_rx.recv_feed_messages().await.unwrap(); let feed_messages = feed_rx.recv_feed_messages().await.unwrap();
assert!(feed_messages.contains(&FeedMessage::Pong { msg: "hello!".to_owned() }), "Expecting pong"); assert!(
feed_messages.contains(&FeedMessage::Pong {
msg: "hello!".to_owned()
}),
"Expecting pong"
);
// Tidy up: // Tidy up:
server.shutdown().await; server.shutdown().await;
@@ -54,57 +65,56 @@ async fn feed_add_and_remove_node() {
let shard_id = server.add_shard().await.unwrap(); let shard_id = server.add_shard().await.unwrap();
// Connect a node to the shard: // Connect a node to the shard:
let (mut node_tx, _node_rx) = server.get_shard(shard_id) let (mut node_tx, _node_rx) = server
.get_shard(shard_id)
.unwrap() .unwrap()
.connect() .connect()
.await .await
.expect("can connect to shard"); .expect("can connect to shard");
// Send a "system connected" message: // Send a "system connected" message:
node_tx.send_json_text(json!( node_tx
{ .send_json_text(json!(
"id":1, {
"ts":"2021-07-12T10:37:47.714666+01:00", "id":1,
"payload": { "ts":"2021-07-12T10:37:47.714666+01:00",
"authority":true, "payload": {
"chain":"Local Testnet", "authority":true,
"config":"", "chain":"Local Testnet",
"genesis_hash": BlockHash::from_low_u64_ne(1), "config":"",
"implementation":"Substrate Node", "genesis_hash": BlockHash::from_low_u64_ne(1),
"msg":"system.connected", "implementation":"Substrate Node",
"name":"Alice", "msg":"system.connected",
"network_id":"12D3KooWEyoppNCUx8Yx66oV9fJnriXwCcXwDDUA2kj6vnc6iDEp", "name":"Alice",
"startup_time":"1625565542717", "network_id":"12D3KooWEyoppNCUx8Yx66oV9fJnriXwCcXwDDUA2kj6vnc6iDEp",
"version":"2.0.0-07a1af348-aarch64-macos" "startup_time":"1625565542717",
}, "version":"2.0.0-07a1af348-aarch64-macos"
} },
)).await.unwrap(); }
))
.await
.unwrap();
// Wait a little for this message to propagate to the core // Wait a little for this message to propagate to the core
// (so that our feed connects after the core knows and not before). // (so that our feed connects after the core knows and not before).
tokio::time::sleep(Duration::from_millis(500)).await; tokio::time::sleep(Duration::from_millis(500)).await;
// Connect a feed. // Connect a feed.
let (_feed_tx, mut feed_rx) = server.get_core().connect() let (_feed_tx, mut feed_rx) = server.get_core().connect().await.unwrap();
.await.unwrap();
let feed_messages = feed_rx.recv_feed_messages().await.unwrap(); let feed_messages = feed_rx.recv_feed_messages().await.unwrap();
assert!(feed_messages.contains( assert!(feed_messages.contains(&FeedMessage::AddedChain {
&FeedMessage::AddedChain { name: "Local Testnet".to_owned(),
name: "Local Testnet".to_owned(), node_count: 1
node_count: 1 }));
}
));
// Disconnect the node: // Disconnect the node:
node_tx.close().await.unwrap(); node_tx.close().await.unwrap();
let feed_messages = feed_rx.recv_feed_messages().await.unwrap(); let feed_messages = feed_rx.recv_feed_messages().await.unwrap();
assert!(feed_messages.contains( assert!(feed_messages.contains(&FeedMessage::RemovedChain {
&FeedMessage::RemovedChain { name: "Local Testnet".to_owned(),
name: "Local Testnet".to_owned(), }));
}
));
// Tidy up: // Tidy up:
server.shutdown().await; server.shutdown().await;
@@ -115,36 +125,40 @@ async fn feed_add_and_remove_shard() {
let mut server = cargo_run_server().await; let mut server = cargo_run_server().await;
let mut shards = vec![]; let mut shards = vec![];
for id in 1 ..= 2 { for id in 1..=2 {
// Add a shard: // Add a shard:
let shard_id = server.add_shard().await.unwrap(); let shard_id = server.add_shard().await.unwrap();
// Connect a node to it: // Connect a node to it:
let (mut node_tx, _node_rx) = server.get_shard(shard_id) let (mut node_tx, _node_rx) = server
.get_shard(shard_id)
.unwrap() .unwrap()
.connect() .connect()
.await .await
.expect("can connect to shard"); .expect("can connect to shard");
// Send a "system connected" message: // Send a "system connected" message:
node_tx.send_json_text(json!( node_tx
{ .send_json_text(json!(
"id":id, {
"ts":"2021-07-12T10:37:47.714666+01:00", "id":id,
"payload": { "ts":"2021-07-12T10:37:47.714666+01:00",
"authority":true, "payload": {
"chain": format!("Local Testnet {}", id), "authority":true,
"config":"", "chain": format!("Local Testnet {}", id),
"genesis_hash": BlockHash::from_low_u64_ne(id), "config":"",
"implementation":"Substrate Node", "genesis_hash": BlockHash::from_low_u64_ne(id),
"msg":"system.connected", "implementation":"Substrate Node",
"name":"Alice", "msg":"system.connected",
"network_id":"12D3KooWEyoppNCUx8Yx66oV9fJnriXwCcXwDDUA2kj6vnc6iDEp", "name":"Alice",
"startup_time":"1625565542717", "network_id":"12D3KooWEyoppNCUx8Yx66oV9fJnriXwCcXwDDUA2kj6vnc6iDEp",
"version":"2.0.0-07a1af348-aarch64-macos" "startup_time":"1625565542717",
}, "version":"2.0.0-07a1af348-aarch64-macos"
} },
)).await.unwrap(); }
))
.await
.unwrap();
// Keep what we need to to keep connection alive and let us kill a shard: // Keep what we need to to keep connection alive and let us kill a shard:
shards.push((shard_id, node_tx)); shards.push((shard_id, node_tx));
@@ -155,30 +169,25 @@ async fn feed_add_and_remove_shard() {
// The feed should be told about both of the chains that we've sent info about: // The feed should be told about both of the chains that we've sent info about:
let feed_messages = feed_rx.recv_feed_messages().await.unwrap(); let feed_messages = feed_rx.recv_feed_messages().await.unwrap();
assert!(feed_messages.contains( assert!(feed_messages.contains(&FeedMessage::AddedChain {
&FeedMessage::AddedChain { name: "Local Testnet 1".to_owned(),
name: "Local Testnet 1".to_owned(), node_count: 1
node_count: 1 }));
} assert!(feed_messages.contains(&FeedMessage::AddedChain {
)); name: "Local Testnet 2".to_owned(),
assert!(feed_messages.contains( node_count: 1
&FeedMessage::AddedChain { }));
name: "Local Testnet 2".to_owned(),
node_count: 1
}
));
// Disconnect the first shard: // Disconnect the first shard:
server.kill_shard(shards[0].0).await; server.kill_shard(shards[0].0).await;
// We should be told about the node connected to that shard disconnecting: // We should be told about the node connected to that shard disconnecting:
let feed_messages = feed_rx.recv_feed_messages().await.unwrap(); let feed_messages = feed_rx.recv_feed_messages().await.unwrap();
assert!(feed_messages.contains( assert!(feed_messages.contains(&FeedMessage::RemovedChain {
&FeedMessage::RemovedChain { name: "Local Testnet 1".to_owned(),
name: "Local Testnet 1".to_owned(), }));
} assert!(!feed_messages.contains(
)); // Spot the "!"; this chain was not removed.
assert!(!feed_messages.contains( // Spot the "!"; this chain was not removed.
&FeedMessage::RemovedChain { &FeedMessage::RemovedChain {
name: "Local Testnet 2".to_owned(), name: "Local Testnet 2".to_owned(),
} }
@@ -199,24 +208,27 @@ async fn feed_can_subscribe_and_unsubscribe_from_chain() {
// Send a "system connected" message for a few nodes/chains: // Send a "system connected" message for a few nodes/chains:
for id in 1..=3 { for id in 1..=3 {
node_tx.send_json_text(json!( node_tx
{ .send_json_text(json!(
"id":id, {
"ts":"2021-07-12T10:37:47.714666+01:00", "id":id,
"payload": { "ts":"2021-07-12T10:37:47.714666+01:00",
"authority":true, "payload": {
"chain":format!("Local Testnet {}", id), "authority":true,
"config":"", "chain":format!("Local Testnet {}", id),
"genesis_hash": BlockHash::from_low_u64_ne(id), "config":"",
"implementation":"Substrate Node", "genesis_hash": BlockHash::from_low_u64_ne(id),
"msg":"system.connected", "implementation":"Substrate Node",
"name":format!("Alice {}", id), "msg":"system.connected",
"network_id":"12D3KooWEyoppNCUx8Yx66oV9fJnriXwCcXwDDUA2kj6vnc6iDEp", "name":format!("Alice {}", id),
"startup_time":"1625565542717", "network_id":"12D3KooWEyoppNCUx8Yx66oV9fJnriXwCcXwDDUA2kj6vnc6iDEp",
"version":"2.0.0-07a1af348-aarch64-macos" "startup_time":"1625565542717",
}, "version":"2.0.0-07a1af348-aarch64-macos"
} },
)).await.unwrap(); }
))
.await
.unwrap();
} }
// Connect a feed // Connect a feed
@@ -226,7 +238,10 @@ async fn feed_can_subscribe_and_unsubscribe_from_chain() {
assert_contains_matches!(feed_messages, AddedChain { name, node_count: 1 } if name == "Local Testnet 1"); assert_contains_matches!(feed_messages, AddedChain { name, node_count: 1 } if name == "Local Testnet 1");
// Subscribe it to a chain // Subscribe it to a chain
feed_tx.send_command("subscribe", "Local Testnet 1").await.unwrap(); feed_tx
.send_command("subscribe", "Local Testnet 1")
.await
.unwrap();
let feed_messages = feed_rx.recv_feed_messages().await.unwrap(); let feed_messages = feed_rx.recv_feed_messages().await.unwrap();
assert_contains_matches!( assert_contains_matches!(
@@ -257,7 +272,10 @@ async fn feed_can_subscribe_and_unsubscribe_from_chain() {
.expect_err("Timeout should elapse since no messages sent"); .expect_err("Timeout should elapse since no messages sent");
// We can change our subscription: // We can change our subscription:
feed_tx.send_command("subscribe", "Local Testnet 2").await.unwrap(); feed_tx
.send_command("subscribe", "Local Testnet 2")
.await
.unwrap();
let feed_messages = feed_rx.recv_feed_messages().await.unwrap(); let feed_messages = feed_rx.recv_feed_messages().await.unwrap();
// We are told about the subscription change and given similar on-subscribe messages to above. // We are told about the subscription change and given similar on-subscribe messages to above.
@@ -282,4 +300,4 @@ async fn feed_can_subscribe_and_unsubscribe_from_chain() {
// Tidy up: // Tidy up:
server.shutdown().await; server.shutdown().await;
} }
+2 -1
View File
@@ -96,7 +96,8 @@ impl Aggregator {
}); });
// Establish a resiliant connection to the core (this retries as needed): // Establish a resiliant connection to the core (this retries as needed):
let tx_to_telemetry_core = create_ws_connection_to_core(tx_from_connection, telemetry_uri).await; let tx_to_telemetry_core =
create_ws_connection_to_core(tx_from_connection, telemetry_uri).await;
// Handle any incoming messages in our handler loop: // Handle any incoming messages in our handler loop:
tokio::spawn(Aggregator::handle_messages( tokio::spawn(Aggregator::handle_messages(
+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; use serde_json::value::RawValue;
#[derive(Debug, PartialEq)] #[derive(Debug, PartialEq)]
@@ -77,7 +79,7 @@ pub enum FeedMessage {
address: String, address: String,
block_number: BlockNumber, block_number: BlockNumber,
block_hash: BlockHash, block_hash: BlockHash,
voter: Option<String> voter: Option<String>,
}, },
AfgReceivedPrecommit { AfgReceivedPrecommit {
address: String, address: String,
@@ -85,7 +87,8 @@ pub enum FeedMessage {
block_hash: BlockHash, block_hash: BlockHash,
voter: Option<String>, 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, a1: String,
a2: String, a2: String,
a3: String, a3: String,
@@ -102,8 +105,8 @@ pub enum FeedMessage {
/// A "special" case when we don't know how to decode an action: /// A "special" case when we don't know how to decode an action:
UnknownValue { UnknownValue {
action: u8, action: u8,
value: String value: String,
} },
} }
#[derive(Debug, PartialEq)] #[derive(Debug, PartialEq)]
@@ -138,22 +141,30 @@ impl FeedMessage {
0 => { 0 => {
let version = serde_json::from_str(raw_val.get())?; let version = serde_json::from_str(raw_val.get())?;
FeedMessage::Version(version) FeedMessage::Version(version)
}, }
// BestBlock // BestBlock
1 => { 1 => {
let (block_number, timestamp, avg_block_time) = serde_json::from_str(raw_val.get())?; let (block_number, timestamp, avg_block_time) =
FeedMessage::BestBlock { block_number, timestamp, avg_block_time } serde_json::from_str(raw_val.get())?;
}, FeedMessage::BestBlock {
block_number,
timestamp,
avg_block_time,
}
}
// BestFinalized // BestFinalized
2 => { 2 => {
let (block_number, block_hash) = serde_json::from_str(raw_val.get())?; 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 // AddNode
3 => { 3 => {
let ( let (
node_id, node_id,
( name, implementation, version, validator, network_id ), (name, implementation, version, validator, network_id),
stats, stats,
io, io,
hardware, hardware,
@@ -163,113 +174,153 @@ impl FeedMessage {
) = serde_json::from_str(raw_val.get())?; ) = serde_json::from_str(raw_val.get())?;
// Give these two types but don't use the results: // Give these two types but don't use the results:
let (_,_): (&RawValue, &RawValue) = (io, hardware); let (_, _): (&RawValue, &RawValue) = (io, hardware);
FeedMessage::AddedNode { FeedMessage::AddedNode {
node_id, node_id,
node: NodeDetails { name, implementation, version, validator, network_id }, node: NodeDetails {
name,
implementation,
version,
validator,
network_id,
},
stats, stats,
block_details, block_details,
location, location,
startup_time, startup_time,
} }
}, }
// RemoveNode // RemoveNode
4 => { 4 => {
let node_id = serde_json::from_str(raw_val.get())?; let node_id = serde_json::from_str(raw_val.get())?;
FeedMessage::RemovedNode { node_id } FeedMessage::RemovedNode { node_id }
}, }
// LocatedNode // LocatedNode
5 => { 5 => {
let (node_id, lat, long, city) = serde_json::from_str(raw_val.get())?; 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 // ImportedBlock
6 => { 6 => {
let (node_id, block_details) = serde_json::from_str(raw_val.get())?; 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 // FinalizedBlock
7 => { 7 => {
let (node_id, block_number, block_hash) = serde_json::from_str(raw_val.get())?; 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 // NodeStatsUpdate
8 => { 8 => {
let (node_id, stats) = serde_json::from_str(raw_val.get())?; let (node_id, stats) = serde_json::from_str(raw_val.get())?;
FeedMessage::NodeStatsUpdate { node_id, stats } FeedMessage::NodeStatsUpdate { node_id, stats }
}, }
// Hardware // Hardware
9 => { 9 => {
let (node_id, _hardware): (_, &RawValue) = serde_json::from_str(raw_val.get())?; let (node_id, _hardware): (_, &RawValue) = serde_json::from_str(raw_val.get())?;
FeedMessage::Hardware { node_id } FeedMessage::Hardware { node_id }
}, }
// TimeSync // TimeSync
10 => { 10 => {
let time = serde_json::from_str(raw_val.get())?; let time = serde_json::from_str(raw_val.get())?;
FeedMessage::TimeSync { time } FeedMessage::TimeSync { time }
}, }
// AddedChain // AddedChain
11 => { 11 => {
let (name, node_count) = serde_json::from_str(raw_val.get())?; let (name, node_count) = serde_json::from_str(raw_val.get())?;
FeedMessage::AddedChain { name, node_count } FeedMessage::AddedChain { name, node_count }
}, }
// RemovedChain // RemovedChain
12 => { 12 => {
let name = serde_json::from_str(raw_val.get())?; let name = serde_json::from_str(raw_val.get())?;
FeedMessage::RemovedChain { name } FeedMessage::RemovedChain { name }
}, }
// SubscribedTo // SubscribedTo
13 => { 13 => {
let name = serde_json::from_str(raw_val.get())?; let name = serde_json::from_str(raw_val.get())?;
FeedMessage::SubscribedTo { name } FeedMessage::SubscribedTo { name }
}, }
// UnsubscribedFrom // UnsubscribedFrom
14 => { 14 => {
let name = serde_json::from_str(raw_val.get())?; let name = serde_json::from_str(raw_val.get())?;
FeedMessage::UnsubscribedFrom { name } FeedMessage::UnsubscribedFrom { name }
}, }
// Pong // Pong
15 => { 15 => {
let msg = serde_json::from_str(raw_val.get())?; let msg = serde_json::from_str(raw_val.get())?;
FeedMessage::Pong { msg } FeedMessage::Pong { msg }
}, }
// AfgFinalized // AfgFinalized
16 => { 16 => {
let (address, block_number, block_hash) = serde_json::from_str(raw_val.get())?; 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 // AfgReceivedPrevote
17 => { 17 => {
let (address, block_number, block_hash, voter) = serde_json::from_str(raw_val.get())?; let (address, block_number, block_hash, voter) =
FeedMessage::AfgReceivedPrevote { address, block_number, block_hash, voter } serde_json::from_str(raw_val.get())?;
}, FeedMessage::AfgReceivedPrevote {
address,
block_number,
block_hash,
voter,
}
}
// AfgReceivedPrecommit // AfgReceivedPrecommit
18 => { 18 => {
let (address, block_number, block_hash, voter) = serde_json::from_str(raw_val.get())?; let (address, block_number, block_hash, voter) =
FeedMessage::AfgReceivedPrecommit { address, block_number, block_hash, voter } serde_json::from_str(raw_val.get())?;
}, FeedMessage::AfgReceivedPrecommit {
address,
block_number,
block_hash,
voter,
}
}
// AfgAuthoritySet // AfgAuthoritySet
19 => { 19 => {
let (a1, a2, a3, block_number, block_hash) = serde_json::from_str(raw_val.get())?; 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 // StaleNode
20 => { 20 => {
let node_id = serde_json::from_str(raw_val.get())?; let node_id = serde_json::from_str(raw_val.get())?;
FeedMessage::StaleNode { node_id } FeedMessage::StaleNode { node_id }
}, }
// NodeIOUpdate // NodeIOUpdate
21 => { 21 => {
// ignore NodeIO for now: // ignore NodeIO for now:
let (node_id, _node_io): (_, &RawValue) = serde_json::from_str(raw_val.get())?; let (node_id, _node_io): (_, &RawValue) = serde_json::from_str(raw_val.get())?;
FeedMessage::NodeIOUpdate { node_id } FeedMessage::NodeIOUpdate { node_id }
}, }
// A catchall for messages we don't know/care about yet: // A catchall for messages we don't know/care about yet:
_ => { _ => {
let value = raw_val.to_string(); let value = raw_val.to_string();
FeedMessage::UnknownValue { action, value } FeedMessage::UnknownValue { action, value }
}, }
}; };
Ok(feed_message) Ok(feed_message)
@@ -288,7 +339,9 @@ mod test {
assert_eq!( assert_eq!(
FeedMessage::from_bytes(msg.as_bytes()).unwrap(), 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!( assert_eq!(
FeedMessage::from_bytes(msg.as_bytes()).unwrap(), FeedMessage::from_bytes(msg.as_bytes()).unwrap(),
vec![ vec![
FeedMessage::RemovedChain { name: "".to_owned() }, FeedMessage::RemovedChain {
FeedMessage::AddedChain { name: "Local Testnet".to_owned(), node_count: 1 }, name: "".to_owned()
},
FeedMessage::AddedChain {
name: "Local Testnet".to_owned(),
node_count: 1
},
] ]
); );
} }
}
}
@@ -33,4 +33,4 @@ fn try_find_workspace_dir() -> Result<PathBuf, std::io::Error> {
let mut dir = std::env::current_dir()?; let mut dir = std::env::current_dir()?;
while !dir.ends_with("backend") && dir.pop() {} while !dir.ends_with("backend") && dir.pop() {}
Ok(dir) Ok(dir)
} }
+83 -30
View File
@@ -1,45 +1,65 @@
use std::time::Duration; use std::time::Duration;
use crate::feed_message_de::FeedMessage;
use crate::ws_client; use crate::ws_client;
use futures::{Sink, SinkExt, Stream, StreamExt}; use futures::{Sink, SinkExt, Stream, StreamExt};
use crate::feed_message_de::FeedMessage;
/// Wrap a `ws_client::Sender` with convenient utility methods for shard connections /// Wrap a `ws_client::Sender` with convenient utility methods for shard connections
pub struct ShardSender(ws_client::Sender); pub struct ShardSender(ws_client::Sender);
impl From<ws_client::Sender> for ShardSender { 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 { impl ShardSender {
/// Close this connection /// 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 self.0.close().await
} }
} }
impl Sink<ws_client::Message> for ShardSender { impl Sink<ws_client::Message> for ShardSender {
type Error = ws_client::SendError; 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) 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) 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) 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) self.0.poll_close_unpin(cx)
} }
} }
impl ShardSender { 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"); let bytes = serde_json::to_vec(&json).expect("valid bytes");
self.send(ws_client::Message::Binary(bytes)).await 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"); let s = serde_json::to_string(&json).expect("valid string");
self.send(ws_client::Message::Text(s)).await self.send(ws_client::Message::Text(s)).await
} }
@@ -49,43 +69,70 @@ impl ShardSender {
pub struct ShardReceiver(ws_client::Receiver); pub struct ShardReceiver(ws_client::Receiver);
impl From<ws_client::Receiver> for ShardReceiver { 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 { impl Stream for ShardReceiver {
type Item = Result<ws_client::Message, ws_client::RecvError>; 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) self.0.poll_next_unpin(cx)
} }
} }
/// Wrap a `ws_client::Sender` with convenient utility methods for feed connections /// Wrap a `ws_client::Sender` with convenient utility methods for feed connections
pub struct FeedSender(ws_client::Sender); pub struct FeedSender(ws_client::Sender);
impl From<ws_client::Sender> for FeedSender { 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 { impl Sink<ws_client::Message> for FeedSender {
type Error = ws_client::SendError; 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) 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) 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) 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) self.0.poll_close_unpin(cx)
} }
} }
impl FeedSender { impl FeedSender {
pub async fn send_command<S: AsRef<str>>(&mut self, command: S, param: S) -> Result<(), ws_client::SendError> { pub async fn send_command<S: AsRef<str>>(
self.send(ws_client::Message::Text(format!("{}:{}", command.as_ref(), param.as_ref()))).await &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); pub struct FeedReceiver(ws_client::Receiver);
impl From<ws_client::Receiver> for FeedReceiver { 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 { impl Stream for FeedReceiver {
type Item = Result<ws_client::Message, ws_client::RecvError>; 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()) 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 /// robust in assuming that messages may not all be delivered at once (unless we are
/// specifically testing which messages are buffered together). /// specifically testing which messages are buffered together).
pub async fn recv_feed_messages_once(&mut self) -> Result<Vec<FeedMessage>, anyhow::Error> { pub async fn recv_feed_messages_once(&mut self) -> Result<Vec<FeedMessage>, anyhow::Error> {
let msg = self.0 let msg = self
.0
.next() .next()
.await .await
.ok_or_else(|| anyhow::anyhow!("Stream closed: no more messages"))??; .ok_or_else(|| anyhow::anyhow!("Stream closed: no more messages"))??;
@@ -120,7 +173,7 @@ impl FeedReceiver {
ws_client::Message::Binary(data) => { ws_client::Message::Binary(data) => {
let messages = FeedMessage::from_bytes(&data)?; let messages = FeedMessage::from_bytes(&data)?;
Ok(messages) Ok(messages)
}, }
ws_client::Message::Text(text) => { ws_client::Message::Text(text) => {
let messages = FeedMessage::from_bytes(text.as_bytes())?; let messages = FeedMessage::from_bytes(text.as_bytes())?;
Ok(messages) Ok(messages)
@@ -135,20 +188,20 @@ impl FeedReceiver {
let mut feed_messages = self.recv_feed_messages_once().await?; 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: // Then, loop a little to make sure we catch any additional messages that are sent soon after:
loop { 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 // Timeout elapsed; return the messages we have so far
Err(_) => { Err(_) => {
break Ok(feed_messages); break Ok(feed_messages);
}, }
// Append messages that come back to our vec // Append messages that come back to our vec
Ok(Ok(mut msgs)) => { Ok(Ok(mut msgs)) => {
feed_messages.append(&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),
} }
} }
} }
} }
+2 -2
View File
@@ -1,6 +1,6 @@
mod utils;
mod server; mod server;
mod utils;
pub mod cargo_run_commands; pub mod cargo_run_commands;
pub mod channels; pub mod channels;
pub use server::*; pub use server::*;
+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::ffi::OsString;
use std::marker::PhantomData; use std::marker::PhantomData;
use crate::ws_client; use tokio::process::{self, Command as TokioCommand};
use tokio::process::{ self, Command as TokioCommand };
use super::{ channels, utils };
use common::{ id_type, DenseMap };
id_type! { id_type! {
/// The ID of a running process. Cannot be constructed externally. /// The ID of a running process. Cannot be constructed externally.
@@ -16,7 +16,7 @@ pub struct StartOpts {
pub shard_command: Command, pub shard_command: Command,
/// Command to run to start a telemetry core process. /// Command to run to start a telemetry core process.
/// The `--listen` and `--log` arguments will be appended within and shouldn't be provided. /// 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 { pub struct ConnectToExistingOpts {
@@ -38,7 +38,9 @@ pub enum Error {
ErrorObtainingPort(anyhow::Error), ErrorObtainingPort(anyhow::Error),
#[error("Whoops; attempt to kill a process we didn't start (and so have no handle to)")] #[error("Whoops; attempt to kill a process we didn't start (and so have no handle to)")]
CannotKillNoHandle, 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, CannotAddShardNoHandle,
} }
@@ -68,20 +70,21 @@ impl Server {
} }
pub fn iter_shards(&self) -> impl Iterator<Item = &ShardProcess> { 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 { pub async fn kill_shard(&mut self, id: ProcessId) -> bool {
let shard = match self.shards.remove(id) { let shard = match self.shards.remove(id) {
Some(shard) => shard, Some(shard) => shard,
None => return false None => return false,
}; };
// With this, killing will complete even if the promise returned is cancelled // With this, killing will complete even if the promise returned is cancelled
// (it should regardless, but just to play it safe..) // (it should regardless, but just to play it safe..)
let _ = tokio::spawn(async move { let _ = tokio::spawn(async move {
let _ = shard.kill().await; let _ = shard.kill().await;
}).await; })
.await;
true true
} }
@@ -91,14 +94,9 @@ impl Server {
// Spawn so we don't need to await cleanup if we don't care. // Spawn so we don't need to await cleanup if we don't care.
// Run all kill futs simultaneously. // Run all kill futs simultaneously.
let handle = tokio::spawn(async move { let handle = tokio::spawn(async move {
let shard_kill_futs = self.shards let shard_kill_futs = self.shards.into_iter().map(|(_, s)| s.kill());
.into_iter()
.map(|(_,s)| s.kill());
let _ = tokio::join!( let _ = tokio::join!(futures::future::join_all(shard_kill_futs), self.core.kill());
futures::future::join_all(shard_kill_futs),
self.core.kill()
);
}); });
// You can wait for cleanup but aren't obliged to: // 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> { pub async fn add_shard(&mut self) -> Result<ProcessId, Error> {
let core_uri = match &self.core_shard_submit_uri { let core_uri = match &self.core_shard_submit_uri {
Some(uri) => 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() .clone()
.ok_or_else(|| Error::CannotAddShardNoHandle)? .ok_or_else(|| Error::CannotAddShardNoHandle)?
.into(); .into();
@@ -141,8 +140,9 @@ impl Server {
let _ = utils::wait_for_line_containing( let _ = utils::wait_for_line_containing(
&mut child_stdout, &mut child_stdout,
"Connected to telemetry core", "Connected to telemetry core",
std::time::Duration::from_secs(5) std::time::Duration::from_secs(5),
).await; )
.await;
// Since we're piping stdout from the child process, we need somewhere for it to go // 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: // else the process will get stuck when it tries to produce output:
@@ -156,7 +156,7 @@ impl Server {
id, id,
handle: Some(shard_process), handle: Some(shard_process),
uri: shard_uri, uri: shard_uri,
_channel_type: PhantomData _channel_type: PhantomData,
}); });
Ok(pid) Ok(pid)
@@ -164,7 +164,6 @@ impl Server {
/// Start a telemetry_core process. From here, we can add/remove shards as needed. /// Start a telemetry_core process. From here, we can add/remove shards as needed.
pub async fn start(opts: StartOpts) -> Result<Server, Error> { pub async fn start(opts: StartOpts) -> Result<Server, Error> {
let mut core_cmd: TokioCommand = opts.core_command.into(); let mut core_cmd: TokioCommand = opts.core_command.into();
let mut child = core_cmd let mut child = core_cmd
@@ -194,16 +193,18 @@ impl Server {
Ok(Server { Ok(Server {
shard_command: Some(opts.shard_command), shard_command: Some(opts.shard_command),
core_shard_submit_uri: Some(format!("http://127.0.0.1:{}/shard_submit", core_port) core_shard_submit_uri: Some(
.parse() format!("http://127.0.0.1:{}/shard_submit", core_port)
.expect("valid shard_submit URI")), .parse()
.expect("valid shard_submit URI"),
),
shards: DenseMap::new(), shards: DenseMap::new(),
core: Process { core: Process {
id: ProcessId(0), id: ProcessId(0),
handle: Some(child), handle: Some(child),
uri: feed_uri, uri: feed_uri,
_channel_type: PhantomData, _channel_type: PhantomData,
} },
}) })
} }
@@ -229,12 +230,11 @@ impl Server {
uri: opts.feed_uri, uri: opts.feed_uri,
handle: None, handle: None,
_channel_type: PhantomData, _channel_type: PhantomData,
} },
} }
} }
} }
/// This represents a running process that we can connect to, which /// This represents a running process that we can connect to, which
/// may be either a `telemetry_shard` or `telemetry_core`. /// may be either a `telemetry_shard` or `telemetry_core`.
pub struct Process<Channel> { pub struct Process<Channel> {
@@ -245,7 +245,7 @@ pub struct Process<Channel> {
/// The URI that we can use to connect to the process socket. /// The URI that we can use to connect to the process socket.
uri: http::Uri, uri: http::Uri,
/// The kind of the process (lets us add methods specific to shard/core). /// 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. /// 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. /// A core process with core-specific methods.
pub type CoreProcess = Process<(channels::FeedSender, channels::FeedReceiver)>; pub type CoreProcess = Process<(channels::FeedSender, channels::FeedReceiver)>;
impl <Channel> Process<Channel> { impl<Channel> Process<Channel> {
/// Get the ID of this process /// Get the ID of this process
pub fn id(&self) -> ProcessId { pub fn id(&self) -> ProcessId {
self.id self.id
@@ -265,25 +265,28 @@ impl <Channel> Process<Channel> {
async fn kill(self) -> Result<(), Error> { async fn kill(self) -> Result<(), Error> {
match self.handle { match self.handle {
Some(mut handle) => Ok(handle.kill().await?), 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 /// Establish a connection to the process
pub async fn connect(&self) -> Result<(Send, Recv), Error> { pub async fn connect(&self) -> Result<(Send, Recv), Error> {
ws_client::connect(&self.uri) ws_client::connect(&self.uri)
.await .await
.map(|(s,r)| (s.into(), r.into())) .map(|(s, r)| (s.into(), r.into()))
.map_err(|e| e.into()) .map_err(|e| e.into())
} }
/// Establish multiple connections to the process /// 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) utils::connect_multiple_to_uri(&self.uri, num_connections)
.await .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()) .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)] #[derive(Clone, Debug)]
pub struct Command { pub struct Command {
command: OsString, command: OsString,
args: Vec<OsString> args: Vec<OsString>,
} }
impl Command { impl Command {
pub fn new<S: Into<OsString>>(command: S) -> Command { pub fn new<S: Into<OsString>>(command: S) -> Command {
Command { Command {
command: command.into(), command: command.into(),
args: Vec::new() args: Vec::new(),
} }
} }
@@ -317,4 +320,4 @@ impl Into<TokioCommand> for Command {
cmd.args(self.args); cmd.args(self.args);
cmd cmd
} }
} }
+28 -16
View File
@@ -1,8 +1,8 @@
use crate::ws_client; use crate::ws_client;
use anyhow::{anyhow, Context};
use tokio::io::BufReader; use tokio::io::BufReader;
use tokio::io::{ AsyncRead, AsyncWrite, AsyncBufReadExt }; use tokio::io::{AsyncBufReadExt, AsyncRead, AsyncWrite};
use tokio::time::Duration; 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, /// 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 /// 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, /// 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 /// such that if we don't see a new line of output within the timeout we bail out
/// and return an error. /// 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 reader = BufReader::new(reader);
let mut reader_lines = reader.lines(); let mut reader_lines = reader.lines();
loop { loop {
let line = tokio::time::timeout( let line = tokio::time::timeout(max_wait_between_lines, reader_lines.next_line()).await;
max_wait_between_lines,
reader_lines.next_line()
).await;
let line = match line { let line = match line {
// timeout expired; couldn't get port: // 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: // 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: // 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! // All OK, and a line is given back; phew!
Ok(Ok(Some(line))) => line Ok(Ok(Some(line))) => line,
}; };
if line.contains(text) { 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. /// 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> { pub async fn connect_multiple_to_uri(
let connect_futs = (0..num_connections) uri: &http::Uri,
.map(|_| ws_client::connect(uri)); num_connections: usize,
let sockets: Result<Vec<_>,_> = futures::future::join_all(connect_futs) ) -> 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 .await
.into_iter() .into_iter()
.collect(); .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) pub fn drain<R, W>(mut reader: R, mut writer: W)
where where
R: AsyncRead + Unpin + Send + 'static, R: AsyncRead + Unpin + Send + 'static,
W: AsyncWrite + Unpin + Send + 'static W: AsyncWrite + Unpin + Send + 'static,
{ {
tokio::spawn(async move { tokio::spawn(async move {
let _ = tokio::io::copy(&mut reader, &mut writer).await; let _ = tokio::io::copy(&mut reader, &mut writer).await;
}); });
} }
+60 -42
View File
@@ -1,17 +1,17 @@
use futures::channel::{ mpsc }; use futures::channel::mpsc;
use soketto::handshake::{Client, ServerResponse};
use tokio_util::compat::{ TokioAsyncReadCompatExt };
use tokio::net::TcpStream;
use futures::{Sink, SinkExt, Stream, StreamExt}; 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 /// Send messages into the connection
#[derive(Clone)] #[derive(Clone)]
pub struct Sender { pub struct Sender {
inner: mpsc::UnboundedSender<SentMessage> inner: mpsc::UnboundedSender<SentMessage>,
} }
impl Sender { impl Sender {
pub async fn close(&mut self) -> Result<(),SendError> { pub async fn close(&mut self) -> Result<(), SendError> {
self.inner.send(SentMessage::Close).await?; self.inner.send(SentMessage::Close).await?;
Ok(()) Ok(())
} }
@@ -20,28 +20,39 @@ impl Sender {
#[derive(thiserror::Error, Debug, Clone)] #[derive(thiserror::Error, Debug, Clone)]
pub enum SendError { pub enum SendError {
#[error("Failed to send message: {0}")] #[error("Failed to send message: {0}")]
ChannelError(#[from] mpsc::SendError) ChannelError(#[from] mpsc::SendError),
} }
impl Sink<Message> for Sender { impl Sink<Message> for Sender {
type Error = SendError; 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()) 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> { 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()) 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()) self.inner.poll_close_unpin(cx).map_err(|e| e.into())
} }
} }
/// Receive messages out of a connection /// Receive messages out of a connection
pub struct Receiver { pub struct Receiver {
inner: mpsc::UnboundedReceiver<Result<Message,RecvError>> inner: mpsc::UnboundedReceiver<Result<Message, RecvError>>,
} }
#[derive(thiserror::Error, Debug)] #[derive(thiserror::Error, Debug)]
@@ -49,12 +60,15 @@ pub enum RecvError {
#[error("Text message contains invalid UTF8: {0}")] #[error("Text message contains invalid UTF8: {0}")]
InvalidUtf8(#[from] std::string::FromUtf8Error), InvalidUtf8(#[from] std::string::FromUtf8Error),
#[error("Stream finished")] #[error("Stream finished")]
StreamFinished StreamFinished,
} }
impl Stream for Receiver { impl Stream for Receiver {
type Item = Result<Message, RecvError>; 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()) 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 /// A message type that can be sent or received from the connection
pub enum Message { pub enum Message {
Text(String), Text(String),
Binary(Vec<u8>) Binary(Vec<u8>),
} }
/// Sent messages can be anything publically visible, or a close message. /// Sent messages can be anything publically visible, or a close message.
enum SentMessage { enum SentMessage {
Message(Message), Message(Message),
Close Close,
} }
#[derive(thiserror::Error, Debug)] #[derive(thiserror::Error, Debug)]
@@ -98,12 +112,10 @@ pub async fn connect(uri: &http::Uri) -> Result<(Sender, Receiver), ConnectError
// Establish a WS connection: // Establish a WS connection:
let mut client = Client::new(socket.compat(), host, &path); let mut client = Client::new(socket.compat(), host, &path);
let (mut ws_to_connection, mut ws_from_connection) = match client.handshake().await? { let (mut ws_to_connection, mut ws_from_connection) = match client.handshake().await? {
ServerResponse::Accepted { .. } => { ServerResponse::Accepted { .. } => client.into_builder().finish(),
client.into_builder().finish()
},
ServerResponse::Redirect { status_code, .. } => { ServerResponse::Redirect { status_code, .. } => {
return Err(ConnectError::ConnectionFailedRedirect { status_code }) return Err(ConnectError::ConnectionFailedRedirect { status_code })
}, }
ServerResponse::Rejected { status_code } => { ServerResponse::Rejected { status_code } => {
return Err(ConnectError::ConnectionFailedRejected { 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) => { Err(e) => {
// Couldn't receive data may mean all senders are gone, so log // Couldn't receive data may mean all senders are gone, so log
// the error and shut this down: // 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; break;
},
Ok(data) => {
data
} }
Ok(data) => data,
}; };
let msg = match message_data { let msg = match message_data {
soketto::Data::Text(_) => { soketto::Data::Text(_) => Ok(Message::Binary(data)),
Ok(Message::Binary(data)) soketto::Data::Binary(_) => String::from_utf8(data)
}, .map(|s| Message::Text(s))
soketto::Data::Binary(_) => { .map_err(|e| e.into()),
String::from_utf8(data)
.map(|s| Message::Text(s))
.map_err(|e| e.into())
},
}; };
data = Vec::with_capacity(128); 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 { if let Err(e) = tx_to_external.send(msg).await {
// Failure to send likely means that the recv has been dropped, // Failure to send likely means that the recv has been dropped,
// so let's drop this loop too. // 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; break;
} }
} }
@@ -161,16 +173,22 @@ pub async fn connect(uri: &http::Uri) -> Result<(Sender, Receiver), ConnectError
match msg { match msg {
SentMessage::Message(Message::Text(s)) => { SentMessage::Message(Message::Text(s)) => {
if let Err(e) = ws_to_connection.send_text_owned(s).await { 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; break;
} }
}, }
SentMessage::Message(Message::Binary(bytes)) => { SentMessage::Message(Message::Binary(bytes)) => {
if let Err(e) = ws_to_connection.send_binary_mut(bytes).await { 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; break;
} }
}, }
SentMessage::Close => { SentMessage::Close => {
if let Err(e) = ws_to_connection.close().await { if let Err(e) = ws_to_connection.close().await {
log::error!("Error attempting to close connection: {}", e); 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 { 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; break;
} }
} }
}); });
Ok(( Ok((Sender { inner: tx_to_ws }, Receiver { inner: rx_from_ws }))
Sender { inner: tx_to_ws },
Receiver { inner: rx_from_ws }
))
} }