1use std::collections::HashMap;
11use std::sync::Mutex;
12use std::sync::atomic::{AtomicU32, Ordering};
13use std::time::{Duration, Instant};
14
15pub struct ConnectionPool {
18 pools: Mutex<HashMap<String, Vec<PooledConn>>>,
19 max_per_host: usize,
20 idle_timeout: Duration,
21}
22
23pub struct PooledConn {
24 host: String,
25 created_at: Instant,
26 healthy: bool,
27}
28
29impl ConnectionPool {
30 pub fn new(max_per_host: usize, idle_timeout: Duration) -> Self {
31 Self {
32 pools: Mutex::new(HashMap::new()),
33 max_per_host,
34 idle_timeout,
35 }
36 }
37
38 pub fn acquire(&self, host: &str) -> Option<PooledConn> {
39 let mut pools = self.pools.lock().unwrap();
40 let pool = pools.entry(host.into()).or_default();
41 let pos = pool
42 .iter()
43 .position(|c| c.healthy && c.created_at.elapsed() < self.idle_timeout)?;
44 Some(pool.remove(pos))
45 }
46
47 pub fn release(&self, conn: PooledConn) {
48 let mut pools = self.pools.lock().unwrap();
49 let pool = pools.entry(conn.host.clone()).or_default();
50 if pool.len() < self.max_per_host {
51 pool.push(conn);
52 }
53 }
54
55 pub fn create(&self, host: &str) -> PooledConn {
56 PooledConn {
57 host: host.into(),
58 created_at: Instant::now(),
59 healthy: true,
60 }
61 }
62
63 pub fn idle_count(&self, host: &str) -> usize {
64 self.pools
65 .lock()
66 .unwrap()
67 .get(host)
68 .map(|v| v.len())
69 .unwrap_or(0)
70 }
71
72 pub fn evict_stale(&self) -> usize {
73 let mut pools = self.pools.lock().unwrap();
74 let mut evicted = 0;
75 for pool in pools.values_mut() {
76 let before = pool.len();
77 pool.retain(|c| c.created_at.elapsed() < self.idle_timeout && c.healthy);
78 evicted += before - pool.len();
79 }
80 evicted
81 }
82}
83
84#[derive(Debug, Clone)]
87pub struct Backend {
88 pub id: String,
89 pub address: String,
90 pub health_score: u32,
91 pub healthy: bool,
92}
93
94pub struct HealthLoadBalancer {
95 backends: Mutex<Vec<Backend>>,
96 current: AtomicU32,
97}
98
99impl HealthLoadBalancer {
100 pub fn new(backends: Vec<Backend>) -> Self {
101 Self {
102 backends: Mutex::new(backends),
103 current: AtomicU32::new(0),
104 }
105 }
106
107 pub fn select(&self) -> Option<Backend> {
108 let backends = self.backends.lock().unwrap();
109 let healthy: Vec<&Backend> = backends.iter().filter(|b| b.healthy).collect();
110 if healthy.is_empty() {
111 return None;
112 }
113 let total_score: u32 = healthy.iter().map(|b| b.health_score).sum();
115 if total_score == 0 {
116 return Some(healthy[0].clone());
117 }
118 let idx = self.current.fetch_add(1, Ordering::SeqCst) as usize % healthy.len();
119 Some(healthy[idx].clone())
120 }
121
122 pub fn update_health(&self, id: &str, score: u32, healthy: bool) {
123 let mut backends = self.backends.lock().unwrap();
124 if let Some(b) = backends.iter_mut().find(|b| b.id == id) {
125 b.health_score = score;
126 b.healthy = healthy;
127 }
128 }
129
130 pub fn healthy_count(&self) -> usize {
131 self.backends
132 .lock()
133 .unwrap()
134 .iter()
135 .filter(|b| b.healthy)
136 .count()
137 }
138
139 pub fn total_count(&self) -> usize {
140 self.backends.lock().unwrap().len()
141 }
142}
143
144pub enum JitterStrategy {
147 Full,
148 Equal,
149 Decorrelated,
150}
151
152pub fn backoff_with_jitter(
153 attempt: u32,
154 base_delay: Duration,
155 max_delay: Duration,
156 strategy: JitterStrategy,
157) -> Duration {
158 let exponential = base_delay.mul_f64(2f64.powi(attempt as i32)).min(max_delay);
159 let mut rng_bytes = [0u8; 4];
160 use rand_core::{OsRng, RngCore};
161 OsRng.fill_bytes(&mut rng_bytes);
162 let rand_val = u32::from_le_bytes(rng_bytes) as f64 / u32::MAX as f64;
163 match strategy {
164 JitterStrategy::Full => {
165 Duration::from_nanos((rand_val * exponential.as_nanos() as f64) as u64)
166 }
167 JitterStrategy::Equal => {
168 let half = exponential.as_nanos() as f64 / 2.0;
169 Duration::from_nanos((half + rand_val * half) as u64)
170 }
171 JitterStrategy::Decorrelated => {
172 Duration::from_nanos((rand_val * max_delay.as_nanos() as f64) as u64).min(max_delay)
173 }
174 }
175}
176
177pub struct DegradationCache<T: Clone> {
180 cache: Mutex<HashMap<String, (T, Instant)>>,
181 max_stale: Duration,
182}
183
184impl<T: Clone> DegradationCache<T> {
185 pub fn new(max_stale: Duration) -> Self {
186 Self {
187 cache: Mutex::new(HashMap::new()),
188 max_stale,
189 }
190 }
191
192 pub fn store(&self, key: &str, value: T) {
193 self.cache
194 .lock()
195 .unwrap()
196 .insert(key.into(), (value, Instant::now()));
197 }
198
199 pub fn get_fresh(&self, key: &str) -> Option<T> {
200 let cache = self.cache.lock().unwrap();
201 cache
202 .get(key)
203 .filter(|(_, t)| t.elapsed() < self.max_stale)
204 .map(|(v, _)| v.clone())
205 }
206
207 pub fn get_stale(&self, key: &str) -> Option<T> {
208 self.cache.lock().unwrap().get(key).map(|(v, _)| v.clone())
209 }
210
211 pub fn is_stale(&self, key: &str) -> bool {
212 let cache = self.cache.lock().unwrap();
213 cache
214 .get(key)
215 .map(|(_, t)| t.elapsed() >= self.max_stale)
216 .unwrap_or(true)
217 }
218
219 pub fn entry_count(&self) -> usize {
220 self.cache.lock().unwrap().len()
221 }
222
223 pub fn evict_stale(&self) -> usize {
224 let mut cache = self.cache.lock().unwrap();
225 let before = cache.len();
226 cache.retain(|_, (_, t)| t.elapsed() < self.max_stale);
227 before - cache.len()
228 }
229}
230
231#[derive(Debug, Clone)]
234pub struct WsSession {
235 pub session_id: String,
236 pub subscribers: Vec<String>,
237 pub last_update: chrono::DateTime<chrono::Utc>,
238}
239
240pub struct WsCoordinator {
241 sessions: Mutex<HashMap<String, WsSession>>,
242}
243
244impl WsCoordinator {
245 pub fn new() -> Self {
246 Self {
247 sessions: Mutex::new(HashMap::new()),
248 }
249 }
250
251 pub fn subscribe(&self, session_id: &str, client_id: &str) {
252 let mut sessions = self.sessions.lock().unwrap();
253 let session = sessions
254 .entry(session_id.into())
255 .or_insert_with(|| WsSession {
256 session_id: session_id.into(),
257 subscribers: Vec::new(),
258 last_update: chrono::Utc::now(),
259 });
260 if !session.subscribers.contains(&client_id.into()) {
261 session.subscribers.push(client_id.into());
262 }
263 }
264
265 pub fn unsubscribe(&self, session_id: &str, client_id: &str) {
266 let mut sessions = self.sessions.lock().unwrap();
267 if let Some(session) = sessions.get_mut(session_id) {
268 session.subscribers.retain(|s| s != client_id);
269 if session.subscribers.is_empty() {
270 sessions.remove(session_id);
271 }
272 }
273 }
274
275 pub fn broadcast(&self, session_id: &str) -> Vec<String> {
276 let mut sessions = self.sessions.lock().unwrap();
277 if let Some(session) = sessions.get_mut(session_id) {
278 session.last_update = chrono::Utc::now();
279 return session.subscribers.clone();
280 }
281 Vec::new()
282 }
283
284 pub fn subscriber_count(&self, session_id: &str) -> usize {
285 self.sessions
286 .lock()
287 .unwrap()
288 .get(session_id)
289 .map(|s| s.subscribers.len())
290 .unwrap_or(0)
291 }
292
293 pub fn active_sessions(&self) -> usize {
294 self.sessions.lock().unwrap().len()
295 }
296}
297
298impl Default for WsCoordinator {
299 fn default() -> Self {
300 Self::new()
301 }
302}
303
304#[derive(Debug, Clone)]
307pub struct WireLabel {
308 pub zero: Vec<u8>,
309 pub one: Vec<u8>,
310}
311
312pub struct GarbledGate {
313 pub input_a: WireLabel,
314 pub input_b: WireLabel,
315 pub output: WireLabel,
316 pub table: [[Vec<u8>; 2]; 2],
317}
318
319pub fn garble_and_gate() -> GarbledGate {
320 use sha2::{Digest, Sha256};
321 let a = WireLabel {
322 zero: rand_bytes(16),
323 one: rand_bytes(16),
324 };
325 let b = WireLabel {
326 zero: rand_bytes(16),
327 one: rand_bytes(16),
328 };
329 let output = WireLabel {
330 zero: rand_bytes(16),
331 one: rand_bytes(16),
332 };
333
334 let a_labels = [&a.zero, &a.one];
337 let b_labels = [&b.zero, &b.one];
338 let mut table = [[vec![], vec![]], [vec![], vec![]]];
339 for (ai, &al) in a_labels.iter().enumerate() {
340 for (bi, &bl) in b_labels.iter().enumerate() {
341 let out_val = if ai == 1 && bi == 1 {
342 &output.one
343 } else {
344 &output.zero
345 };
346 let mut h = Sha256::new();
347 h.update(b"garble");
348 h.update(al);
349 h.update(bl);
350 let key = h.finalize();
351 let encrypted: Vec<u8> = out_val.iter().zip(key.iter()).map(|(o, k)| o ^ k).collect();
352 table[ai][bi] = encrypted;
353 }
354 }
355 GarbledGate {
356 input_a: a,
357 input_b: b,
358 output,
359 table,
360 }
361}
362
363pub fn evaluate_gate(gate: &GarbledGate, a: &[u8], b: &[u8]) -> Option<Vec<u8>> {
364 use sha2::{Digest, Sha256};
365 for row in &gate.table {
366 for cell in row {
367 let mut h = Sha256::new();
368 h.update(b"garble");
369 h.update(a);
370 h.update(b);
371 let key = h.finalize();
372 let decrypted: Vec<u8> = cell.iter().zip(key.iter()).map(|(c, k)| c ^ k).collect();
373 if decrypted == gate.output.zero || decrypted == gate.output.one {
374 return Some(decrypted);
375 }
376 }
377 }
378 None
379}
380
381pub struct BulkSignPipeline {
384 batch_size: usize,
385 pending: Mutex<Vec<BulkSignItem>>,
386}
387
388#[derive(Debug, Clone)]
389pub struct BulkSignItem {
390 pub message_hash: Vec<u8>,
391 pub quorum_id: String,
392}
393
394#[derive(Debug, Clone)]
395pub struct BulkSignResult {
396 pub signatures: Vec<Vec<u8>>,
397 pub batch_count: usize,
398}
399
400impl BulkSignPipeline {
401 pub fn new(batch_size: usize) -> Self {
402 Self {
403 batch_size,
404 pending: Mutex::new(Vec::new()),
405 }
406 }
407
408 pub fn enqueue(&self, item: BulkSignItem) -> bool {
409 let mut pending = self.pending.lock().unwrap();
410 pending.push(item);
411 pending.len() >= self.batch_size
412 }
413
414 pub fn flush(&self) -> Vec<BulkSignItem> {
415 let mut pending = self.pending.lock().unwrap();
416 std::mem::take(&mut *pending)
417 }
418
419 pub fn pending_count(&self) -> usize {
420 self.pending.lock().unwrap().len()
421 }
422
423 pub fn batch_and_sign<F>(&self, sign_fn: F) -> BulkSignResult
424 where
425 F: Fn(&[BulkSignItem]) -> Vec<Vec<u8>>,
426 {
427 let batch = self.flush();
428 let sigs = sign_fn(&batch);
429 BulkSignResult {
430 signatures: sigs,
431 batch_count: batch.len(),
432 }
433 }
434}
435
436fn rand_bytes(n: usize) -> Vec<u8> {
437 use rand_core::{OsRng, RngCore};
438 let mut buf = vec![0u8; n];
439 OsRng.fill_bytes(&mut buf);
440 buf
441}
442
443#[cfg(test)]
444mod tests {
445 use super::*;
446
447 #[test]
449 fn pool_create_and_release() {
450 let pool = ConnectionPool::new(5, Duration::from_secs(60));
451 let conn = pool.create("host1");
452 pool.release(conn);
453 assert_eq!(pool.idle_count("host1"), 1);
454 }
455
456 #[test]
457 fn pool_acquire_reuses() {
458 let pool = ConnectionPool::new(5, Duration::from_secs(60));
459 let conn = pool.create("host1");
460 pool.release(conn);
461 let acquired = pool.acquire("host1");
462 assert!(acquired.is_some());
463 assert_eq!(pool.idle_count("host1"), 0);
464 }
465
466 #[test]
467 fn pool_evict_stale() {
468 let pool = ConnectionPool::new(5, Duration::from_millis(1));
469 let conn = pool.create("host1");
470 pool.release(conn);
471 std::thread::sleep(Duration::from_millis(5));
472 assert_eq!(pool.evict_stale(), 1);
473 }
474
475 #[test]
477 fn lb_selects_healthy() {
478 let lb = HealthLoadBalancer::new(vec![
479 Backend {
480 id: "a".into(),
481 address: "a:80".into(),
482 health_score: 100,
483 healthy: true,
484 },
485 Backend {
486 id: "b".into(),
487 address: "b:80".into(),
488 health_score: 50,
489 healthy: false,
490 },
491 ]);
492 let selected = lb.select().unwrap();
493 assert_eq!(selected.id, "a");
494 }
495
496 #[test]
497 fn lb_no_healthy_returns_none() {
498 let lb = HealthLoadBalancer::new(vec![Backend {
499 id: "a".into(),
500 address: "a:80".into(),
501 health_score: 100,
502 healthy: false,
503 }]);
504 assert!(lb.select().is_none());
505 }
506
507 #[test]
508 fn lb_update_health() {
509 let lb = HealthLoadBalancer::new(vec![Backend {
510 id: "a".into(),
511 address: "a:80".into(),
512 health_score: 100,
513 healthy: true,
514 }]);
515 lb.update_health("a", 50, false);
516 assert_eq!(lb.healthy_count(), 0);
517 }
518
519 #[test]
521 fn backoff_returns_delay() {
522 let delay = backoff_with_jitter(
523 0,
524 Duration::from_millis(100),
525 Duration::from_secs(10),
526 JitterStrategy::Full,
527 );
528 assert!(delay <= Duration::from_millis(100));
529 }
530
531 #[test]
532 fn backoff_capped_at_max() {
533 let delay = backoff_with_jitter(
534 20,
535 Duration::from_millis(100),
536 Duration::from_secs(1),
537 JitterStrategy::Full,
538 );
539 assert!(delay <= Duration::from_secs(1));
540 }
541
542 #[test]
544 fn cache_store_get_fresh() {
545 let cache = DegradationCache::<String>::new(Duration::from_secs(60));
546 cache.store("k1", "value".into());
547 assert_eq!(cache.get_fresh("k1"), Some("value".into()));
548 assert!(!cache.is_stale("k1"));
549 }
550
551 #[test]
552 fn cache_stale_after_ttl() {
553 let cache = DegradationCache::<String>::new(Duration::from_millis(1));
554 cache.store("k1", "value".into());
555 std::thread::sleep(Duration::from_millis(5));
556 assert!(cache.is_stale("k1"));
557 assert_eq!(cache.get_stale("k1"), Some("value".into()));
558 assert!(cache.get_fresh("k1").is_none());
559 }
560
561 #[test]
563 fn ws_subscribe_unsubscribe() {
564 let ws = WsCoordinator::new();
565 ws.subscribe("s1", "c1");
566 ws.subscribe("s1", "c2");
567 assert_eq!(ws.subscriber_count("s1"), 2);
568 ws.unsubscribe("s1", "c1");
569 assert_eq!(ws.subscriber_count("s1"), 1);
570 }
571
572 #[test]
573 fn ws_broadcast_returns_subscribers() {
574 let ws = WsCoordinator::new();
575 ws.subscribe("s1", "c1");
576 ws.subscribe("s1", "c2");
577 let recipients = ws.broadcast("s1");
578 assert_eq!(recipients.len(), 2);
579 }
580
581 #[test]
583 fn garble_and_evaluate_and() {
584 let gate = garble_and_gate();
585 let result00 = evaluate_gate(&gate, &gate.input_a.zero, &gate.input_b.zero);
586 let result11 = evaluate_gate(&gate, &gate.input_a.one, &gate.input_b.one);
587 assert!(result00.is_some());
588 assert!(result11.is_some());
589 assert_eq!(result00.unwrap(), gate.output.zero);
590 assert_eq!(result11.unwrap(), gate.output.one);
591 }
592
593 #[test]
595 fn bulk_enqueue_and_flush() {
596 let pipeline = BulkSignPipeline::new(3);
597 assert!(!pipeline.enqueue(BulkSignItem {
598 message_hash: vec![1],
599 quorum_id: "q".into()
600 }));
601 assert!(!pipeline.enqueue(BulkSignItem {
602 message_hash: vec![2],
603 quorum_id: "q".into()
604 }));
605 assert!(pipeline.enqueue(BulkSignItem {
606 message_hash: vec![3],
607 quorum_id: "q".into()
608 }));
609 let batch = pipeline.flush();
610 assert_eq!(batch.len(), 3);
611 }
612
613 #[test]
614 fn bulk_batch_and_sign() {
615 let pipeline = BulkSignPipeline::new(2);
616 pipeline.enqueue(BulkSignItem {
617 message_hash: vec![1],
618 quorum_id: "q".into(),
619 });
620 pipeline.enqueue(BulkSignItem {
621 message_hash: vec![2],
622 quorum_id: "q".into(),
623 });
624 let result =
625 pipeline.batch_and_sign(|items| items.iter().map(|i| i.message_hash.clone()).collect());
626 assert_eq!(result.batch_count, 2);
627 assert_eq!(result.signatures.len(), 2);
628 }
629}