diff --git a/common/bandwidth-controller/src/event.rs b/common/bandwidth-controller/src/event.rs new file mode 100644 index 0000000000..287598f959 --- /dev/null +++ b/common/bandwidth-controller/src/event.rs @@ -0,0 +1,13 @@ +// Copyright 2024 - Nym Technologies SA +// SPDX-License-Identifier: Apache-2.0 + +// See other comments for other TaskStatus message enumds about abusing the Error trait when we +// should have a new trait for TaskStatus messages +#[derive(Debug, thiserror::Error)] +pub enum BandwidthStatusMessage { + #[error("remaining bandwidth: {0}")] + RemainingBandwidth(i64), + + #[error("no bandwidth left")] + NoBandwidth, +} diff --git a/common/bandwidth-controller/src/lib.rs b/common/bandwidth-controller/src/lib.rs index c7aed707f9..91201b9a92 100644 --- a/common/bandwidth-controller/src/lib.rs +++ b/common/bandwidth-controller/src/lib.rs @@ -14,8 +14,11 @@ use nym_validator_client::coconut::all_coconut_api_clients; use nym_validator_client::nym_api::EpochId; use nym_validator_client::nyxd::contract_traits::DkgQueryClient; +pub use event::BandwidthStatusMessage; + pub mod acquire; pub mod error; +mod event; mod utils; #[derive(Debug)] diff --git a/common/client-libs/gateway-client/src/client.rs b/common/client-libs/gateway-client/src/client.rs index d8892e856f..91c320bb46 100644 --- a/common/client-libs/gateway-client/src/client.rs +++ b/common/client-libs/gateway-client/src/client.rs @@ -11,7 +11,7 @@ use crate::traits::GatewayPacketRouter; use crate::{cleanup_socket_message, try_decrypt_binary_message}; use futures::{SinkExt, StreamExt}; use log::*; -use nym_bandwidth_controller::BandwidthController; +use nym_bandwidth_controller::{BandwidthController, BandwidthStatusMessage}; use nym_credential_storage::ephemeral_storage::EphemeralStorage as EphemeralCredentialStorage; use nym_credential_storage::storage::Storage as CredentialStorage; use nym_credentials::CredentialSpendingData; @@ -105,8 +105,8 @@ pub struct GatewayClient { // currently unused (but populated) negotiated_protocol: Option, - /// Listen to shutdown messages. - shutdown: TaskClient, + /// Listen to shutdown messages and send notifications back to the task manager + task_client: TaskClient, } impl GatewayClient { @@ -117,7 +117,7 @@ impl GatewayClient { shared_key: Option>, packet_router: PacketRouter, bandwidth_controller: Option>, - shutdown: TaskClient, + task_client: TaskClient, ) -> Self { GatewayClient { authenticated: false, @@ -135,7 +135,7 @@ impl GatewayClient { reconnection_attempts: DEFAULT_RECONNECTION_ATTEMPTS, reconnection_backoff: DEFAULT_RECONNECTION_BACKOFF, negotiated_protocol: None, - shutdown, + task_client, } } @@ -299,7 +299,7 @@ impl GatewayClient { loop { tokio::select! { - _ = self.shutdown.recv() => { + _ = self.task_client.recv() => { log::trace!("GatewayClient control response: Received shutdown"); log::debug!("GatewayClient control response: Exiting"); break Err(GatewayClientError::ConnectionClosedGatewayShutdown); @@ -540,6 +540,9 @@ impl GatewayClient { self.bandwidth_remaining = bandwidth_remaining; self.negotiated_protocol = protocol_version; log::debug!("authenticated: {status}, bandwidth remaining: {bandwidth_remaining}"); + self.task_client.send_status_msg(Box::new( + BandwidthStatusMessage::RemainingBandwidth(bandwidth_remaining), + )); Ok(()) } ServerResponse::Error { message } => Err(GatewayClientError::GatewayError(message)), @@ -805,7 +808,7 @@ impl GatewayClient { .as_ref() .expect("no shared key present even though we're authenticated!"), ), - self.shutdown.clone(), + self.task_client.clone(), ) } _ => unreachable!(), @@ -879,8 +882,8 @@ impl GatewayClient { // perfectly fine here, because it's not meant to be used let (ack_tx, _) = mpsc::unbounded(); let (mix_tx, _) = mpsc::unbounded(); - let shutdown = TaskClient::dummy(); - let packet_router = PacketRouter::new(ack_tx, mix_tx, shutdown.clone()); + let task_client = TaskClient::dummy(); + let packet_router = PacketRouter::new(ack_tx, mix_tx, task_client.clone()); GatewayClient { authenticated: false, @@ -898,7 +901,7 @@ impl GatewayClient { reconnection_attempts: DEFAULT_RECONNECTION_ATTEMPTS, reconnection_backoff: DEFAULT_RECONNECTION_BACKOFF, negotiated_protocol: None, - shutdown, + task_client, } } @@ -906,7 +909,7 @@ impl GatewayClient { self, packet_router: PacketRouter, bandwidth_controller: Option>, - shutdown: TaskClient, + task_client: TaskClient, ) -> GatewayClient { // invariants that can't be broken // (unless somebody decided to expose some field that wasn't meant to be exposed) @@ -930,7 +933,7 @@ impl GatewayClient { reconnection_attempts: self.reconnection_attempts, reconnection_backoff: self.reconnection_backoff, negotiated_protocol: self.negotiated_protocol, - shutdown, + task_client, } } }