polyfill_rs/
connection_manager.rs

1//! Connection management for maintaining warm HTTP connections
2//!
3//! This module provides functionality to keep connections alive and prevent
4//! connection drops that cause 200ms+ reconnection overhead.
5
6use reqwest::Client;
7use std::sync::Arc;
8use std::sync::atomic::{AtomicBool, Ordering};
9use std::time::Duration;
10use tokio::sync::Mutex;
11use tokio::task::JoinHandle;
12
13/// Connection keep-alive manager
14pub struct ConnectionManager {
15    client: Client,
16    base_url: String,
17    running: Arc<AtomicBool>,
18    handle: Arc<Mutex<Option<JoinHandle<()>>>>,
19}
20
21impl ConnectionManager {
22    /// Create a new connection manager
23    pub fn new(client: Client, base_url: String) -> Self {
24        Self {
25            client,
26            base_url,
27            running: Arc::new(AtomicBool::new(false)),
28            handle: Arc::new(Mutex::new(None)),
29        }
30    }
31
32    /// Start the keep-alive background task
33    /// Sends periodic lightweight requests to keep the connection warm
34    pub async fn start_keepalive(&self, interval: Duration) {
35        // If already running, return
36        if self.running.load(Ordering::Relaxed) {
37            return;
38        }
39
40        self.running.store(true, Ordering::Relaxed);
41        
42        let client = self.client.clone();
43        let base_url = self.base_url.clone();
44        let running = self.running.clone();
45
46        let handle = tokio::spawn(async move {
47            while running.load(Ordering::Relaxed) {
48                // Send a lightweight request to keep connection alive
49                // Use /time endpoint as it's fast and doesn't require auth
50                let _ = client
51                    .get(format!("{}/time", base_url))
52                    .timeout(Duration::from_secs(5))
53                    .send()
54                    .await;
55
56                // Wait for next interval
57                tokio::time::sleep(interval).await;
58            }
59        });
60
61        let mut handle_guard = self.handle.lock().await;
62        *handle_guard = Some(handle);
63    }
64
65    /// Stop the keep-alive background task
66    pub async fn stop_keepalive(&self) {
67        self.running.store(false, Ordering::Relaxed);
68        
69        let mut handle_guard = self.handle.lock().await;
70        if let Some(handle) = handle_guard.take() {
71            handle.abort();
72        }
73    }
74
75    /// Check if keep-alive is running
76    pub fn is_running(&self) -> bool {
77        self.running.load(Ordering::Relaxed)
78    }
79
80    /// Send a single keep-alive ping
81    pub async fn ping(&self) -> Result<(), reqwest::Error> {
82        self.client
83            .get(format!("{}/time", self.base_url))
84            .timeout(Duration::from_secs(5))
85            .send()
86            .await?;
87        Ok(())
88    }
89}
90
91impl Drop for ConnectionManager {
92    fn drop(&mut self) {
93        self.running.store(false, Ordering::Relaxed);
94    }
95}
96
97#[cfg(test)]
98mod tests {
99    use super::*;
100
101    #[tokio::test]
102    async fn test_connection_manager_creation() {
103        let client = Client::new();
104        let manager = ConnectionManager::new(client, "https://clob.polymarket.com".to_string());
105        assert!(!manager.is_running());
106    }
107
108    #[tokio::test]
109    async fn test_keepalive_start_stop() {
110        let client = Client::new();
111        let manager = ConnectionManager::new(client, "https://clob.polymarket.com".to_string());
112        
113        manager.start_keepalive(Duration::from_secs(30)).await;
114        tokio::time::sleep(Duration::from_millis(100)).await;
115        assert!(manager.is_running());
116        
117        manager.stop_keepalive().await;
118        tokio::time::sleep(Duration::from_millis(100)).await;
119        assert!(!manager.is_running());
120    }
121}
122
polyfill_rs/connection_manager.rs

polyfill_rs/
connection_manager.rs
