From bf13e175cb72b8f4d94c5f4809cd7da52e8f8ec9 Mon Sep 17 00:00:00 2001 From: pranc1ngpegasus Date: Tue, 15 Oct 2024 22:26:35 +0900 Subject: [PATCH 1/3] feat: implement size-limited JSON requests for save metrics --- agent/src/services/studio.rs | 88 ++++++++++++++++++++++-------------- 1 file changed, 53 insertions(+), 35 deletions(-) diff --git a/agent/src/services/studio.rs b/agent/src/services/studio.rs index 3e80a7da..90c6876c 100644 --- a/agent/src/services/studio.rs +++ b/agent/src/services/studio.rs @@ -24,6 +24,10 @@ use std::collections::VecDeque; use serde::{Deserialize, Serialize}; use serde_json::{json, Value}; +// The maximum JSON body size is actually 1MB +// We reserve 100KB as a buffer for Verifiable Credential capacity +const JSON_BODY_MAX_SIZE: usize = 900_000; + #[derive(Deserialize)] pub struct EmptyResponse {} @@ -423,46 +427,60 @@ struct MetricsWithTimestampStr { impl MetricStoreRepository for Studio { async fn save(&self, request: VecDeque) -> anyhow::Result<()> { - let metrics_str = request - .into_iter() - .map(|m| MetricsWithTimestampStr { - timestamp: m.timestamp, - metrics: m - .metrics - .into_iter() - .map(|metric| MetricStr { - metric_type: metric.metric_type.to_string(), - value: metric.value, - }) - .collect::>(), - }) - .collect::>(); - - let my_did = self.did_accessor.get_my_did(); - let my_keyring = self.did_accessor.get_my_keyring(); + let mut metrics = request; + while !metrics.is_empty() { + let my_did = self.did_accessor.get_my_did(); + let my_keyring = self.did_accessor.get_my_keyring(); + let mut metrics_str = Vec::new(); + let mut current_size = 0; + + while let Some(m) = metrics.pop_front() { + let metrics_with_timestamp_str = MetricsWithTimestampStr { + timestamp: m.timestamp, + metrics: m + .metrics + .iter() + .map(|metric| MetricStr { + metric_type: metric.metric_type.to_string(), + value: metric.value, + }) + .collect::>(), + }; + + let item_size = serde_json::to_string(&metrics_with_timestamp_str)?.len(); + if current_size + item_size > JSON_BODY_MAX_SIZE { + metrics.push_front(m); + break; + } + current_size += item_size; + metrics_str.push(metrics_with_timestamp_str); + } - let model = VerifiableCredentials::new(my_did, json!(metrics_str), chrono::Utc::now()); - let payload = DidVcService::generate(&self.did_repository, model, &my_keyring) - .context("failed to generate payload")?; + let model = VerifiableCredentials::new(my_did, json!(metrics_str), chrono::Utc::now()); + let payload = DidVcService::generate(&self.did_repository, model, &my_keyring) + .context("failed to generate payload")?; - let payload = serde_json::to_string(&payload).context("failed to serialize")?; - let res = self.http_client.post("/v1/metrics", &payload).await?; + let payload = serde_json::to_string(&payload).context("failed to serialize")?; + let res = self.http_client.post("/v1/metrics", &payload).await?; - let status = res.status(); - let json: Value = res.json().await.context("Failed to read response body")?; - let message = if let Some(message) = json.get("message").map(|v| v.to_string()) { - message - } else { - "".to_string() - }; - match status { - reqwest::StatusCode::OK => Ok(()), - reqwest::StatusCode::NOT_FOUND => anyhow::bail!("StatusCode=404, {}", message), - reqwest::StatusCode::INTERNAL_SERVER_ERROR => { - anyhow::bail!("StatusCode=500, {}", message); + let status = res.status(); + let json: Value = res.json().await.context("Failed to read response body")?; + let message = if let Some(message) = json.get("message").map(|v| v.to_string()) { + message + } else { + "".to_string() + }; + match status { + reqwest::StatusCode::OK => continue, + reqwest::StatusCode::NOT_FOUND => anyhow::bail!("StatusCode=404, {}", message), + reqwest::StatusCode::INTERNAL_SERVER_ERROR => { + anyhow::bail!("StatusCode=500, {}", message); + } + other => anyhow::bail!("StatusCode={other}, {}", message), } - other => anyhow::bail!("StatusCode={other}, {}", message), } + + Ok(()) } } From 4fb86fd495c4b7bafb8553b65d2b5699b6d87f0c Mon Sep 17 00:00:00 2001 From: pranc1ngpegasus Date: Fri, 18 Oct 2024 00:03:37 +0900 Subject: [PATCH 2/3] fix: simplify --- agent/src/repository/metric_repository.rs | 1 + agent/src/services/studio.rs | 28 ++--------------------- 2 files changed, 3 insertions(+), 26 deletions(-) diff --git a/agent/src/repository/metric_repository.rs b/agent/src/repository/metric_repository.rs index 6dbc1208..6eb03d3e 100644 --- a/agent/src/repository/metric_repository.rs +++ b/agent/src/repository/metric_repository.rs @@ -34,6 +34,7 @@ pub trait MetricStoreRepository { } #[derive(Debug, Serialize, Clone, PartialEq)] +#[serde(rename_all = "snake_case")] pub enum MetricType { CpuUsage, MemoryUsage, diff --git a/agent/src/services/studio.rs b/agent/src/services/studio.rs index 90c6876c..a0696ca2 100644 --- a/agent/src/services/studio.rs +++ b/agent/src/services/studio.rs @@ -413,18 +413,6 @@ impl MessageActivityRepository for Studio { } } -#[derive(Serialize)] -struct MetricStr { - metric_type: String, - value: f32, -} - -#[derive(Serialize)] -struct MetricsWithTimestampStr { - timestamp: chrono::DateTime, - metrics: Vec, -} - impl MetricStoreRepository for Studio { async fn save(&self, request: VecDeque) -> anyhow::Result<()> { let mut metrics = request; @@ -435,25 +423,13 @@ impl MetricStoreRepository for Studio { let mut current_size = 0; while let Some(m) = metrics.pop_front() { - let metrics_with_timestamp_str = MetricsWithTimestampStr { - timestamp: m.timestamp, - metrics: m - .metrics - .iter() - .map(|metric| MetricStr { - metric_type: metric.metric_type.to_string(), - value: metric.value, - }) - .collect::>(), - }; - - let item_size = serde_json::to_string(&metrics_with_timestamp_str)?.len(); + let item_size = serde_json::to_string(&m)?.len(); if current_size + item_size > JSON_BODY_MAX_SIZE { metrics.push_front(m); break; } current_size += item_size; - metrics_str.push(metrics_with_timestamp_str); + metrics_str.push(m); } let model = VerifiableCredentials::new(my_did, json!(metrics_str), chrono::Utc::now()); From 1952571883f561aca0cc4a0ee2a96737d2409efc Mon Sep 17 00:00:00 2001 From: pranc1ngpegasus Date: Sat, 19 Oct 2024 22:45:40 +0900 Subject: [PATCH 3/3] fix: add guard clause --- agent/src/services/studio.rs | 3 +++ 1 file changed, 3 insertions(+) diff --git a/agent/src/services/studio.rs b/agent/src/services/studio.rs index a0696ca2..a0a3abbf 100644 --- a/agent/src/services/studio.rs +++ b/agent/src/services/studio.rs @@ -424,6 +424,9 @@ impl MetricStoreRepository for Studio { while let Some(m) = metrics.pop_front() { let item_size = serde_json::to_string(&m)?.len(); + if item_size > JSON_BODY_MAX_SIZE { + anyhow::bail!("invalid item size: JSON body size too large") + } if current_size + item_size > JSON_BODY_MAX_SIZE { metrics.push_front(m); break;