From f939cae3d9bed0d4d85b526e8003f6bf254e4aa0 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Bogdan-=C8=98tefan=20Neac=C5=9Fu?= <bogdan@nymtech.net>
Date: Tue, 6 Aug 2024 18:14:29 +0200
Subject: [PATCH] Update peer refresh value (#4754)

* Use a more proper timeout value

* Move const to wireguard types
---
 common/wireguard-types/src/lib.rs       |  7 +++++++
 common/wireguard/src/peer_controller.rs | 19 +++++++++++++------
 2 files changed, 20 insertions(+), 6 deletions(-)

diff --git a/common/wireguard-types/src/lib.rs b/common/wireguard-types/src/lib.rs
index 18a01cd9e7..bdb8b575b7 100644
--- a/common/wireguard-types/src/lib.rs
+++ b/common/wireguard-types/src/lib.rs
@@ -6,10 +6,17 @@ pub mod error;
 pub mod public_key;
 pub mod registration;
 
+use std::time::Duration;
+
 pub use config::Config;
 pub use error::Error;
 pub use public_key::PeerPublicKey;
 pub use registration::{ClientMac, ClientMessage, GatewayClient, InitMessage, Nonce};
 
+// To avoid any problems, keep this stale check time bigger (>2x) then the bandwidth cap
+// reset time (currently that one is 24h, at UTC midnight)
+pub const DEFAULT_PEER_TIMEOUT: Duration = Duration::from_secs(60 * 60 * 24 * 3); // 3 days
+pub const DEFAULT_PEER_TIMEOUT_CHECK: Duration = Duration::from_secs(5); // 5 seconds
+
 #[cfg(feature = "verify")]
 pub use registration::HmacSha256;
diff --git a/common/wireguard/src/peer_controller.rs b/common/wireguard/src/peer_controller.rs
index 7a48a20857..9097f155ac 100644
--- a/common/wireguard/src/peer_controller.rs
+++ b/common/wireguard/src/peer_controller.rs
@@ -5,19 +5,15 @@ use chrono::{Timelike, Utc};
 use defguard_wireguard_rs::{host::Peer, key::Key, WireguardInterfaceApi};
 use nym_gateway_storage::Storage;
 use nym_wireguard_types::registration::{RemainingBandwidthData, BANDWIDTH_CAP_PER_DAY};
+use nym_wireguard_types::{DEFAULT_PEER_TIMEOUT, DEFAULT_PEER_TIMEOUT_CHECK};
 use std::time::SystemTime;
-use std::{collections::HashMap, sync::Arc, time::Duration};
+use std::{collections::HashMap, sync::Arc};
 use tokio::sync::mpsc;
 use tokio_stream::{wrappers::IntervalStream, StreamExt};
 
 use crate::error::Error;
 use crate::WgApiWrapper;
 
-// To avoid any problems, keep this stale check time bigger (>2x) then the bandwidth cap
-// reset time (currently that one is 24h, at UTC midnight)
-const DEFAULT_PEER_TIMEOUT: Duration = Duration::from_secs(60 * 60 * 24 * 3); // 3 days
-const DEFAULT_PEER_TIMEOUT_CHECK: Duration = Duration::from_secs(60); // 1 minute
-
 pub enum PeerControlRequest {
     AddPeer(Peer),
     RemovePeer(Key),
@@ -50,6 +46,7 @@ pub struct PeerController<St: Storage> {
     active_peers: HashMap<Key, Peer>,
     suspended_peers: HashMap<Key, Peer>,
     last_seen_bandwidth: HashMap<Key, u64>,
+    timeout_count: u8,
 }
 
 impl<St: Storage> PeerController<St> {
@@ -82,6 +79,7 @@ impl<St: Storage> PeerController<St> {
             active_peers,
             suspended_peers,
             last_seen_bandwidth: HashMap::new(),
+            timeout_count: 0,
         }
     }
 
@@ -144,6 +142,15 @@ impl<St: Storage> PeerController<St> {
             .iter()
             .map(|(key, peer)| (key.clone(), peer.rx_bytes + peer.tx_bytes))
             .collect();
+
+        // Do in-memory updates of bandwidth every DEFAULT_PEER_TIMEOUT_CHECK
+        // and storage updates every 5 * DEFAULT_PEER_TIMEOUT_CHECK, because in-memory
+        // is more important for client query preciseness
+        self.timeout_count = self.timeout_count % 5 + 1;
+        if !reset && self.timeout_count < 5 {
+            return Ok(());
+        }
+
         if reset {
             self.active_peers = host.peers;
             for peer in self.active_peers.values() {