Merge pull request #157 from caibirdme/feat-cache-arc

feat: use Arc<T> to store data in cache

Author: caibirdme

src/logquery/labels.rs

@@ -1,3 +1,5 @@
+use std::sync::Arc;
+
 use super::*;
 use crate::{errors::AppError, state::AppState};
 use axum::{
@@ -31,7 +33,7 @@ pub async fn query_labels(
 	};
 	if should_cache {
 		let d = serde_json::to_vec(&resp).unwrap();
-		cache.insert(label_cache_key().to_string(), d);
+		cache.insert(label_cache_key().to_string(), Arc::new(d));
 	}
 	Ok(resp)
 }
@@ -96,7 +98,7 @@ pub async fn query_label_values(
 	};
 	if should_cache {
 		let d = serde_json::to_vec(&resp).unwrap();
-		cache.insert(cache_key, d);
+		cache.insert(cache_key, Arc::new(d));
 	}
 	Ok(resp)
 }
@@ -140,7 +142,7 @@ pub async fn query_series(
 		.await?;
 	if !values.is_empty() {
 		let d = serde_json::to_vec(&values).unwrap();
-		state.cache.insert(series_cache_key(), d);
+		state.cache.insert(series_cache_key(), Arc::new(d));
 	}
 	Ok(Json(QuerySeriesResponse {
 		status: ResponseStatus::Success,

src/logquery/query_range.rs

@@ -9,7 +9,7 @@ use axum_valid::Valid;
 use itertools::Itertools;
 use logql::parser;
 use moka::sync::Cache;
-use std::collections::HashMap;
+use std::{collections::HashMap, sync::Arc};
 
 pub async fn query_range(
 	State(state): State<AppState>,
@@ -31,7 +31,7 @@ pub async fn query_range(
 	};
 	if let Ok(inner) = &resp {
 		let d = serde_json::to_vec(inner).unwrap();
-		state.cache.insert(cache_key, d);
+		state.cache.insert(cache_key, Arc::new(d));
 	}
 	resp
 }
@@ -52,7 +52,7 @@ pub async fn loki_is_working() -> Result<QueryRangeResponse, AppError> {
 
 fn get_cached_query(
 	key: &str,
-	cache: Cache<String, Vec<u8>>,
+	cache: Cache<String, Arc<Vec<u8>>>,
 ) -> Option<QueryRangeResponse> {
 	if let Some(v) = cache.get(key) {
 		if let Ok(d) = serde_json::from_slice(&v) {

src/state.rs

@@ -10,6 +10,6 @@ pub struct AppState {
 	pub config: Arc<config::AppConfig>,
 	pub log_handle: Box<dyn LogStorage>,
 	pub trace_handle: Box<dyn TraceStorage>,
-	pub cache: Cache<String, Vec<u8>>,
+	pub cache: Cache<String, Arc<Vec<u8>>>,
 	pub metrics: Arc<metrics::Instrumentations>,
 }

src/trace/traceid.rs

@@ -1,3 +1,5 @@
+use std::sync::Arc;
+
 use super::*;
 use crate::{
 	errors::AppError, proto::tempopb::Trace, state::AppState,
@@ -97,21 +99,25 @@ pub async fn get_trace_by_id(
 	Ok(val)
 }
 
-fn cache_trace(trace_id: &str, trace: &Trace, cache: Cache<String, Vec<u8>>) {
+fn cache_trace(
+	trace_id: &str,
+	trace: &Trace,
+	cache: Cache<String, Arc<Vec<u8>>>,
+) {
 	let d = trace.encode_to_vec();
 	let key = get_trace_cache_key(trace_id);
-	cache.insert(key.clone(), d.clone());
+	cache.insert(key.clone(), Arc::new(d.clone()));
 }
 
 fn get_cached_trace(
 	trace_id: &str,
-	cache: Cache<String, Vec<u8>>,
+	cache: Cache<String, Arc<Vec<u8>>>,
 ) -> Result<Option<Trace>, AppError> {
 	let data = cache.get(get_trace_cache_key(trace_id).as_str());
 	match data {
 		Some(data) => {
-			let trace =
-				Message::decode(Bytes::from(data)).map_err(|e| anyhow!(e))?;
+			let s: &[u8] = data.as_ref();
+			let trace = Message::decode(s).map_err(|e| anyhow!(e))?;
 			Ok(Some(trace))
 		}
 		None => Ok(None),