vantage_api_pool/matcher/
mod.rs

1use std::{
2    collections::HashMap,
3    sync::{atomic::AtomicUsize, Arc},
4};
5
6use reqwest::{Request, Response};
7use tokio::sync::{mpsc, oneshot, Mutex};
8use tracing::{warn, Instrument as _};
9
10use crate::EventualRequest;
11
12pub struct EventualRequestMatcher<T: Send + Sync + Sized> {
13    thread_handle: tokio::task::JoinHandle<()>,
14    request_sender: mpsc::Sender<EventualRequest<T>>,
15    pending_requests: Arc<Mutex<HashMap<usize, oneshot::Sender<EventualRequest<T>>>>>,
16    seq_id: AtomicUsize,
17}
18
19impl<T: Send + Sync + Sized + 'static> EventualRequestMatcher<T> {
20    pub fn new(
21        request_sender: mpsc::Sender<EventualRequest<T>>,
22        response_receiver: mpsc::Receiver<EventualRequest<T>>,
23    ) -> Self {
24        let pending_requests = Arc::new(Mutex::new(HashMap::new()));
25        let thread_handle = tokio::spawn(
26            Self::worker_thread(pending_requests.clone(), response_receiver).in_current_span(),
27        );
28
29        Self {
30            request_sender,
31            thread_handle,
32            pending_requests,
33            seq_id: AtomicUsize::new(0),
34        }
35    }
36
37    pub fn seq_id(&self) -> usize {
38        self.seq_id
39            .fetch_add(1, std::sync::atomic::Ordering::Relaxed)
40    }
41
42    pub async fn worker_thread(
43        ch: Arc<Mutex<HashMap<usize, oneshot::Sender<EventualRequest<T>>>>>,
44        mut response_receiver: mpsc::Receiver<EventualRequest<T>>,
45    ) {
46        loop {
47            let Some(response) = response_receiver.recv().await else {
48                warn!("response channel closed, matcher shutting down");
49                break;
50            };
51
52            match ch.lock().await.remove(&response.get_id().unwrap()) {
53                Some(sender) => sender.send(response).unwrap_or_else(|response| {
54                    warn!(id = ?response.get_id(), "failed to send response (caller dropped receiver)");
55                }),
56                None => warn!(id = ?response.get_id(), "no pending request found for response"),
57            }
58        }
59    }
60
61    pub async fn send(
62        &self,
63        request: impl Into<Request>,
64        metadata: Option<T>,
65    ) -> Result<Response, String> {
66        let mut request = EventualRequest::new(request, metadata);
67        let receiver = request
68            .register(self.seq_id(), self.pending_requests.clone())
69            .await;
70
71        self.request_sender
72            .send(request)
73            .await
74            .map_err(|e| e.to_string())?;
75
76        let mut result = receiver.await.map_err(|e| e.to_string())?;
77
78        result.response().ok_or("No response".to_string())
79    }
80
81    pub async fn shutdown(self) -> Result<(), tokio::task::JoinError> {
82        // Close channel, this will cause worker to finish
83        drop(self.request_sender);
84
85        // wait for worker to finish
86        self.thread_handle.await
87    }
88}
vantage_api_pool/matcher/mod.rs

vantage_api_pool/matcher/
mod.rs
