Use OnDemandHashMap for servername_ratelimiter

This way, semaphores are actually cleaned up eventually.
This commit is contained in:
Lambda 2024-05-27 19:27:08 +00:00
parent 2c19abc535
commit 07b5233980
2 changed files with 8 additions and 22 deletions

View file

@ -34,7 +34,7 @@ use trust_dns_resolver::TokioAsyncResolver;
use crate::{ use crate::{
api::server_server::FedDest, observability::FilterReloadHandles, services, api::server_server::FedDest, observability::FilterReloadHandles, services,
Config, Error, Result, utils::on_demand_hashmap::OnDemandHashMap, Config, Error, Result,
}; };
type WellKnownMap = HashMap<OwnedServerName, (FedDest, String)>; type WellKnownMap = HashMap<OwnedServerName, (FedDest, String)>;
@ -66,7 +66,7 @@ pub(crate) struct Service {
pub(crate) bad_query_ratelimiter: pub(crate) bad_query_ratelimiter:
Arc<RwLock<HashMap<OwnedServerName, RateLimitState>>>, Arc<RwLock<HashMap<OwnedServerName, RateLimitState>>>,
pub(crate) servername_ratelimiter: pub(crate) servername_ratelimiter:
Arc<RwLock<HashMap<OwnedServerName, Arc<Semaphore>>>>, OnDemandHashMap<OwnedServerName, Semaphore>,
pub(crate) roomid_mutex_insert: pub(crate) roomid_mutex_insert:
RwLock<HashMap<OwnedRoomId, Arc<Mutex<()>>>>, RwLock<HashMap<OwnedRoomId, Arc<Mutex<()>>>>,
pub(crate) roomid_mutex_state: RwLock<HashMap<OwnedRoomId, Arc<Mutex<()>>>>, pub(crate) roomid_mutex_state: RwLock<HashMap<OwnedRoomId, Arc<Mutex<()>>>>,
@ -263,7 +263,9 @@ impl Service {
bad_event_ratelimiter: Arc::new(RwLock::new(HashMap::new())), bad_event_ratelimiter: Arc::new(RwLock::new(HashMap::new())),
bad_signature_ratelimiter: Arc::new(RwLock::new(HashMap::new())), bad_signature_ratelimiter: Arc::new(RwLock::new(HashMap::new())),
bad_query_ratelimiter: Arc::new(RwLock::new(HashMap::new())), bad_query_ratelimiter: Arc::new(RwLock::new(HashMap::new())),
servername_ratelimiter: Arc::new(RwLock::new(HashMap::new())), servername_ratelimiter: OnDemandHashMap::new(
"servername_ratelimiter".to_owned(),
),
roomid_mutex_state: RwLock::new(HashMap::new()), roomid_mutex_state: RwLock::new(HashMap::new()),
roomid_mutex_insert: RwLock::new(HashMap::new()), roomid_mutex_insert: RwLock::new(HashMap::new()),
roomid_mutex_federation: RwLock::new(HashMap::new()), roomid_mutex_federation: RwLock::new(HashMap::new()),

View file

@ -1935,25 +1935,9 @@ impl Service {
let permit = services() let permit = services()
.globals .globals
.servername_ratelimiter .servername_ratelimiter
.read() .get_or_insert_with(origin.to_owned(), || Semaphore::new(1))
.await .await;
.get(origin) let permit = permit.acquire().await;
.map(|s| Arc::clone(s).acquire_owned());
let permit = if let Some(p) = permit {
p
} else {
let mut write =
services().globals.servername_ratelimiter.write().await;
let s = Arc::clone(
write
.entry(origin.to_owned())
.or_insert_with(|| Arc::new(Semaphore::new(1))),
);
s.acquire_owned()
}
.await;
let back_off = |id| async { let back_off = |id| async {
match services() match services()