improved bandwidth information propagation within the client

This commit is contained in:
Jędrzej Stuczyński
2024-07-23 09:56:19 +01:00
parent aa00711138
commit 4630eb6cd4
6 changed files with 147 additions and 52 deletions
@@ -0,0 +1,83 @@
// Copyright 2024 - Nym Technologies SA <contact@nymtech.net>
// SPDX-License-Identifier: Apache-2.0
use si_scale::helpers::bibytes2;
use std::sync::atomic::{AtomicI64, Ordering};
use std::sync::Arc;
use std::time::Duration;
use time::OffsetDateTime;
#[derive(Clone, Default)]
pub struct ClientBandwidth {
inner: Arc<ClientBandwidthInner>,
}
#[derive(Default)]
struct ClientBandwidthInner {
/// the actual bandwidth amount (in bytes) available
available: AtomicI64,
/// defines the timestamp when the bandwidth information has been logged to the logs stream
last_logged_ts: AtomicI64,
/// defines the timestamp when the bandwidth value was last updated
last_updated_ts: AtomicI64,
}
impl ClientBandwidth {
pub(crate) fn new_empty() -> Self {
ClientBandwidth {
inner: Arc::new(ClientBandwidthInner {
available: AtomicI64::new(0),
last_logged_ts: AtomicI64::new(0),
last_updated_ts: AtomicI64::new(0),
}),
}
}
pub(crate) fn remaining(&self) -> i64 {
self.inner.available.load(Ordering::Acquire)
}
pub(crate) fn maybe_log_bandwidth(&self, now: Option<OffsetDateTime>) {
let last = self.last_logged();
let now = now.unwrap_or_else(OffsetDateTime::now_utc);
if last + Duration::from_secs(10) < now {
self.log_bandwidth(Some(now))
}
}
pub(crate) fn log_bandwidth(&self, now: Option<OffsetDateTime>) {
let now = now.unwrap_or_else(OffsetDateTime::now_utc);
let remaining_bi2 = bibytes2(self.inner.available.load(Ordering::Relaxed) as f64);
log::info!("remaining bandwidth: {remaining_bi2}");
self.inner
.last_logged_ts
.store(now.unix_timestamp(), Ordering::Relaxed)
}
pub(crate) fn update_and_maybe_log(&self, remaining: i64) {
let now = OffsetDateTime::now_utc();
self.inner.available.store(remaining, Ordering::Release);
self.inner
.last_updated_ts
.store(now.unix_timestamp(), Ordering::Relaxed);
self.maybe_log_bandwidth(Some(now))
}
pub(crate) fn update_and_log(&self, remaining: i64) {
let now = OffsetDateTime::now_utc();
self.inner.available.store(remaining, Ordering::Release);
self.inner
.last_updated_ts
.store(now.unix_timestamp(), Ordering::Relaxed);
self.log_bandwidth(Some(now))
}
fn last_logged(&self) -> OffsetDateTime {
// SAFETY: this value is always populated with valid timestamps
OffsetDateTime::from_unix_timestamp(self.inner.last_logged_ts.load(Ordering::Relaxed))
.unwrap()
}
}
+21 -17
View File
@@ -28,7 +28,6 @@ use nym_sphinx::forwarding::packet::MixPacket;
use nym_task::TaskClient;
use nym_validator_client::nyxd::contract_traits::DkgQueryClient;
use rand::rngs::OsRng;
use std::sync::atomic::{AtomicI64, Ordering};
use std::sync::Arc;
use std::time::Duration;
use tungstenite::protocol::Message;
@@ -41,6 +40,7 @@ use tokio::time::sleep;
#[cfg(not(target_arch = "wasm32"))]
use tokio_tungstenite::connect_async;
use crate::bandwidth::ClientBandwidth;
#[cfg(not(unix))]
use std::os::raw::c_int as RawFd;
#[cfg(target_arch = "wasm32")]
@@ -79,11 +79,10 @@ impl GatewayConfig {
}
// TODO: this should be refactored into a state machine that keeps track of its authentication state
#[derive(Debug)]
pub struct GatewayClient<C, St = EphemeralCredentialStorage> {
authenticated: bool,
disabled_credentials_mode: bool,
bandwidth_remaining: Arc<AtomicI64>,
bandwidth: ClientBandwidth,
gateway_address: String,
gateway_identity: identity::PublicKey,
local_identity: Arc<identity::KeyPair>,
@@ -122,7 +121,7 @@ impl<C, St> GatewayClient<C, St> {
GatewayClient {
authenticated: false,
disabled_credentials_mode: true,
bandwidth_remaining: Arc::new(AtomicI64::new(0)),
bandwidth: ClientBandwidth::new_empty(),
gateway_address: config.gateway_listener,
gateway_identity: config.gateway_identity,
local_identity,
@@ -182,7 +181,7 @@ impl<C, St> GatewayClient<C, St> {
}
pub fn remaining_bandwidth(&self) -> i64 {
self.bandwidth_remaining.load(Ordering::Acquire)
self.bandwidth.remaining()
}
#[cfg(not(target_arch = "wasm32"))]
@@ -537,8 +536,8 @@ impl<C, St> GatewayClient<C, St> {
} => {
self.check_gateway_protocol(protocol_version)?;
self.authenticated = status;
self.bandwidth_remaining
.store(bandwidth_remaining, Ordering::Release);
self.bandwidth.update_and_maybe_log(bandwidth_remaining);
self.negotiated_protocol = protocol_version;
log::debug!("authenticated: {status}, bandwidth remaining: {bandwidth_remaining}");
self.task_client.send_status_msg(Box::new(
@@ -595,8 +594,11 @@ impl<C, St> GatewayClient<C, St> {
ServerResponse::Error { message } => Err(GatewayClientError::GatewayError(message)),
_ => Err(GatewayClientError::UnexpectedResponse),
}?;
self.bandwidth_remaining
.store(bandwidth_remaining, Ordering::Relaxed);
// TODO: create tracing span
info!("managed to claim ecash bandwidth");
self.bandwidth.update_and_log(bandwidth_remaining);
Ok(())
}
@@ -607,8 +609,10 @@ impl<C, St> GatewayClient<C, St> {
ServerResponse::Error { message } => Err(GatewayClientError::GatewayError(message)),
_ => Err(GatewayClientError::UnexpectedResponse),
}?;
self.bandwidth_remaining
.store(bandwidth_remaining, Ordering::Release);
info!("managed to claim testnet bandwidth");
self.bandwidth.update_and_log(bandwidth_remaining);
Ok(())
}
@@ -683,7 +687,7 @@ impl<C, St> GatewayClient<C, St> {
if !self.authenticated {
return Err(GatewayClientError::NotAuthenticated);
}
let bandwidth_remaining = self.bandwidth_remaining.load(Ordering::Acquire);
let bandwidth_remaining = self.bandwidth.remaining();
if bandwidth_remaining < REMAINING_BANDWIDTH_THRESHOLD {
self.claim_bandwidth().await?;
}
@@ -755,7 +759,7 @@ impl<C, St> GatewayClient<C, St> {
if !self.authenticated {
return Err(GatewayClientError::NotAuthenticated);
}
let bandwidth_remaining = self.bandwidth_remaining.load(Ordering::Acquire);
let bandwidth_remaining = self.bandwidth.remaining();
if bandwidth_remaining < REMAINING_BANDWIDTH_THRESHOLD {
self.claim_bandwidth().await?;
}
@@ -813,7 +817,7 @@ impl<C, St> GatewayClient<C, St> {
.as_ref()
.expect("no shared key present even though we're authenticated!"),
),
self.bandwidth_remaining.clone(),
self.bandwidth.clone(),
self.task_client.clone(),
)
}
@@ -854,7 +858,7 @@ impl<C, St> GatewayClient<C, St> {
self.establish_connection().await?;
}
let shared_key = self.perform_initial_authentication().await?;
let bandwidth_remaining = self.bandwidth_remaining.load(Ordering::Acquire);
let bandwidth_remaining = self.bandwidth.remaining();
if bandwidth_remaining < REMAINING_BANDWIDTH_THRESHOLD {
info!("Claiming more bandwidth with existing credentials. Stop the process now if you don't want that to happen.");
self.claim_bandwidth().await?;
@@ -893,7 +897,7 @@ impl GatewayClient<InitOnly, EphemeralCredentialStorage> {
GatewayClient {
authenticated: false,
disabled_credentials_mode: true,
bandwidth_remaining: Arc::new(AtomicI64::new(0)),
bandwidth: ClientBandwidth::new_empty(),
gateway_address: gateway_listener.to_string(),
gateway_identity,
local_identity,
@@ -925,7 +929,7 @@ impl GatewayClient<InitOnly, EphemeralCredentialStorage> {
GatewayClient {
authenticated: self.authenticated,
disabled_credentials_mode: self.disabled_credentials_mode,
bandwidth_remaining: self.bandwidth_remaining,
bandwidth: self.bandwidth,
gateway_address: self.gateway_address,
gateway_identity: self.gateway_identity,
local_identity: self.local_identity,
@@ -1,21 +1,26 @@
// Copyright 2021 - Nym Technologies SA <contact@nymtech.net>
// SPDX-License-Identifier: Apache-2.0
#[cfg(target_arch = "wasm32")]
use gloo_utils::errors::JsError;
use nym_gateway_requests::registration::handshake::error::HandshakeError;
use nym_gateway_requests::SimpleGatewayRequestsError;
use std::io;
use thiserror::Error;
use tungstenite::Error as WsError;
#[cfg(target_arch = "wasm32")]
use gloo_utils::errors::JsError;
#[derive(Debug, Error)]
pub enum GatewayClientError {
#[error("Connection to the gateway is not established")]
ConnectionNotEstablished,
#[error("Gateway returned an error response: {0}")]
#[error("gateway returned an error response: {0}")]
GatewayError(String),
#[error("gateway returned an error response: {0}")]
TypedGatewayError(SimpleGatewayRequestsError),
#[error("There was a network error: {0}")]
NetworkError(#[from] WsError),
@@ -14,6 +14,7 @@ pub use packet_router::{
};
pub use traits::GatewayPacketRouter;
mod bandwidth;
pub mod client;
pub mod error;
pub mod packet_router;
@@ -12,12 +12,8 @@ use log::*;
use nym_gateway_requests::registration::handshake::SharedKeys;
use nym_gateway_requests::ServerResponse;
use nym_task::TaskClient;
use si_scale::helpers::bibytes2;
use std::os::raw::c_int as RawFd;
use std::sync::atomic::{AtomicI64, Ordering};
use std::sync::Arc;
use std::time::Duration;
use time::OffsetDateTime;
use tungstenite::Message;
#[cfg(unix)]
@@ -27,6 +23,7 @@ use tokio::net::TcpStream;
#[cfg(not(target_arch = "wasm32"))]
use tokio_tungstenite::{MaybeTlsStream, WebSocketStream};
use crate::bandwidth::ClientBandwidth;
#[cfg(target_arch = "wasm32")]
use wasm_utils::websocket::JSWebsocket;
@@ -53,21 +50,6 @@ pub(crate) fn ws_fd(_conn: &WsConn) -> Option<RawFd> {
None
}
// disgusting? absolutely, but does the trick for now
static LAST_LOGGED_BANDWIDTH_TS: AtomicI64 = AtomicI64::new(0);
fn maybe_log_bandwidth(remaining: i64) {
// SAFETY: this value is always populated with valid timestamps
let last =
OffsetDateTime::from_unix_timestamp(LAST_LOGGED_BANDWIDTH_TS.load(Ordering::Relaxed))
.unwrap();
let now = OffsetDateTime::now_utc();
if last + Duration::from_secs(10) < now {
log::info!("remaining bandwidth: {}", bibytes2(remaining as f64));
LAST_LOGGED_BANDWIDTH_TS.store(now.unix_timestamp(), Ordering::Relaxed)
}
}
#[derive(Debug)]
pub(crate) struct PartiallyDelegated {
sink_half: SplitSink<WsConn, Message>,
@@ -76,10 +58,12 @@ pub(crate) struct PartiallyDelegated {
}
impl PartiallyDelegated {
// fn try_recover_plaintext(ws_message: Message, shared_key: &SharedKeys, )
fn recover_received_plaintexts(
ws_msgs: Vec<Message>,
shared_key: &SharedKeys,
bandwidth_remaining: Arc<AtomicI64>,
client_bandwidth: ClientBandwidth,
) -> Result<Vec<Vec<u8>>, GatewayClientError> {
let mut plaintexts = Vec::with_capacity(ws_msgs.len());
for ws_msg in ws_msgs {
@@ -105,15 +89,15 @@ impl PartiallyDelegated {
{
ServerResponse::Send {
remaining_bandwidth,
} => {
maybe_log_bandwidth(remaining_bandwidth);
bandwidth_remaining
.store(remaining_bandwidth, std::sync::atomic::Ordering::Release)
}
} => client_bandwidth.update_and_maybe_log(remaining_bandwidth),
ServerResponse::Error { message } => {
error!("gateway failure: {message}");
error!("[1] gateway failure: {message}");
return Err(GatewayClientError::GatewayError(message));
}
ServerResponse::TypedError { error } => {
error!("[2] gateway failure: {error}");
return Err(GatewayClientError::TypedGatewayError(error));
}
other => {
warn!(
"received illegal message of type {} in an authenticated client",
@@ -134,10 +118,9 @@ impl PartiallyDelegated {
ws_msgs: Vec<Message>,
packet_router: &PacketRouter,
shared_key: &SharedKeys,
bandwidth_remaining: Arc<AtomicI64>,
client_bandwidth: ClientBandwidth,
) -> Result<(), GatewayClientError> {
let plaintexts =
Self::recover_received_plaintexts(ws_msgs, shared_key, bandwidth_remaining)?;
let plaintexts = Self::recover_received_plaintexts(ws_msgs, shared_key, client_bandwidth)?;
packet_router.route_received(plaintexts)
}
@@ -145,7 +128,7 @@ impl PartiallyDelegated {
conn: WsConn,
mut packet_router: PacketRouter,
shared_key: Arc<SharedKeys>,
bandwidth_remaining: Arc<AtomicI64>,
client_bandwidth: ClientBandwidth,
mut shutdown: TaskClient,
) -> Self {
// when called for, it NEEDS TO yield back the stream so that we could merge it and
@@ -177,12 +160,12 @@ impl PartiallyDelegated {
Ok(msgs) => msgs
};
if let Err(err) = Self::route_socket_messages(ws_msgs, &packet_router, shared_key.as_ref(), bandwidth_remaining.clone()) {
if let Err(err) = Self::route_socket_messages(ws_msgs, &packet_router, shared_key.as_ref(), client_bandwidth.clone()) {
log::error!("Route socket messages failed: {err}");
break Err(err)
}
}
};
}
};
if match ret_err {
+19
View File
@@ -84,6 +84,14 @@ impl TryInto<String> for RegistrationHandshake {
}
}
// specific errors (that should not be nested!!) for clients to match on
#[derive(Debug, Error, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum SimpleGatewayRequestsError {
#[error("insufficient bandwidth available to process the request. required: {required}B, available: {available}B")]
OutOfBandwidth { required: i64, available: i64 },
}
#[derive(Debug, Error)]
pub enum GatewayRequestsError {
#[error("the request is too short")]
@@ -136,6 +144,10 @@ pub enum GatewayRequestsError {
#[error("the provided [v1] credential has invalid number of parameters - {0}")]
InvalidNumberOfEmbededParameters(u32),
// variant to catch legacy errors
#[error("{0}")]
Other(String),
}
#[derive(Serialize, Deserialize, Debug)]
@@ -283,9 +295,15 @@ pub enum ServerResponse {
Send {
remaining_bandwidth: i64,
},
// Generic error
Error {
message: String,
},
// Specific typed errors
// so that clients could match on this variant without doing naive string matching
TypedError {
error: SimpleGatewayRequestsError,
},
}
impl ServerResponse {
@@ -296,6 +314,7 @@ impl ServerResponse {
ServerResponse::Bandwidth { .. } => "Bandwidth".to_string(),
ServerResponse::Send { .. } => "Send".to_string(),
ServerResponse::Error { .. } => "Error".to_string(),
ServerResponse::TypedError { .. } => "TypedError".to_string(),
}
}
pub fn new_error<S: Into<String>>(msg: S) -> Self {