Skip to content

fix(guard): use moka cache for in flight reqs #3629

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
MasterPtato wants to merge 1 commit into from
+6 −8
Closed

fix(guard): use moka cache for in flight reqs #3629

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
14 changes: 6 additions & 8 deletions engine/packages/guard-core/src/proxy_service.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,7 @@ use serde_json;
use rivet_runner_protocol as protocol;
use std::{
borrow::Cow,
collections::{HashMap as StdHashMap, HashSet},
collections::HashMap as StdHashMap,
net::SocketAddr,
sync::Arc,
time::{Duration, Instant},
Expand Down Expand Up @@ -350,7 +350,7 @@ pub struct ProxyState {
route_cache: RouteCache,
rate_limiters: Cache<(Id, std::net::IpAddr), Arc<Mutex<RateLimiter>>>,
in_flight_counters: Cache<(Id, std::net::IpAddr), Arc<Mutex<InFlightCounter>>>,
in_flight_requests: Arc<Mutex<HashSet<protocol::RequestId>>>,
in_flight_requests: Cache<protocol::RequestId, ()>,
port_type: PortType,
clickhouse_inserter: Option<clickhouse_inserter::ClickHouseInserterHandle>,
tasks: Arc<TaskGroup>,
Expand Down Expand Up @@ -379,7 +379,7 @@ impl ProxyState {
.max_capacity(10_000)
.time_to_live(PROXY_STATE_CACHE_TTL)
.build(),
in_flight_requests: Arc::new(Mutex::new(HashSet::new())),
in_flight_requests: Cache::builder().max_capacity(10_000_000).build(),
port_type,
clickhouse_inserter,
tasks: TaskGroup::new(),
Expand Down Expand Up @@ -660,22 +660,20 @@ impl ProxyState {
}

// Release request ID
let mut requests = self.in_flight_requests.lock().await;
requests.remove(&request_id);
self.in_flight_requests.invalidate(&request_id).await;
}

/// Generate a unique request ID that is not currently in flight
async fn generate_unique_request_id(&self) -> Result<protocol::RequestId> {
const MAX_TRIES: u32 = 100;
let mut requests = self.in_flight_requests.lock().await;

for attempt in 0..MAX_TRIES {
let request_id = protocol::util::generate_request_id();

// Check if this ID is already in use
if !requests.contains(&request_id) {
if self.in_flight_requests.get(&request_id).await.is_none() {
// Insert the ID and return it
requests.insert(request_id);
self.in_flight_requests.insert(request_id, ()).await;
return Ok(request_id);
}

Expand Down
Loading