From 72bab9567d16c425640169b0119a2f3dccf05853 Mon Sep 17 00:00:00 2001
From: William Smith <williamprincesmith@gmail.com>
Date: Fri, 21 Jun 2024 16:47:05 -0700
Subject: [PATCH] count everything as spam on rpc nodes

---
 crates/sui-core/src/authority_server.rs       |  9 +-------
 crates/sui-core/src/traffic_controller/mod.rs | 12 ++--------
 crates/sui-json-rpc/src/axum_router.rs        | 22 ++++++++-----------
 crates/sui-types/src/traffic_control.rs       |  7 +++++-
 4 files changed, 18 insertions(+), 32 deletions(-)

diff --git a/crates/sui-core/src/authority_server.rs b/crates/sui-core/src/authority_server.rs
index 578108ab67150..d5ba65c35c692 100644
--- a/crates/sui-core/src/authority_server.rs
+++ b/crates/sui-core/src/authority_server.rs
@@ -984,14 +984,7 @@ fn normalize(err: SuiError) -> Weight {
 macro_rules! handle_with_decoration {
     ($self:ident, $func_name:ident, $request:ident) => {{
         if $self.client_id_source.is_none() {
-            return match $self.$func_name($request).await {
-                Ok((result, _)) => {
-                    Ok(result)
-                }
-                Err(err) => {
-                    Err(err)
-                }
-            }
+            return $self.$func_name($request).await.map(|(result, _)| result);
         }
 
         let client = match $self.client_id_source.as_ref().unwrap() {
diff --git a/crates/sui-core/src/traffic_controller/mod.rs b/crates/sui-core/src/traffic_controller/mod.rs
index ff186670e5fbf..e10f3b3e5928a 100644
--- a/crates/sui-core/src/traffic_controller/mod.rs
+++ b/crates/sui-core/src/traffic_controller/mod.rs
@@ -296,7 +296,7 @@ async fn handle_error_tally(
     metrics: Arc<TrafficControllerMetrics>,
     mem_drainfile_present: bool,
 ) -> Result<(), reqwest::Error> {
-    if !tally.error_weight.is_sampled().await {
+    if !tally.error_weight.is_sampled() {
         return Ok(());
     }
     let resp = policy.handle_tally(tally.clone());
@@ -330,15 +330,7 @@ async fn handle_spam_tally(
     metrics: Arc<TrafficControllerMetrics>,
     mem_drainfile_present: bool,
 ) -> Result<(), reqwest::Error> {
-    let sampled = futures::future::join_all(vec![
-        tally.spam_weight.is_sampled(),
-        policy_config.spam_sample_rate.is_sampled(),
-    ])
-    .await
-    .into_iter()
-    .all(|sampled| sampled);
-
-    if !sampled {
+    if !(tally.spam_weight.is_sampled() && policy_config.spam_sample_rate.is_sampled()) {
         return Ok(());
     }
     let resp = policy.handle_tally(tally.clone());
diff --git a/crates/sui-json-rpc/src/axum_router.rs b/crates/sui-json-rpc/src/axum_router.rs
index c724ae1f4c8d0..f7ef9d7693b67 100644
--- a/crates/sui-json-rpc/src/axum_router.rs
+++ b/crates/sui-json-rpc/src/axum_router.rs
@@ -167,10 +167,9 @@ async fn process_raw_request<L: Logger>(
         }
 
         // handle response tallying
-        let method = request.method.to_string();
         let response = process_request(request, api_version, service.call_data()).await;
         if let Some(traffic_controller) = &service.traffic_controller {
-            handle_traffic_resp(traffic_controller.clone(), client, &response, method);
+            handle_traffic_resp(traffic_controller.clone(), client, &response);
         }
 
         response
@@ -199,27 +198,24 @@ async fn handle_traffic_req(
     }
 }
 
-fn spam_weight_for_method(method: String) -> Weight {
-    // unless request requires gas payment, count against
-    // spam tally
-    match method.as_str() {
-        "sui_executeTransactionBlock" | "sui_devInspectTransactionBlock" => Weight::zero(),
-        _ => Weight::one(),
-    }
-}
-
 fn handle_traffic_resp(
     traffic_controller: Arc<TrafficController>,
     client: Option<IpAddr>,
     response: &MethodResponse,
-    method: String,
 ) {
     let error = response.error_code.map(ErrorCode::from);
     traffic_controller.tally(TrafficTally {
         direct: client,
         through_fullnode: None,
         error_weight: error.map(normalize).unwrap_or(Weight::zero()),
-        spam_weight: spam_weight_for_method(method),
+        // For now, count everything as spam with equal weight
+        // on the rpc node side, including gas-charging endpoints
+        // such as `sui_executeTransactionBlock`, as this can enable
+        // node operators who wish to rate limit their transcation
+        // traffic and incentivize high volume clients to choose a
+        // suitable rpc provider (or run their own). Later we may want
+        // to provide a weight distribution based on the method being called.
+        spam_weight: Weight::one(),
         timestamp: SystemTime::now(),
     });
 }
diff --git a/crates/sui-types/src/traffic_control.rs b/crates/sui-types/src/traffic_control.rs
index c5d5d41116f0a..e976081abcecd 100644
--- a/crates/sui-types/src/traffic_control.rs
+++ b/crates/sui-types/src/traffic_control.rs
@@ -53,7 +53,7 @@ impl Weight {
         self.0
     }
 
-    pub async fn is_sampled(&self) -> bool {
+    pub fn is_sampled(&self) -> bool {
         let mut rng = rand::thread_rng();
         let sample = rand::distributions::Uniform::new(0.0, 1.0).sample(&mut rng);
         sample <= self.value()
@@ -201,6 +201,11 @@ pub struct PolicyConfig {
     #[serde(default = "default_channel_capacity")]
     pub channel_capacity: usize,
     #[serde(default = "default_spam_sample_rate")]
+    /// Note that this sample policy is applied on top of the
+    /// endpoint-specific sample policy (not configurable) which
+    /// weighs endpoints by the relative effort required to serve
+    /// them. Therefore a sample rate of N will yield an actual
+    /// sample rate <= N.
     pub spam_sample_rate: Weight,
     #[serde(default = "default_dry_run")]
     pub dry_run: bool,