1use crate::error::*;
2use crate::host_identity::{HostConnectError, HostExpectation, HostProcess};
3use crate::private::PrivateSecret;
4use crate::record::{try_read_record, ServiceRecord, PROTOCOL_MAJOR};
5use crate::suite::{Suite, SuiteError};
6use crate::transport::{connect, connect_expecting, Conn};
7use rightkit_sidecar::{read_json_line, write_json_line, FrameError};
8use serde_json::{json, Value};
9use std::io::{BufReader, Read, Write};
10use std::sync::atomic::{AtomicBool, AtomicU64, AtomicUsize, Ordering};
11use std::sync::{mpsc, Arc, Mutex};
12use std::time::{Duration, Instant};
13
14const LIMIT: usize = 16 * 1024 * 1024;
15
16pub const DEFAULT_DEADLINE: Duration = Duration::from_secs(120);
19
20struct Watchdog {
24 fired: Arc<AtomicBool>,
25 _disarm: mpsc::Sender<()>,
26}
27
28impl Watchdog {
29 fn arm(after: Option<Duration>, shutdown: Arc<dyn Fn() + Send + Sync>) -> Watchdog {
30 let fired = Arc::new(AtomicBool::new(false));
31 let (tx, rx) = mpsc::channel::<()>();
32 if let Some(d) = after {
33 let f = fired.clone();
34 std::thread::spawn(move || {
35 if let Err(mpsc::RecvTimeoutError::Timeout) = rx.recv_timeout(d) {
36 f.store(true, Ordering::SeqCst);
37 shutdown();
38 }
39 });
40 }
41 Watchdog { fired, _disarm: tx }
42 }
43 fn fired(&self) -> bool {
44 self.fired.load(Ordering::SeqCst)
45 }
46}
47
48fn timed_out(what: &str, d: Option<Duration>) -> ServiceError {
49 ServiceError::new(
50 ERROR_TIMEOUT,
51 format!(
52 "{what} exceeded the {} ms deadline; connection aborted",
53 d.map(|d| d.as_millis()).unwrap_or(0)
54 ),
55 )
56}
57
58fn lock_until<T>(m: &Mutex<T>, deadline: Option<Instant>) -> Option<std::sync::MutexGuard<'_, T>> {
61 use std::sync::TryLockError;
62 let Some(deadline) = deadline else {
63 return Some(m.lock().unwrap_or_else(|e| e.into_inner()));
64 };
65 let mut pause = Duration::from_micros(200);
66 loop {
67 match m.try_lock() {
68 Ok(g) => return Some(g),
69 Err(TryLockError::Poisoned(e)) => return Some(e.into_inner()),
70 Err(TryLockError::WouldBlock) => {}
71 }
72 let left = deadline.saturating_duration_since(Instant::now());
73 if left.is_zero() {
74 return None;
75 }
76 std::thread::sleep(pause.min(left));
77 pause = (pause * 2).min(Duration::from_millis(5));
78 }
79}
80
81struct Io {
82 reader: BufReader<Box<dyn Read + Send>>,
83 writer: Box<dyn Write + Send>,
84 buf: Vec<u8>,
85}
86
87#[derive(Debug, Clone)]
106pub struct ClientConfig {
107 service: String,
108 app: String,
109 suite: Option<Suite>,
110 deadline: Option<Duration>,
111 ignore_env_overrides: bool,
112 secret: Option<PrivateSecret>,
113}
114
115impl ClientConfig {
116 pub fn new(service: impl Into<String>, app: impl Into<String>) -> Self {
120 Self {
121 service: service.into(),
122 app: app.into(),
123 suite: None,
124 deadline: Some(DEFAULT_DEADLINE),
125 ignore_env_overrides: false,
126 secret: None,
127 }
128 }
129
130 pub fn secret(mut self, secret: PrivateSecret) -> Self {
134 self.secret = Some(secret);
135 self
136 }
137
138 pub fn suite(mut self, suite: Suite) -> Self {
140 self.suite = Some(suite);
141 self
142 }
143
144 pub fn deadline(mut self, deadline: Option<Duration>) -> Self {
146 self.deadline = deadline;
147 self
148 }
149
150 pub fn ignore_env_overrides(mut self, ignore: bool) -> Self {
155 self.ignore_env_overrides = ignore;
156 self
157 }
158
159 pub fn service(&self) -> &str {
160 &self.service
161 }
162
163 pub fn app(&self) -> &str {
164 &self.app
165 }
166
167 pub fn ignores_env_overrides(&self) -> bool {
169 self.ignore_env_overrides
170 }
171
172 #[allow(deprecated)] pub fn resolve_suite(&self) -> Result<Suite, SuiteError> {
175 match (&self.suite, self.ignore_env_overrides) {
176 (Some(suite), true) if suite.env_derived() => Err(SuiteError::Insecure(format!(
177 "the suite {} was derived from the environment, which this client ignores; \
178 pass Suite::at(..) or leave the suite unset",
179 suite.root().display()
180 ))),
181 (Some(suite), _) => Ok(suite.clone()),
182 (None, ignore) => Suite::resolve_ignoring_env_overrides(ignore),
183 }
184 }
185}
186
187fn suite_error(e: SuiteError) -> ServiceError {
188 ServiceError::new(ERROR_SUITE_UNRESOLVED, e.to_string())
189}
190
191pub struct Client {
196 io: Mutex<Io>,
197 next_id: AtomicU64,
198 dead: AtomicBool,
199 hello: Value,
200 record: ServiceRecord,
201 shutdown: Arc<dyn Fn() + Send + Sync>,
202 deadline_ms: AtomicU64,
204 host: Option<HostProcess>,
206 secret: Option<PrivateSecret>,
208}
209
210fn with_secret(mut frame: Value, secret: Option<&PrivateSecret>) -> Value {
212 if let Some(secret) = secret {
213 frame["secret"] = Value::String(secret.expose().to_string());
214 }
215 frame
216}
217
218fn closed() -> ServiceError {
219 ServiceError::retryable(
220 ERROR_CONNECTION_CLOSED,
221 "connection closed before a response arrived",
222 )
223}
224
225impl Client {
226 pub fn connect(suite: &Suite, service: &str, app: &str) -> Result<Client, ServiceError> {
236 Self::connect_with(suite, service, app, Some(DEFAULT_DEADLINE))
237 }
238
239 pub fn connect_with(
243 suite: &Suite,
244 service: &str,
245 app: &str,
246 deadline: Option<Duration>,
247 ) -> Result<Client, ServiceError> {
248 Self::connect_bounded(suite, service, app, deadline, None, None)
249 .map_err(ConnectError::into_service_error)
250 }
251
252 pub fn connect_config(config: &ClientConfig) -> Result<Client, ServiceError> {
256 let suite = config.resolve_suite().map_err(suite_error)?;
257 Self::connect_bounded(
258 &suite,
259 &config.service,
260 &config.app,
261 config.deadline,
262 None,
263 config.secret.clone(),
264 )
265 .map_err(ConnectError::into_service_error)
266 }
267
268 pub fn connect_expecting(
291 config: &ClientConfig,
292 expect: &HostExpectation,
293 ) -> Result<Client, ConnectError> {
294 let suite = config
295 .resolve_suite()
296 .map_err(|e| ConnectError::Service(suite_error(e)))?;
297 Self::connect_bounded(
298 &suite,
299 &config.service,
300 &config.app,
301 config.deadline,
302 Some(expect.clone()),
303 config.secret.clone(),
304 )
305 }
306
307 fn connect_bounded(
311 suite: &Suite,
312 service: &str,
313 app: &str,
314 deadline: Option<Duration>,
315 expect: Option<HostExpectation>,
316 secret: Option<PrivateSecret>,
317 ) -> Result<Client, ConnectError> {
318 let Some(d) = deadline else {
319 return Self::handshake(suite, service, app, None, expect.as_ref(), secret);
320 };
321 let (tx, rx) = mpsc::channel();
322 let (s, sv, ap) = (suite.clone(), service.to_string(), app.to_string());
323 std::thread::spawn(move || {
324 let _ = tx.send(Self::handshake(
325 &s,
326 &sv,
327 &ap,
328 Some(d),
329 expect.as_ref(),
330 secret,
331 ));
332 });
333 match rx.recv_timeout(d + Duration::from_millis(250)) {
334 Ok(r) => r,
335 Err(_) => Err(ConnectError::Service(timed_out("connect", deadline))),
336 }
337 }
338
339 fn handshake(
340 suite: &Suite,
341 service: &str,
342 app: &str,
343 deadline: Option<Duration>,
344 expect: Option<&HostExpectation>,
345 secret: Option<PrivateSecret>,
346 ) -> Result<Client, ConnectError> {
347 let started = Instant::now();
348 let record = match try_read_record(suite, service) {
351 Ok(Some(record)) => record,
352 Ok(None) => {
353 return Err(ServiceError::retryable(
354 ERROR_UNAVAILABLE,
355 format!("no service.json for {service}"),
356 )
357 .into())
358 }
359 Err(e) => {
360 return Err(ServiceError::retryable(
361 ERROR_UNAVAILABLE,
362 format!("could not read service.json for {service} (try again): {e}"),
363 )
364 .into())
365 }
366 };
367 if record.protocol_major != PROTOCOL_MAJOR {
368 return Err(ServiceError::new(
369 "protocol_mismatch",
370 format!("host speaks protocol {}", record.protocol_major),
371 )
372 .into());
373 }
374 let unreachable_host = |e: std::io::Error| {
375 ServiceError::retryable(
376 ERROR_UNAVAILABLE,
377 format!("{service} host unreachable: {e}"),
378 )
379 };
380 let (conn, host) = match expect {
382 None => (connect(&record.endpoint).map_err(unreachable_host)?, None),
383 Some(expect) => match connect_expecting(&record.endpoint, expect) {
384 Ok((conn, host)) => (conn, Some(host)),
385 Err(HostConnectError::Io(e)) => return Err(unreachable_host(e).into()),
386 Err(HostConnectError::Identity(e)) => return Err(ConnectError::HostIdentity(e)),
387 },
388 };
389 let Conn {
390 reader,
391 writer,
392 shutdown,
393 ..
394 } = conn;
395 let mut io = Io {
396 reader: BufReader::new(reader),
397 writer,
398 buf: Vec::new(),
399 };
400 let wd = Watchdog::arm(
401 deadline.map(|d| d.saturating_sub(started.elapsed())),
402 shutdown.clone(),
403 );
404 let hello_res = (|| {
405 write_json_line(
406 &mut io.writer,
407 &with_secret(
408 json!({"type": "hello", "app": app, "protocol_major": PROTOCOL_MAJOR}),
409 secret.as_ref(),
410 ),
411 LIMIT,
412 )
413 .map_err(|e| ServiceError::new("io", e.to_string()))?;
414 io.writer.flush()?;
415 read_json_line::<_, Value>(&mut io.reader, &mut io.buf, LIMIT)
416 .map_err(|e| ServiceError::new("io", e.to_string()))?
417 .ok_or_else(closed)
418 })();
419 let hello = match hello_res {
420 Ok(h) => h,
421 Err(e) => {
422 let t = wd.fired();
423 (shutdown)();
424 let e = if t { timed_out("hello", deadline) } else { e };
425 return Err(e.into());
426 }
427 };
428 drop(wd);
429 if hello["ok"] != true {
430 return Err(ServiceError::from_value(&hello["error"]).into());
431 }
432 Ok(Client {
433 io: Mutex::new(io),
434 next_id: AtomicU64::new(1),
435 dead: AtomicBool::new(false),
436 hello,
437 record,
438 shutdown,
439 deadline_ms: AtomicU64::new(deadline.map(|d| d.as_millis() as u64).unwrap_or(0)),
440 host,
441 secret,
442 })
443 }
444
445 pub fn verified_host(&self) -> Option<&HostProcess> {
449 self.host.as_ref()
450 }
451
452 pub fn set_deadline(&self, deadline: Option<Duration>) {
454 self.deadline_ms.store(
455 deadline.map(|d| (d.as_millis() as u64).max(1)).unwrap_or(0),
456 Ordering::SeqCst,
457 );
458 }
459
460 pub fn record(&self) -> &ServiceRecord {
461 &self.record
462 }
463 pub fn hello(&self) -> &Value {
465 &self.hello
466 }
467
468 pub fn abort(&self) {
471 self.dead.store(true, Ordering::SeqCst);
472 (self.shutdown)();
473 }
474
475 pub fn is_closed(&self) -> bool {
479 self.dead.load(Ordering::SeqCst)
480 }
481
482 pub fn call(&self, method: &str, params: Value) -> Result<Value, ServiceError> {
483 self.call_events(method, params, &mut |_| {})
484 }
485
486 pub fn call_events(
491 &self,
492 method: &str,
493 params: Value,
494 on_event: &mut dyn FnMut(&Value),
495 ) -> Result<Value, ServiceError> {
496 let deadline = match self.deadline_ms.load(Ordering::SeqCst) {
497 0 => None,
498 ms => Some(Instant::now() + Duration::from_millis(ms)),
499 };
500 self.call_events_until(method, params, deadline, on_event)
501 }
502
503 pub fn call_events_until(
508 &self,
509 method: &str,
510 params: Value,
511 deadline: Option<Instant>,
512 on_event: &mut dyn FnMut(&Value),
513 ) -> Result<Value, ServiceError> {
514 let mut frame = self.call_frame_until(method, params, false, deadline, on_event)?;
515 Ok(frame["result"].take())
516 }
517
518 fn call_frame_until(
521 &self,
522 method: &str,
523 params: Value,
524 chunked_ok: bool,
525 deadline: Option<Instant>,
526 on_event: &mut dyn FnMut(&Value),
527 ) -> Result<Value, ServiceError> {
528 let started = Instant::now();
529 let total = deadline.map(|d| d.saturating_duration_since(started));
530 if self.dead.load(Ordering::SeqCst) {
531 return Err(closed());
532 }
533 let id = self.next_id.fetch_add(1, Ordering::SeqCst);
534 let Some(mut io) = lock_until(&self.io, deadline) else {
535 return Err(timed_out(
536 &format!("request `{method}` waiting for the connection"),
537 total,
538 ));
539 };
540 if self.dead.load(Ordering::SeqCst) {
541 return Err(closed());
542 }
543 let wd = Watchdog::arm(
544 deadline.map(|d| d.saturating_duration_since(Instant::now())),
545 self.shutdown.clone(),
546 );
547 let res = (|| {
548 let mut frame =
549 json!({"type": "request", "id": id, "method": method, "params": params});
550 if chunked_ok {
551 frame["chunked_ok"] = Value::Bool(true);
552 }
553 write_json_line(
554 &mut io.writer,
555 &with_secret(frame, self.secret.as_ref()),
556 LIMIT,
557 )
558 .map_err(|e| match e {
559 FrameError::TooLarge { .. } => ServiceError::new(
560 ERROR_REQUEST_TOO_LARGE,
561 "request exceeds the frame limit (Client::call_large sends it in pieces)",
562 ),
563 _ => closed(),
564 })?;
565 io.writer.flush().map_err(|_| closed())?;
566 loop {
567 let Io { reader, buf, .. } = &mut *io;
568 match read_json_line::<_, Value>(reader, buf, LIMIT) {
569 Ok(Some(v)) if v["id"].as_u64() == Some(id) => {
570 if v["type"] == "event" {
571 on_event(&v["event"]);
572 } else if v["type"] == "response" {
573 return if v["ok"] == true {
574 Ok(v)
575 } else {
576 Err(ServiceError::from_value(&v["error"]))
577 };
578 }
579 }
580 Ok(Some(_)) => {} Ok(None) | Err(FrameError::Io(_)) => return Err(closed()),
582 Err(FrameError::TooLarge { .. }) => {
586 return Err(ServiceError::new(
587 ERROR_RESPONSE_TOO_LARGE,
588 "the host's reply exceeds the client's frame limit",
589 ));
590 }
591 Err(_) => {}
592 }
593 }
594 })();
595 if wd.fired() && res.is_err() {
596 self.dead.store(true, Ordering::SeqCst);
597 return Err(timed_out(&format!("request `{method}`"), total));
598 }
599 if matches!(&res, Err(e) if e.code == ERROR_CONNECTION_CLOSED) {
600 self.dead.store(true, Ordering::SeqCst);
601 }
602 res
603 }
604}
605
606impl Client {
607 fn call_step(
609 &self,
610 method: &str,
611 params: Value,
612 chunked_ok: bool,
613 on_event: &mut dyn FnMut(&Value),
614 ) -> Result<Value, ServiceError> {
615 let deadline = match self.deadline_ms.load(Ordering::SeqCst) {
616 0 => None,
617 ms => Some(Instant::now() + Duration::from_millis(ms)),
618 };
619 self.call_frame_until(method, params, chunked_ok, deadline, on_event)
620 }
621
622 pub fn call_large(
643 &self,
644 method: &str,
645 params: Value,
646 on_event: &mut dyn FnMut(&Value),
647 ) -> Result<Value, ServiceError> {
648 self.call_large_until(method, params, None, on_event)
649 }
650
651 fn call_large_until(
654 &self,
655 method: &str,
656 params: Value,
657 deadline: Option<Option<Instant>>,
658 on_event: &mut dyn FnMut(&Value),
659 ) -> Result<Value, ServiceError> {
660 let step =
661 |method: &str, params: Value, chunked, sink: &mut dyn FnMut(&Value)| match deadline {
662 Some(limit) => {
663 if limit.is_some_and(|d| Instant::now() >= d) {
664 return Err(timed_out("pool request", Some(Duration::ZERO)));
665 }
666 self.call_frame_until(method, params, chunked, limit, sink)
667 }
668 None => self.call_step(method, params, chunked, sink),
669 };
670 use crate::chunk::{boundaries, CHUNK_TEXT_BYTES, METHOD_FETCH, METHOD_UPLOAD};
671 let advertised = |key: &str| {
672 self.hello[key]
673 .as_u64()
674 .and_then(|n| usize::try_from(n).ok())
675 };
676 let max_message = advertised("max_message").unwrap_or(0);
677 let frame_limit = advertised("frame_limit").unwrap_or(LIMIT).min(LIMIT);
678 let chunk = advertised("chunk_bytes")
679 .unwrap_or(CHUNK_TEXT_BYTES)
680 .clamp(16, CHUNK_TEXT_BYTES);
681 let text = serde_json::to_string(¶ms)
682 .map_err(|e| ServiceError::new("bad_request", e.to_string()))?;
683 let single_max = if max_message == 0 {
684 frame_limit - (64 * 1024).min(frame_limit / 2)
688 } else {
689 frame_limit / 2
690 };
691 let mut frame = if text.len() <= single_max {
692 drop(text);
693 step(method, params, max_message > 0, &mut *on_event)?
694 } else {
695 if max_message == 0 || text.len() > max_message {
696 return Err(ServiceError::new(
697 ERROR_REQUEST_TOO_LARGE,
698 if max_message == 0 {
699 "the request exceeds the frame limit and the host does not chunk"
700 .to_string()
701 } else {
702 format!(
703 "the request is {} bytes, over the host's {max_message}-byte message limit",
704 text.len()
705 )
706 },
707 ));
708 }
709 drop(params);
710 let ends = boundaries(&text, chunk);
711 let upload_id = format!("u{}", self.next_id.fetch_add(1, Ordering::SeqCst));
712 let mut start = 0;
713 let mut answer = Value::Null;
714 for (seq, end) in ends.iter().enumerate() {
715 let last = seq + 1 == ends.len();
716 let mut piece = json!({
717 "id": upload_id, "seq": seq, "last": last, "text": &text[start..*end],
718 });
719 if seq == 0 {
720 piece["method"] = Value::String(method.to_string());
721 }
722 start = *end;
723 let frame = step(METHOD_UPLOAD, piece, last, &mut *on_event)?;
724 if last {
725 answer = frame;
726 }
727 }
728 answer
729 };
730 let held = frame.get_mut("chunked").map(Value::take);
731 let Some(held) = held else {
732 return Ok(frame["result"].take());
733 };
734 let bad = |why: &str| ServiceError::new(ERROR_BAD_CHUNK, why.to_string());
736 let id = held["id"]
737 .as_str()
738 .ok_or_else(|| bad("the chunked result has no id"))?
739 .to_string();
740 let bytes = held["bytes"]
741 .as_u64()
742 .and_then(|n| usize::try_from(n).ok())
743 .ok_or_else(|| bad("the chunked result has no size"))?;
744 let count = held["count"]
745 .as_u64()
746 .ok_or_else(|| bad("the chunked result has no count"))?;
747 if bytes > max_message || count > bytes.max(1) as u64 {
748 return Err(bad(
749 "the chunked result is larger than the host's own message limit",
750 ));
751 }
752 let mut joined = String::with_capacity(bytes);
753 for seq in 0..count {
754 let request = json!({"id": id, "seq": seq});
755 let piece = step(METHOD_FETCH, request, false, &mut *on_event)?;
756 let part = piece["result"]["text"]
757 .as_str()
758 .ok_or_else(|| bad("a result piece has no text"))?;
759 if joined.len() + part.len() > bytes {
760 return Err(bad("the result pieces add up to more than announced"));
761 }
762 joined.push_str(part);
763 }
764 if joined.len() != bytes {
765 return Err(bad("the result pieces add up to less than announced"));
766 }
767 serde_json::from_str(&joined)
768 .map_err(|e| bad(&format!("the assembled result is not JSON: {e}")))
769 }
770}
771
772#[derive(Debug, Clone, Copy)]
774pub struct RetryPolicy {
775 pub max_attempts: u32,
776 pub backoff: Duration,
777 pub give_up_after: Duration,
780 pub request_deadline: Option<Duration>,
782 pub idempotent: bool,
788}
789
790impl Default for RetryPolicy {
791 fn default() -> Self {
792 Self {
793 max_attempts: 8,
794 backoff: Duration::from_millis(150),
795 give_up_after: Duration::from_secs(120),
796 request_deadline: Some(DEFAULT_DEADLINE),
797 idempotent: false,
798 }
799 }
800}
801
802fn run_attempt(
805 c: &Client,
806 method: &str,
807 params: &Value,
808 stop: Option<&(dyn Fn() -> bool + Sync)>,
809 on_event: &mut dyn FnMut(&Value),
810) -> Result<Value, ServiceError> {
811 let Some(stop) = stop else {
812 return c.call_large(method, params.clone(), on_event);
813 };
814 let done = AtomicBool::new(false);
815 std::thread::scope(|scope| {
816 scope.spawn(|| {
817 while !done.load(Ordering::SeqCst) {
818 if stop() {
819 c.abort();
820 break;
821 }
822 std::thread::sleep(Duration::from_millis(25));
823 }
824 });
825 let result = c.call_large(method, params.clone(), on_event);
826 done.store(true, Ordering::SeqCst);
827 result
828 })
829}
830
831fn sleep_unless_stopped(d: Duration, stop: Option<&(dyn Fn() -> bool + Sync)>) {
834 let Some(stop) = stop else {
835 std::thread::sleep(d);
836 return;
837 };
838 let end = Instant::now() + d;
839 loop {
840 if stop() {
841 return;
842 }
843 let now = Instant::now();
844 if now >= end {
845 return;
846 }
847 std::thread::sleep((end - now).min(Duration::from_millis(25)));
848 }
849}
850
851impl Client {
852 pub fn call_retrying(
856 suite: &Suite,
857 service: &str,
858 app: &str,
859 method: &str,
860 params: Value,
861 policy: RetryPolicy,
862 on_event: &mut dyn FnMut(&Value),
863 ) -> Result<Value, ServiceError> {
864 Self::retry_loop(
865 &|deadline| Self::connect_bounded(suite, service, app, deadline, None, None),
866 method,
867 params,
868 policy,
869 None,
870 on_event,
871 )
872 }
873
874 pub fn call_retrying_expecting(
880 config: &ClientConfig,
881 expect: &HostExpectation,
882 method: &str,
883 params: Value,
884 policy: RetryPolicy,
885 on_event: &mut dyn FnMut(&Value),
886 ) -> Result<Value, ServiceError> {
887 let suite = config.resolve_suite().map_err(suite_error)?;
888 Self::retry_loop(
889 &|deadline| {
890 Self::connect_bounded(
891 &suite,
892 &config.service,
893 &config.app,
894 deadline,
895 Some(expect.clone()),
896 config.secret.clone(),
897 )
898 },
899 method,
900 params,
901 policy,
902 None,
903 on_event,
904 )
905 }
906
907 pub fn call_retrying_over<F>(
914 connect: F,
915 method: &str,
916 params: Value,
917 policy: RetryPolicy,
918 on_event: &mut dyn FnMut(&Value),
919 ) -> Result<Value, ServiceError>
920 where
921 F: Fn(Option<Duration>) -> Result<Client, ConnectError>,
922 {
923 Self::retry_loop(&connect, method, params, policy, None, on_event)
924 }
925
926 pub fn call_retrying_until<F>(
933 connect: F,
934 method: &str,
935 params: Value,
936 policy: RetryPolicy,
937 stop: &(dyn Fn() -> bool + Sync),
938 on_event: &mut dyn FnMut(&Value),
939 ) -> Result<Value, ServiceError>
940 where
941 F: Fn(Option<Duration>) -> Result<Client, ConnectError>,
942 {
943 Self::retry_loop(&connect, method, params, policy, Some(stop), on_event)
944 }
945
946 fn retry_loop(
947 connect: &dyn Fn(Option<Duration>) -> Result<Client, ConnectError>,
948 method: &str,
949 params: Value,
950 policy: RetryPolicy,
951 stop: Option<&(dyn Fn() -> bool + Sync)>,
952 on_event: &mut dyn FnMut(&Value),
953 ) -> Result<Value, ServiceError> {
954 let started = Instant::now();
955 let stopped = || stop.is_some_and(|f| f());
956 let mut last = ServiceError::retryable(ERROR_UNAVAILABLE, "no attempt made");
957 for attempt in 0..policy.max_attempts.max(1) {
958 if attempt > 0 {
959 if stopped() || started.elapsed() > policy.give_up_after {
960 break;
961 }
962 sleep_unless_stopped(policy.backoff * attempt.min(8), stop);
963 if stopped() {
964 break;
965 }
966 }
967 let c = match connect(policy.request_deadline) {
968 Ok(c) => c,
969 Err(ConnectError::Service(e)) if e.retryable => {
970 last = e; continue;
972 }
973 Err(e) => return Err(e.into_service_error()),
974 };
975 match run_attempt(&c, method, ¶ms, stop, on_event) {
976 Ok(v) => return Ok(v),
977 Err(e)
980 if (e.retryable && (e.code == ERROR_HOST_HANDOFF || policy.idempotent))
981 || (e.code == ERROR_TIMEOUT && policy.idempotent) =>
982 {
983 last = e
984 }
985 Err(e) => return Err(e),
986 }
987 }
988 Err(last)
989 }
990}
991
992pub struct ClientPool {
1002 slots: Vec<Mutex<Option<Client>>>,
1003 next: AtomicUsize,
1004 connect: Box<dyn Fn(Option<Duration>) -> Result<Client, ConnectError> + Send + Sync>,
1005}
1006
1007impl ClientPool {
1008 pub fn new<F>(size: usize, connect: F) -> Self
1011 where
1012 F: Fn(Option<Duration>) -> Result<Client, ConnectError> + Send + Sync + 'static,
1013 {
1014 Self {
1015 slots: (0..size.max(1)).map(|_| Mutex::new(None)).collect(),
1016 next: AtomicUsize::new(0),
1017 connect: Box::new(connect),
1018 }
1019 }
1020
1021 pub fn discard_all(&self) {
1025 for slot in &self.slots {
1026 *slot.lock().unwrap_or_else(|e| e.into_inner()) = None;
1027 }
1028 }
1029
1030 fn pick(&self, deadline: Option<Instant>) -> Option<std::sync::MutexGuard<'_, Option<Client>>> {
1032 let n = self.slots.len();
1033 let start = self.next.fetch_add(1, Ordering::Relaxed);
1034 let mut pause = Duration::from_micros(200);
1035 loop {
1036 if deadline.is_some_and(|d| Instant::now() >= d) {
1037 return None;
1038 }
1039 for i in 0..n {
1040 match self.slots[start.wrapping_add(i) % n].try_lock() {
1041 Ok(g) => return Some(g),
1042 Err(std::sync::TryLockError::Poisoned(e)) => return Some(e.into_inner()),
1043 Err(std::sync::TryLockError::WouldBlock) => {}
1044 }
1045 }
1046 let wait = deadline.map_or(pause, |d| {
1047 pause.min(d.saturating_duration_since(Instant::now()))
1048 });
1049 std::thread::sleep(wait);
1050 pause = (pause * 2).min(Duration::from_millis(5));
1051 }
1052 }
1053
1054 pub fn call(&self, method: &str, params: Value) -> Result<Value, ServiceError> {
1057 self.call_with_policy(method, params, RetryPolicy::default(), &mut |_| {})
1058 }
1059
1060 pub fn call_with_policy(
1063 &self,
1064 method: &str,
1065 params: Value,
1066 policy: RetryPolicy,
1067 on_event: &mut dyn FnMut(&Value),
1068 ) -> Result<Value, ServiceError> {
1069 let started = Instant::now();
1070 let total_deadline = started.checked_add(policy.give_up_after);
1071 let bounded_attempt = |start: Instant| match (
1072 total_deadline,
1073 policy.request_deadline.and_then(|d| start.checked_add(d)),
1074 ) {
1075 (Some(total), Some(request)) => Some(total.min(request)),
1076 (total, request) => total.or(request),
1077 };
1078 let mut attempt_deadline = bounded_attempt(started);
1079 let Some(mut slot) = self.pick(attempt_deadline) else {
1080 return Err(timed_out("pool slot wait", policy.request_deadline));
1081 };
1082 let mut last = ServiceError::retryable(ERROR_UNAVAILABLE, "no attempt made");
1083 for attempt in 0..policy.max_attempts.max(1) {
1084 if attempt > 0 {
1085 let left = total_deadline.map(|d| d.saturating_duration_since(Instant::now()));
1086 if left.is_some_and(|d| d.is_zero()) {
1087 return Err(timed_out("pool call", Some(policy.give_up_after)));
1088 }
1089 let backoff = policy.backoff.saturating_mul(attempt.min(8));
1090 std::thread::sleep(left.map_or(backoff, |d| backoff.min(d)));
1091 attempt_deadline = bounded_attempt(Instant::now());
1092 }
1093 if attempt_deadline.is_some_and(|d| Instant::now() >= d) {
1094 return Err(timed_out("pool call", Some(policy.give_up_after)));
1095 }
1096 if slot.as_ref().is_some_and(Client::is_closed) {
1098 *slot = None;
1099 }
1100 if slot.is_none() {
1101 match (self.connect)(
1102 attempt_deadline.map(|d| d.saturating_duration_since(Instant::now())),
1103 ) {
1104 Ok(c) => *slot = Some(c),
1105 Err(ConnectError::Service(e)) if e.retryable => {
1106 last = e; continue;
1108 }
1109 Err(e) => return Err(e.into_service_error()),
1110 }
1111 }
1112 let result = match slot.as_ref() {
1113 Some(client) => client.call_large_until(
1114 method,
1115 params.clone(),
1116 Some(attempt_deadline),
1117 &mut *on_event,
1118 ),
1119 None => continue,
1120 };
1121 match result {
1122 Ok(v) => return Ok(v),
1123 Err(e) if e.code == ERROR_CONNECTION_CLOSED || e.code == ERROR_TIMEOUT => {
1126 *slot = None;
1127 if !policy.idempotent {
1128 return Err(e);
1129 }
1130 last = e;
1131 }
1132 Err(e) if e.retryable && e.code == ERROR_HOST_HANDOFF => {
1133 *slot = None;
1134 last = e;
1135 }
1136 Err(e) => return Err(e),
1137 }
1138 }
1139 Err(last)
1140 }
1141}
1142
1143#[cfg(test)]
1144mod tests {
1145 use super::*;
1146
1147 #[test]
1148 fn defaults_are_the_unhardened_resolve_with_the_default_deadline() {
1149 let cfg = ClientConfig::new("svc", "app");
1150 assert_eq!((cfg.service(), cfg.app()), ("svc", "app"));
1151 assert!(!cfg.ignores_env_overrides());
1152 assert_eq!(cfg.deadline, Some(DEFAULT_DEADLINE));
1153 }
1154
1155 #[test]
1156 fn an_explicit_suite_is_used_as_is_unless_it_came_from_the_environment_and_the_switch_is_on() {
1157 let dir = tempfile::tempdir().unwrap();
1158 let explicit = Suite::at(dir.path()).unwrap();
1159 for ignore in [false, true] {
1160 let cfg = ClientConfig::new("svc", "app")
1161 .suite(explicit.clone())
1162 .ignore_env_overrides(ignore);
1163 assert_eq!(cfg.resolve_suite().unwrap().root(), explicit.root());
1164 }
1165 #[cfg(feature = "test-mode")]
1167 if let Ok(from_env) = Suite::resolve() {
1168 let lenient = ClientConfig::new("svc", "app").suite(from_env);
1169 assert!(lenient.resolve_suite().is_ok());
1170 let strict = lenient.ignore_env_overrides(true);
1171 assert!(matches!(
1172 strict.resolve_suite(),
1173 Err(SuiteError::Insecure(_))
1174 ));
1175 }
1176 }
1177
1178 struct Echo;
1179
1180 impl crate::server::Handler for Echo {
1181 fn call(
1182 &self,
1183 _cx: &crate::server::CallContext,
1184 method: &str,
1185 params: &Value,
1186 _emit: &crate::server::Emit,
1187 ) -> Result<Value, ServiceError> {
1188 match method {
1189 "echo" => Ok(params.clone()),
1190 _ => Err(ServiceError::new("unknown_method", method)),
1191 }
1192 }
1193 }
1194
1195 #[test]
1196 fn the_pool_reuses_a_healthy_client_discards_a_closed_one_and_retries_a_failed_connect() {
1197 use crate::record::Allowlist;
1198 use crate::server::{host, ServiceConfig};
1199 use std::sync::atomic::AtomicU32;
1200 let dir = tempfile::tempdir().unwrap();
1201 let suite = Suite::at(dir.path().join("suite")).unwrap();
1202 let config = ServiceConfig::new(
1203 "pool-svc",
1204 "engine",
1205 1,
1206 "1.0",
1207 Allowlist::new().allow("app", ["echo"]),
1208 );
1209 let _hosted = host(&suite, config, Arc::new(Echo)).unwrap();
1210 let connects = Arc::new(AtomicU32::new(0));
1211 let counter = connects.clone();
1212 let pool_suite = suite.clone();
1213 let pool = ClientPool::new(1, move |deadline| {
1214 if counter.fetch_add(1, Ordering::SeqCst) == 0 {
1217 return Err(ConnectError::Service(ServiceError::retryable(
1218 ERROR_UNAVAILABLE,
1219 "not up yet",
1220 )));
1221 }
1222 Client::connect_bounded(&pool_suite, "pool-svc", "app", deadline, None, None)
1223 });
1224 let policy = RetryPolicy {
1225 max_attempts: 3,
1226 backoff: Duration::from_millis(1),
1227 give_up_after: Duration::from_secs(30),
1228 request_deadline: Some(Duration::from_secs(10)),
1229 idempotent: false,
1230 };
1231 let first = pool
1232 .call_with_policy("echo", json!({"n": 1}), policy, &mut |_| {})
1233 .unwrap();
1234 assert_eq!(first, json!({"n": 1}));
1235 assert_eq!(
1236 connects.load(Ordering::SeqCst),
1237 2,
1238 "one failure, one connect"
1239 );
1240 assert_eq!(pool.call("echo", json!({"n": 2})).unwrap(), json!({"n": 2}));
1242 assert_eq!(connects.load(Ordering::SeqCst), 2);
1243 {
1246 let slot = pool.slots[0].lock().unwrap();
1247 slot.as_ref().unwrap().abort();
1248 assert!(slot.as_ref().unwrap().is_closed());
1249 }
1250 assert_eq!(pool.call("echo", json!({"n": 3})).unwrap(), json!({"n": 3}));
1251 assert_eq!(
1252 connects.load(Ordering::SeqCst),
1253 3,
1254 "reconnected after discard"
1255 );
1256 pool.discard_all();
1258 assert_eq!(pool.call("echo", json!({"n": 4})).unwrap(), json!({"n": 4}));
1259 assert_eq!(connects.load(Ordering::SeqCst), 4);
1260 }
1261
1262 struct PoolGate {
1263 entered: mpsc::Sender<()>,
1264 release: Mutex<mpsc::Receiver<()>>,
1265 calls: AtomicUsize,
1266 }
1267
1268 impl crate::server::Handler for PoolGate {
1269 fn call(
1270 &self,
1271 _cx: &crate::server::CallContext,
1272 method: &str,
1273 params: &Value,
1274 _emit: &crate::server::Emit,
1275 ) -> Result<Value, ServiceError> {
1276 if method == "block" {
1277 self.calls.fetch_add(1, Ordering::SeqCst);
1278 self.entered.send(()).unwrap();
1279 let _ = self
1280 .release
1281 .lock()
1282 .unwrap()
1283 .recv_timeout(Duration::from_secs(3));
1284 }
1285 Ok(params.clone())
1286 }
1287 }
1288
1289 #[test]
1290 fn pooled_reuse_obeys_current_short_deadline_without_replaying_uncertain_request() {
1291 use crate::record::Allowlist;
1292 use crate::server::{host, ServiceConfig};
1293 for (previous, request_deadline, give_up_after) in [
1294 (
1295 None,
1296 Some(Duration::from_millis(50)),
1297 Duration::from_secs(2),
1298 ),
1299 (
1300 Some(Duration::from_secs(60)),
1301 Some(Duration::from_millis(50)),
1302 Duration::from_secs(2),
1303 ),
1304 (None, None, Duration::from_millis(50)),
1305 ] {
1306 let dir = tempfile::tempdir().unwrap();
1307 let suite = Suite::at(dir.path().join("suite")).unwrap();
1308 let (entered, _entered_rx) = mpsc::channel();
1309 let (release_tx, release) = mpsc::channel();
1310 let handler = Arc::new(PoolGate {
1311 entered,
1312 release: Mutex::new(release),
1313 calls: AtomicUsize::new(0),
1314 });
1315 let config = ServiceConfig::new(
1316 "pool-budget",
1317 "engine",
1318 1,
1319 "1.0",
1320 Allowlist::new().allow("app", ["echo", "block"]),
1321 );
1322 let _hosted = host(&suite, config, handler.clone()).unwrap();
1323 let connects = Arc::new(AtomicUsize::new(0));
1324 let counter = connects.clone();
1325 let pool = ClientPool::new(1, move |_| {
1326 counter.fetch_add(1, Ordering::SeqCst);
1327 Client::connect_bounded(&suite, "pool-budget", "app", previous, None, None)
1328 });
1329 pool.call_with_policy(
1330 "echo",
1331 Value::Null,
1332 RetryPolicy {
1333 request_deadline: previous,
1334 ..RetryPolicy::default()
1335 },
1336 &mut |_| {},
1337 )
1338 .unwrap();
1339 let started = Instant::now();
1340 let result = pool.call_with_policy(
1341 "block",
1342 Value::Null,
1343 RetryPolicy {
1344 request_deadline,
1345 give_up_after,
1346 max_attempts: 3,
1347 idempotent: false,
1348 ..RetryPolicy::default()
1349 },
1350 &mut |_| {},
1351 );
1352 let elapsed = started.elapsed();
1353 release_tx.send(()).unwrap();
1354 assert_eq!(result.unwrap_err().code, ERROR_TIMEOUT);
1355 assert!(elapsed < Duration::from_secs(1));
1356 assert_eq!(connects.load(Ordering::SeqCst), 1);
1357 assert_eq!(handler.calls.load(Ordering::SeqCst), 1);
1358 }
1359 }
1360
1361 #[test]
1362 fn busy_pool_wait_is_bounded_without_aborting_active_stream() {
1363 use crate::record::Allowlist;
1364 use crate::server::{host, ServiceConfig};
1365 let dir = tempfile::tempdir().unwrap();
1366 let suite = Suite::at(dir.path().join("suite")).unwrap();
1367 let (entered, entered_rx) = mpsc::channel();
1368 let (release_tx, release) = mpsc::channel();
1369 let handler = Arc::new(PoolGate {
1370 entered,
1371 release: Mutex::new(release),
1372 calls: AtomicUsize::new(0),
1373 });
1374 let config = ServiceConfig::new(
1375 "pool-busy",
1376 "engine",
1377 1,
1378 "1.0",
1379 Allowlist::new().allow("app", ["echo", "block"]),
1380 );
1381 let _hosted = host(&suite, config, handler).unwrap();
1382 let pool = Arc::new(ClientPool::new(1, move |deadline| {
1383 Client::connect_bounded(&suite, "pool-busy", "app", deadline, None, None)
1384 }));
1385 let active_pool = pool.clone();
1386 let active = std::thread::spawn(move || active_pool.call("block", json!("original")));
1387 entered_rx.recv_timeout(Duration::from_secs(2)).unwrap();
1388 for request_deadline in [None, Some(Duration::from_millis(30))] {
1389 let started = Instant::now();
1390 let result = pool.call_with_policy(
1391 "echo",
1392 Value::Null,
1393 RetryPolicy {
1394 request_deadline,
1395 give_up_after: Duration::from_millis(60),
1396 ..RetryPolicy::default()
1397 },
1398 &mut |_| {},
1399 );
1400 assert_eq!(result.unwrap_err().code, ERROR_TIMEOUT);
1401 assert!(started.elapsed() < Duration::from_secs(1));
1402 }
1403 release_tx.send(()).unwrap();
1404 assert_eq!(active.join().unwrap().unwrap(), json!("original"));
1405 assert_eq!(
1406 pool.call("echo", json!("healthy")).unwrap(),
1407 json!("healthy")
1408 );
1409 }
1410
1411 #[test]
1412 fn pool_retry_backoff_cannot_exceed_total_budget() {
1413 let connects = Arc::new(AtomicUsize::new(0));
1414 let counter = connects.clone();
1415 let pool = ClientPool::new(1, move |_| {
1416 counter.fetch_add(1, Ordering::SeqCst);
1417 Err(ConnectError::Service(ServiceError::retryable(
1418 ERROR_UNAVAILABLE,
1419 "offline",
1420 )))
1421 });
1422 let started = Instant::now();
1423 let result = pool.call_with_policy(
1424 "echo",
1425 Value::Null,
1426 RetryPolicy {
1427 give_up_after: Duration::from_millis(40),
1428 backoff: Duration::from_secs(10),
1429 max_attempts: 3,
1430 request_deadline: None,
1431 ..RetryPolicy::default()
1432 },
1433 &mut |_| {},
1434 );
1435 assert_eq!(result.unwrap_err().code, ERROR_TIMEOUT);
1436 assert!(started.elapsed() < Duration::from_secs(1));
1437 assert_eq!(connects.load(Ordering::SeqCst), 1);
1438 }
1439
1440 #[test]
1441 fn a_missing_record_is_a_retryable_service_error_not_an_identity_error() {
1442 let dir = tempfile::tempdir().unwrap();
1443 let cfg = ClientConfig::new("nobody-home", "app").suite(Suite::at(dir.path()).unwrap());
1444 let err = Client::connect_expecting(&cfg, &HostExpectation::executable("/x"))
1445 .err()
1446 .expect("there is no host");
1447 match err {
1448 ConnectError::Service(e) => {
1449 assert_eq!(e.code, ERROR_UNAVAILABLE);
1450 assert!(e.retryable);
1451 }
1452 other => panic!("expected a service error, got {other:?}"),
1453 }
1454 }
1455}