summaryrefslogtreecommitdiffstats
path: root/src/transport
diff options
context:
space:
mode:
Diffstat (limited to 'src/transport')
-rw-r--r--src/transport/firewall.rs4
-rw-r--r--src/transport/mod.rs27
-rw-r--r--src/transport/mycelium.rs21
-rw-r--r--src/transport/traffic.rs199
4 files changed, 243 insertions, 8 deletions
diff --git a/src/transport/firewall.rs b/src/transport/firewall.rs
index 84fdc59..b8e05cd 100644
--- a/src/transport/firewall.rs
+++ b/src/transport/firewall.rs
@@ -89,12 +89,12 @@ impl Transport for FirewallTransport {
let firewall = Arc::clone(&self.firewall);
self.inner
- .start_consumer(Box::new(move |sender, message| {
+ .start_consumer(Box::new(move |sender, message, payload_len| {
if firewall.blocks_inbound(&sender) {
tracing::trace!(peer = %sender, "Firewall dropped inbound message");
return Box::pin(async { Ok(()) });
}
- handler(sender, message)
+ handler(sender, message, payload_len)
}))
.await
}
diff --git a/src/transport/mod.rs b/src/transport/mod.rs
index 19c4126..b2f5c9d 100644
--- a/src/transport/mod.rs
+++ b/src/transport/mod.rs
@@ -11,6 +11,7 @@ pub use crate::document::replica::Epoch;
pub mod firewall;
pub mod mycelium;
+pub mod traffic;
/// Random identifier of one run of a node (from start to exit).
pub type RunId = [u8; 16];
@@ -101,7 +102,27 @@ pub enum ParsedMessage {
Pong(Keepalive),
}
+/// What a message is, without its contents.
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum MessageKind {
+ Hello,
+ HelloReply,
+ Sync,
+ Ping,
+ Pong,
+}
+
impl ParsedMessage {
+ pub fn kind(&self) -> MessageKind {
+ match self {
+ ParsedMessage::Hello(_) => MessageKind::Hello,
+ ParsedMessage::HelloReply(_) => MessageKind::HelloReply,
+ ParsedMessage::Sync(_) => MessageKind::Sync,
+ ParsedMessage::Ping(_) => MessageKind::Ping,
+ ParsedMessage::Pong(_) => MessageKind::Pong,
+ }
+ }
+
/// Every message carries the sender's epoch.
pub fn epoch(&self) -> Epoch {
match self {
@@ -161,9 +182,13 @@ pub struct SyncEnvelope {
pub struct EncodedMessage {
pub peer: String,
pub body: bytes::Bytes,
+ pub kind: MessageKind,
+ /// Bytes the message takes on the wire (without the transport's own framing).
+ pub payload_len: usize,
}
-pub type HandlerFn = dyn Fn(String, ParsedMessage) -> Pin<Box<dyn Future<Output = Result<()>> + Send>>
+/// Called with the sender, the message and its `payload_len` (as in `EncodedMessage`).
+pub type HandlerFn = dyn Fn(String, ParsedMessage, usize) -> Pin<Box<dyn Future<Output = Result<()>> + Send>>
+ Send
+ Sync
+ 'static;
diff --git a/src/transport/mycelium.rs b/src/transport/mycelium.rs
index 9075ae8..cc7f25f 100644
--- a/src/transport/mycelium.rs
+++ b/src/transport/mycelium.rs
@@ -559,10 +559,12 @@ impl Transport for MyceliumTransport {
}
fn encode(&self, peer: &str, message: &ParsedMessage) -> Result<EncodedMessage> {
- let outbound = format_message(peer, message)?;
+ let (outbound, payload_len) = format_message(peer, message)?;
Ok(EncodedMessage {
peer: peer.to_owned(),
body: serde_json::to_vec(&outbound)?.into(),
+ kind: message.kind(),
+ payload_len,
})
}
@@ -685,7 +687,8 @@ impl Transport for MyceliumTransport {
//payload_bytes = raw_bytes.len(),
"Received message"
);
- if let Err(error) = handler(sender, parsed).await {
+ let payload_len = base64_decoded_len(&inbound_message.payload);
+ if let Err(error) = handler(sender, parsed, payload_len).await {
tracing::warn!(error = %format!("{error:#}"), "Failed to process inbound message");
}
}
@@ -720,6 +723,12 @@ fn encode_payload(payload: &Vec<u8>) -> String {
BASE64.encode(payload)
}
+/// Length of what padded base64 `encoded` decodes to.
+fn base64_decoded_len(encoded: &str) -> usize {
+ let padding = encoded.bytes().rev().take_while(|&b| b == b'=').count();
+ (encoded.len() / 4 * 3).saturating_sub(padding)
+}
+
fn decode_payload(topic_encoded: &str, payload: &str) -> Result<ParsedMessage, TransportError> {
let topic: &str = &decode_topic(topic_encoded)?;
let bytes = BASE64.decode(payload)?;
@@ -754,10 +763,11 @@ fn parse_message(message: &InboundMessage) -> Result<(String, ParsedMessage), Tr
Ok((sender, parsed))
}
+/// The message for Mycelium's API, and the length of its payload.
fn format_message(
recipient: &str,
message: &ParsedMessage,
-) -> Result<OutboundMessage, TransportError> {
+) -> Result<(OutboundMessage, usize), TransportError> {
let (topic, payload_bytes) = match message {
ParsedMessage::Hello(req) => (TOPIC_HELLO, hello_request_to_bytes(req)?),
ParsedMessage::HelloReply(res) => (TOPIC_HELLO_REPLY, hello_reply_to_bytes(res)?),
@@ -765,13 +775,14 @@ fn format_message(
ParsedMessage::Ping(k) => (TOPIC_PING, keepalive_to_bytes(k)?),
ParsedMessage::Pong(k) => (TOPIC_PONG, keepalive_to_bytes(k)?),
};
- Ok(OutboundMessage {
+ let outbound = OutboundMessage {
dst: Destination {
pk: recipient.to_owned(),
},
topic: encode_topic(topic),
payload: encode_payload(&payload_bytes),
- })
+ };
+ Ok((outbound, payload_bytes.len()))
}
// =====================================================================
diff --git a/src/transport/traffic.rs b/src/transport/traffic.rs
new file mode 100644
index 0000000..bb93d48
--- /dev/null
+++ b/src/transport/traffic.rs
@@ -0,0 +1,199 @@
+//! Messages and bytes exchanged with each peer, by kind: counted where every message
+//! passes anyway, at the transport. No messages of our own.
+//!
+//! Bytes are payload bytes, what the Mycelium message carries, without its framing,
+//! the topic, or the base64 of its HTTP API. Counters only grow (from the start of the
+//! process); rates are for whoever reads them.
+//!
+//! A message dropped by the debug firewall still counts as sent (the firewall plays a
+//! lossy network) but not as received.
+
+use std::collections::HashMap;
+use std::sync::{Arc, Mutex, MutexGuard};
+
+use anyhow::Result;
+use async_trait::async_trait;
+use serde::Serialize;
+use utoipa::ToSchema;
+
+use crate::transport::{EncodedMessage, MessageHandler, MessageKind, ParsedMessage, Transport};
+
+#[derive(Debug, Default, Clone, Copy, Serialize, ToSchema)]
+pub struct Counter {
+ pub messages: u64,
+ pub bytes: u64,
+}
+
+impl Counter {
+ fn add(&mut self, bytes: usize) {
+ self.messages += 1;
+ self.bytes += bytes as u64;
+ }
+}
+
+#[derive(Debug, Default, Clone, Copy, Serialize, ToSchema)]
+#[serde(rename_all = "camelCase")]
+pub struct ByKind {
+ pub hello: Counter,
+ pub hello_reply: Counter,
+ pub sync: Counter,
+ pub ping: Counter,
+ pub pong: Counter,
+}
+
+impl ByKind {
+ fn get_mut(&mut self, kind: MessageKind) -> &mut Counter {
+ match kind {
+ MessageKind::Hello => &mut self.hello,
+ MessageKind::HelloReply => &mut self.hello_reply,
+ MessageKind::Sync => &mut self.sync,
+ MessageKind::Ping => &mut self.ping,
+ MessageKind::Pong => &mut self.pong,
+ }
+ }
+}
+
+#[derive(Debug, Default, Clone, Copy, Serialize, ToSchema)]
+#[serde(rename_all = "camelCase")]
+pub struct PeerTraffic {
+ /// Handed to Mycelium (it accepted them).
+ pub sent: ByKind,
+ pub received: ByKind,
+ /// Sends Mycelium refused (e.g. no route to the peer); not counted as sent.
+ pub refused: ByKind,
+}
+
+#[derive(Debug, Default)]
+struct Table {
+ peers: HashMap<String, PeerTraffic>,
+ /// All peers, including ones no longer listed.
+ total: PeerTraffic,
+}
+
+#[derive(Debug, Default)]
+pub struct Traffic {
+ table: Mutex<Table>,
+}
+
+impl Traffic {
+ fn lock(&self) -> MutexGuard<'_, Table> {
+ self.table
+ .lock()
+ .unwrap_or_else(|poisoned| poisoned.into_inner())
+ }
+
+ fn count(&self, peer: &str, bytes: usize, pick: impl Fn(&mut PeerTraffic) -> &mut Counter) {
+ let mut guard = self.lock();
+ let table = &mut *guard;
+ pick(&mut table.total).add(bytes);
+ match table.peers.get_mut(peer) {
+ Some(traffic) => pick(traffic).add(bytes),
+ None => pick(table.peers.entry(peer.to_owned()).or_default()).add(bytes),
+ }
+ }
+
+ fn sent(&self, peer: &str, kind: MessageKind, bytes: usize) {
+ self.count(peer, bytes, |t| t.sent.get_mut(kind));
+ }
+
+ fn refused(&self, peer: &str, kind: MessageKind, bytes: usize) {
+ self.count(peer, bytes, |t| t.refused.get_mut(kind));
+ }
+
+ fn received(&self, peer: &str, kind: MessageKind, bytes: usize) {
+ self.count(peer, bytes, |t| t.received.get_mut(kind));
+ }
+
+ pub fn peer(&self, peer: &str) -> Option<PeerTraffic> {
+ self.lock().peers.get(peer).copied()
+ }
+
+ pub fn total(&self) -> PeerTraffic {
+ self.lock().total
+ }
+
+ /// Drops the counters of peers not kept (anyone can send us messages; the totals
+ /// keep their share).
+ pub fn retain(&self, mut keep: impl FnMut(&str) -> bool) {
+ self.lock().peers.retain(|peer, _| keep(peer));
+ }
+}
+
+/// Counts what passes through `inner`.
+pub struct CountingTransport {
+ inner: Arc<dyn Transport>,
+ traffic: Arc<Traffic>,
+}
+
+impl CountingTransport {
+ pub fn new(inner: Arc<dyn Transport>, traffic: Arc<Traffic>) -> Self {
+ Self { inner, traffic }
+ }
+}
+
+#[async_trait]
+impl Transport for CountingTransport {
+ fn local_id(&self) -> &str {
+ self.inner.local_id()
+ }
+
+ async fn resolve_seed(&self, seed: &str) -> Result<String> {
+ self.inner.resolve_seed(seed).await
+ }
+
+ fn encode(&self, peer: &str, message: &ParsedMessage) -> Result<EncodedMessage> {
+ self.inner.encode(peer, message)
+ }
+
+ async fn send_encoded(&self, message: &EncodedMessage) -> Result<()> {
+ let result = self.inner.send_encoded(message).await;
+ if result.is_ok() {
+ self.traffic
+ .sent(&message.peer, message.kind, message.payload_len);
+ } else {
+ self.traffic
+ .refused(&message.peer, message.kind, message.payload_len);
+ }
+ result
+ }
+
+ async fn start_consumer(&self, handler: MessageHandler) -> Result<()> {
+ let traffic = Arc::clone(&self.traffic);
+ self.inner
+ .start_consumer(Box::new(move |sender, message, payload_len| {
+ traffic.received(&sender, message.kind(), payload_len);
+ handler(sender, message, payload_len)
+ }))
+ .await
+ }
+
+ async fn wait_for_exit(&self) -> std::io::Result<std::process::ExitStatus> {
+ self.inner.wait_for_exit().await
+ }
+}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+
+ #[test]
+ fn totals_survive_retain() {
+ let traffic = Traffic::default();
+ traffic.sent("a", MessageKind::Ping, 100);
+ traffic.sent("b", MessageKind::Sync, 1000);
+ traffic.received("b", MessageKind::Pong, 90);
+ traffic.refused("a", MessageKind::Sync, 500);
+
+ let b = traffic.peer("b").unwrap();
+ assert_eq!(b.sent.sync.messages, 1);
+ assert_eq!(b.received.pong.bytes, 90);
+
+ traffic.retain(|peer| peer == "b");
+ assert!(traffic.peer("a").is_none());
+
+ let total = traffic.total();
+ assert_eq!(total.sent.ping.bytes, 100);
+ assert_eq!(total.sent.sync.bytes, 1000);
+ assert_eq!(total.refused.sync.bytes, 500);
+ }
+}
generated by cgit v1.3.1 (git 2.54.0) at 2026-10-24 23:30:37 +0000