use std::collections::HashMap;
use std::future::Future;
use std::hash::{Hash, Hasher};
use std::pin::Pin;
use std::sync::Arc;
use std::time::{Duration, Instant};
use http::{Method, Uri};
use parking_lot::Mutex;
use tokio::sync::broadcast;
use crate::error::{Error, Result};
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct RequestKey {
method: String,
uri: String,
body_hash: Option<u64>,
}
impl RequestKey {
pub fn new(method: &Method, uri: &Uri) -> Self {
Self {
method: method.to_string(),
uri: uri.to_string(),
body_hash: None,
}
}
pub fn with_body(method: &Method, uri: &Uri, body: &[u8]) -> Self {
let mut hasher = std::collections::hash_map::DefaultHasher::new();
body.hash(&mut hasher);
Self {
method: method.to_string(),
uri: uri.to_string(),
body_hash: Some(hasher.finish()),
}
}
}
#[derive(Debug, Clone)]
pub struct SharedResult {
pub status: u16,
pub headers: Vec<(String, String)>,
pub body: Vec<u8>,
}
#[derive(Debug)]
struct InFlight {
started: Instant,
sender: broadcast::Sender<SharedResult>,
}
#[derive(Debug, Clone)]
pub struct DedupConfig {
pub enabled: bool,
pub max_wait: Duration,
pub max_waiters: usize,
pub methods: Vec<Method>,
}
impl Default for DedupConfig {
fn default() -> Self {
Self {
enabled: true,
max_wait: Duration::from_secs(30),
max_waiters: 100,
methods: vec![Method::GET, Method::HEAD],
}
}
}
impl DedupConfig {
pub fn new() -> Self {
Self::default()
}
#[must_use]
pub fn disabled() -> Self {
Self {
enabled: false,
..Default::default()
}
}
#[must_use]
pub fn with_max_wait(mut self, duration: Duration) -> Self {
self.max_wait = duration;
self
}
#[must_use]
pub fn with_max_waiters(mut self, max: usize) -> Self {
self.max_waiters = max;
self
}
#[must_use]
pub fn with_method(mut self, method: Method) -> Self {
if !self.methods.contains(&method) {
self.methods.push(method);
}
self
}
pub fn should_dedup(&self, method: &Method) -> bool {
self.enabled && self.methods.contains(method)
}
}
#[derive(Debug)]
pub struct Deduplicator {
config: DedupConfig,
in_flight: Arc<Mutex<HashMap<RequestKey, InFlight>>>,
}
impl Default for Deduplicator {
fn default() -> Self {
Self::new(DedupConfig::default())
}
}
impl Deduplicator {
pub fn new(config: DedupConfig) -> Self {
Self {
config,
in_flight: Arc::new(Mutex::new(HashMap::new())),
}
}
pub fn try_coalesce(
&self,
key: &RequestKey,
) -> Result<Option<broadcast::Receiver<SharedResult>>> {
if !self.config.enabled {
return Ok(None);
}
let mut in_flight = self.in_flight.lock();
if let Some(flight) = in_flight.get(key) {
if flight.started.elapsed() < self.config.max_wait {
if flight.sender.receiver_count() >= self.config.max_waiters {
return Err(Error::Circuit {
message: "Too many waiters for coalesced request".to_string(),
source: None,
});
}
return Ok(Some(flight.sender.subscribe()));
}
in_flight.remove(key);
}
Ok(None)
}
pub fn register(&self, key: RequestKey) -> broadcast::Sender<SharedResult> {
let (sender, _) = broadcast::channel(1);
let flight = InFlight {
started: Instant::now(),
sender: sender.clone(),
};
self.in_flight.lock().insert(key, flight);
sender
}
pub fn complete(&self, key: &RequestKey, result: SharedResult) {
let mut in_flight = self.in_flight.lock();
if let Some(flight) = in_flight.remove(key) {
let _ = flight.sender.send(result);
}
}
pub fn cancel(&self, key: &RequestKey) {
self.in_flight.lock().remove(key);
}
pub fn in_flight_count(&self) -> usize {
self.in_flight.lock().len()
}
pub fn cleanup(&self) {
let mut in_flight = self.in_flight.lock();
in_flight.retain(|_, flight| flight.started.elapsed() < self.config.max_wait);
}
pub async fn execute<F, Fut>(&self, key: RequestKey, make_request: F) -> Result<SharedResult>
where
F: FnOnce() -> Fut,
Fut: Future<Output = Result<SharedResult>>,
{
if let Some(mut receiver) = self.try_coalesce(&key)? {
return receiver.recv().await.map_err(|e| Error::Circuit {
message: format!("Coalesced request was cancelled: {}", e),
source: None,
});
}
let _sender = self.register(key.clone());
let result = make_request().await;
match result {
Ok(shared) => {
self.complete(&key, shared.clone());
Ok(shared)
}
Err(e) => {
self.cancel(&key);
Err(e)
}
}
}
}
pub type CoalescedFuture<'a> = Pin<Box<dyn Future<Output = Result<SharedResult>> + Send + 'a>>;
#[cfg(test)]
mod tests {
#![allow(clippy::unwrap_used)]
use super::*;
#[test]
fn test_request_key() {
let key1 = RequestKey::new(&Method::GET, &"/api/test".parse().unwrap());
let key2 = RequestKey::new(&Method::GET, &"/api/test".parse().unwrap());
let key3 = RequestKey::new(&Method::POST, &"/api/test".parse().unwrap());
assert_eq!(key1, key2);
assert_ne!(key1, key3);
}
#[test]
fn test_request_key_with_body() {
let body = b"test body";
let key1 = RequestKey::with_body(&Method::POST, &"/api".parse().unwrap(), body);
let key2 = RequestKey::with_body(&Method::POST, &"/api".parse().unwrap(), body);
let key3 = RequestKey::with_body(&Method::POST, &"/api".parse().unwrap(), b"other");
assert_eq!(key1, key2);
assert_ne!(key1, key3);
}
#[test]
fn test_dedup_config() {
let config = DedupConfig::new()
.with_max_wait(Duration::from_secs(60))
.with_max_waiters(50)
.with_method(Method::POST);
assert!(config.should_dedup(&Method::GET));
assert!(config.should_dedup(&Method::POST));
assert!(!config.should_dedup(&Method::PUT));
}
#[test]
fn test_deduplicator_register() {
let dedup = Deduplicator::default();
let key = RequestKey::new(&Method::GET, &"/test".parse().unwrap());
assert_eq!(dedup.in_flight_count(), 0);
let _sender = dedup.register(key.clone());
assert_eq!(dedup.in_flight_count(), 1);
dedup.cancel(&key);
assert_eq!(dedup.in_flight_count(), 0);
}
#[tokio::test]
async fn test_deduplicator_coalesce() {
let dedup = Deduplicator::default();
let key = RequestKey::new(&Method::GET, &"/test".parse().unwrap());
assert!(dedup.try_coalesce(&key).unwrap().is_none());
let sender = dedup.register(key.clone());
let receiver = dedup.try_coalesce(&key).unwrap();
assert!(receiver.is_some());
let result = SharedResult {
status: 200,
headers: vec![],
body: b"test".to_vec(),
};
let _ = sender.send(result.clone());
let mut rx = receiver.unwrap();
let received = rx.recv().await.unwrap();
assert_eq!(received.status, 200);
}
}