1use crate::error::*;
2use crate::host_identity::{HostConnectError, HostExpectation, HostProcess};
3use crate::record::{try_read_record, ServiceRecord, PROTOCOL_MAJOR};
4use crate::suite::{Suite, SuiteError};
5use crate::transport::{connect, connect_expecting, Conn};
6use rightkit_sidecar::{read_json_line, write_json_line, FrameError};
7use serde_json::{json, Value};
8use std::io::{BufReader, Read, Write};
9use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
10use std::sync::{mpsc, Arc, Mutex};
11use std::time::{Duration, Instant};
12
13const LIMIT: usize = 16 * 1024 * 1024;
14
15pub const DEFAULT_DEADLINE: Duration = Duration::from_secs(120);
18
19struct Watchdog {
23 fired: Arc<AtomicBool>,
24 _disarm: mpsc::Sender<()>,
25}
26
27impl Watchdog {
28 fn arm(after: Option<Duration>, shutdown: Arc<dyn Fn() + Send + Sync>) -> Watchdog {
29 let fired = Arc::new(AtomicBool::new(false));
30 let (tx, rx) = mpsc::channel::<()>();
31 if let Some(d) = after {
32 let f = fired.clone();
33 std::thread::spawn(move || {
34 if let Err(mpsc::RecvTimeoutError::Timeout) = rx.recv_timeout(d) {
35 f.store(true, Ordering::SeqCst);
36 shutdown();
37 }
38 });
39 }
40 Watchdog { fired, _disarm: tx }
41 }
42 fn fired(&self) -> bool {
43 self.fired.load(Ordering::SeqCst)
44 }
45}
46
47fn timed_out(what: &str, d: Option<Duration>) -> ServiceError {
48 ServiceError::new(
49 ERROR_TIMEOUT,
50 format!(
51 "{what} exceeded the {} ms deadline; connection aborted",
52 d.map(|d| d.as_millis()).unwrap_or(0)
53 ),
54 )
55}
56
57fn lock_until<T>(m: &Mutex<T>, deadline: Option<Instant>) -> Option<std::sync::MutexGuard<'_, T>> {
60 use std::sync::TryLockError;
61 let Some(deadline) = deadline else {
62 return Some(m.lock().unwrap_or_else(|e| e.into_inner()));
63 };
64 let mut pause = Duration::from_micros(200);
65 loop {
66 match m.try_lock() {
67 Ok(g) => return Some(g),
68 Err(TryLockError::Poisoned(e)) => return Some(e.into_inner()),
69 Err(TryLockError::WouldBlock) => {}
70 }
71 let left = deadline.saturating_duration_since(Instant::now());
72 if left.is_zero() {
73 return None;
74 }
75 std::thread::sleep(pause.min(left));
76 pause = (pause * 2).min(Duration::from_millis(5));
77 }
78}
79
80struct Io {
81 reader: BufReader<Box<dyn Read + Send>>,
82 writer: Box<dyn Write + Send>,
83 buf: Vec<u8>,
84}
85
86#[derive(Debug, Clone)]
105pub struct ClientConfig {
106 service: String,
107 app: String,
108 suite: Option<Suite>,
109 deadline: Option<Duration>,
110 ignore_env_overrides: bool,
111}
112
113impl ClientConfig {
114 pub fn new(service: impl Into<String>, app: impl Into<String>) -> Self {
118 Self {
119 service: service.into(),
120 app: app.into(),
121 suite: None,
122 deadline: Some(DEFAULT_DEADLINE),
123 ignore_env_overrides: false,
124 }
125 }
126
127 pub fn suite(mut self, suite: Suite) -> Self {
129 self.suite = Some(suite);
130 self
131 }
132
133 pub fn deadline(mut self, deadline: Option<Duration>) -> Self {
135 self.deadline = deadline;
136 self
137 }
138
139 pub fn ignore_env_overrides(mut self, ignore: bool) -> Self {
144 self.ignore_env_overrides = ignore;
145 self
146 }
147
148 pub fn service(&self) -> &str {
149 &self.service
150 }
151
152 pub fn app(&self) -> &str {
153 &self.app
154 }
155
156 pub fn ignores_env_overrides(&self) -> bool {
158 self.ignore_env_overrides
159 }
160
161 #[allow(deprecated)] pub fn resolve_suite(&self) -> Result<Suite, SuiteError> {
164 match (&self.suite, self.ignore_env_overrides) {
165 (Some(suite), true) if suite.env_derived() => Err(SuiteError::Insecure(format!(
166 "the suite {} was derived from the environment, which this client ignores; \
167 pass Suite::at(..) or leave the suite unset",
168 suite.root().display()
169 ))),
170 (Some(suite), _) => Ok(suite.clone()),
171 (None, ignore) => Suite::resolve_ignoring_env_overrides(ignore),
172 }
173 }
174}
175
176fn suite_error(e: SuiteError) -> ServiceError {
177 ServiceError::new(ERROR_SUITE_UNRESOLVED, e.to_string())
178}
179
180pub struct Client {
185 io: Mutex<Io>,
186 next_id: AtomicU64,
187 dead: AtomicBool,
188 hello: Value,
189 record: ServiceRecord,
190 shutdown: Arc<dyn Fn() + Send + Sync>,
191 deadline_ms: AtomicU64,
193 host: Option<HostProcess>,
195}
196
197fn closed() -> ServiceError {
198 ServiceError::retryable(
199 ERROR_CONNECTION_CLOSED,
200 "connection closed before a response arrived",
201 )
202}
203
204impl Client {
205 pub fn connect(suite: &Suite, service: &str, app: &str) -> Result<Client, ServiceError> {
215 Self::connect_with(suite, service, app, Some(DEFAULT_DEADLINE))
216 }
217
218 pub fn connect_with(
222 suite: &Suite,
223 service: &str,
224 app: &str,
225 deadline: Option<Duration>,
226 ) -> Result<Client, ServiceError> {
227 Self::connect_bounded(suite, service, app, deadline, None)
228 .map_err(ConnectError::into_service_error)
229 }
230
231 pub fn connect_config(config: &ClientConfig) -> Result<Client, ServiceError> {
235 let suite = config.resolve_suite().map_err(suite_error)?;
236 Self::connect_bounded(&suite, &config.service, &config.app, config.deadline, None)
237 .map_err(ConnectError::into_service_error)
238 }
239
240 pub fn connect_expecting(
263 config: &ClientConfig,
264 expect: &HostExpectation,
265 ) -> Result<Client, ConnectError> {
266 let suite = config
267 .resolve_suite()
268 .map_err(|e| ConnectError::Service(suite_error(e)))?;
269 Self::connect_bounded(
270 &suite,
271 &config.service,
272 &config.app,
273 config.deadline,
274 Some(expect.clone()),
275 )
276 }
277
278 fn connect_bounded(
282 suite: &Suite,
283 service: &str,
284 app: &str,
285 deadline: Option<Duration>,
286 expect: Option<HostExpectation>,
287 ) -> Result<Client, ConnectError> {
288 let Some(d) = deadline else {
289 return Self::handshake(suite, service, app, None, expect.as_ref());
290 };
291 let (tx, rx) = mpsc::channel();
292 let (s, sv, ap) = (suite.clone(), service.to_string(), app.to_string());
293 std::thread::spawn(move || {
294 let _ = tx.send(Self::handshake(&s, &sv, &ap, Some(d), expect.as_ref()));
295 });
296 match rx.recv_timeout(d + Duration::from_millis(250)) {
297 Ok(r) => r,
298 Err(_) => Err(ConnectError::Service(timed_out("connect", deadline))),
299 }
300 }
301
302 fn handshake(
303 suite: &Suite,
304 service: &str,
305 app: &str,
306 deadline: Option<Duration>,
307 expect: Option<&HostExpectation>,
308 ) -> Result<Client, ConnectError> {
309 let started = Instant::now();
310 let record = match try_read_record(suite, service) {
313 Ok(Some(record)) => record,
314 Ok(None) => {
315 return Err(ServiceError::retryable(
316 ERROR_UNAVAILABLE,
317 format!("no service.json for {service}"),
318 )
319 .into())
320 }
321 Err(e) => {
322 return Err(ServiceError::retryable(
323 ERROR_UNAVAILABLE,
324 format!("could not read service.json for {service} (try again): {e}"),
325 )
326 .into())
327 }
328 };
329 if record.protocol_major != PROTOCOL_MAJOR {
330 return Err(ServiceError::new(
331 "protocol_mismatch",
332 format!("host speaks protocol {}", record.protocol_major),
333 )
334 .into());
335 }
336 let unreachable_host = |e: std::io::Error| {
337 ServiceError::retryable(
338 ERROR_UNAVAILABLE,
339 format!("{service} host unreachable: {e}"),
340 )
341 };
342 let (conn, host) = match expect {
344 None => (connect(&record.endpoint).map_err(unreachable_host)?, None),
345 Some(expect) => match connect_expecting(&record.endpoint, expect) {
346 Ok((conn, host)) => (conn, Some(host)),
347 Err(HostConnectError::Io(e)) => return Err(unreachable_host(e).into()),
348 Err(HostConnectError::Identity(e)) => return Err(ConnectError::HostIdentity(e)),
349 },
350 };
351 let Conn {
352 reader,
353 writer,
354 shutdown,
355 ..
356 } = conn;
357 let mut io = Io {
358 reader: BufReader::new(reader),
359 writer,
360 buf: Vec::new(),
361 };
362 let wd = Watchdog::arm(
363 deadline.map(|d| d.saturating_sub(started.elapsed())),
364 shutdown.clone(),
365 );
366 let hello_res = (|| {
367 write_json_line(
368 &mut io.writer,
369 &json!({"type": "hello", "app": app, "protocol_major": PROTOCOL_MAJOR}),
370 LIMIT,
371 )
372 .map_err(|e| ServiceError::new("io", e.to_string()))?;
373 io.writer.flush()?;
374 read_json_line::<_, Value>(&mut io.reader, &mut io.buf, LIMIT)
375 .map_err(|e| ServiceError::new("io", e.to_string()))?
376 .ok_or_else(closed)
377 })();
378 let hello = match hello_res {
379 Ok(h) => h,
380 Err(e) => {
381 let t = wd.fired();
382 (shutdown)();
383 let e = if t { timed_out("hello", deadline) } else { e };
384 return Err(e.into());
385 }
386 };
387 drop(wd);
388 if hello["ok"] != true {
389 return Err(ServiceError::from_value(&hello["error"]).into());
390 }
391 Ok(Client {
392 io: Mutex::new(io),
393 next_id: AtomicU64::new(1),
394 dead: AtomicBool::new(false),
395 hello,
396 record,
397 shutdown,
398 deadline_ms: AtomicU64::new(deadline.map(|d| d.as_millis() as u64).unwrap_or(0)),
399 host,
400 })
401 }
402
403 pub fn verified_host(&self) -> Option<&HostProcess> {
407 self.host.as_ref()
408 }
409
410 pub fn set_deadline(&self, deadline: Option<Duration>) {
412 self.deadline_ms.store(
413 deadline.map(|d| (d.as_millis() as u64).max(1)).unwrap_or(0),
414 Ordering::SeqCst,
415 );
416 }
417
418 pub fn record(&self) -> &ServiceRecord {
419 &self.record
420 }
421 pub fn hello(&self) -> &Value {
423 &self.hello
424 }
425
426 pub fn abort(&self) {
429 self.dead.store(true, Ordering::SeqCst);
430 (self.shutdown)();
431 }
432
433 pub fn call(&self, method: &str, params: Value) -> Result<Value, ServiceError> {
434 self.call_events(method, params, &mut |_| {})
435 }
436
437 pub fn call_events(
442 &self,
443 method: &str,
444 params: Value,
445 on_event: &mut dyn FnMut(&Value),
446 ) -> Result<Value, ServiceError> {
447 let deadline = match self.deadline_ms.load(Ordering::SeqCst) {
448 0 => None,
449 ms => Some(Instant::now() + Duration::from_millis(ms)),
450 };
451 self.call_events_until(method, params, deadline, on_event)
452 }
453
454 pub fn call_events_until(
459 &self,
460 method: &str,
461 params: Value,
462 deadline: Option<Instant>,
463 on_event: &mut dyn FnMut(&Value),
464 ) -> Result<Value, ServiceError> {
465 let started = Instant::now();
466 let total = deadline.map(|d| d.saturating_duration_since(started));
467 if self.dead.load(Ordering::SeqCst) {
468 return Err(closed());
469 }
470 let id = self.next_id.fetch_add(1, Ordering::SeqCst);
471 let Some(mut io) = lock_until(&self.io, deadline) else {
472 return Err(timed_out(
473 &format!("request `{method}` waiting for the connection"),
474 total,
475 ));
476 };
477 if self.dead.load(Ordering::SeqCst) {
478 return Err(closed());
479 }
480 let wd = Watchdog::arm(
481 deadline.map(|d| d.saturating_duration_since(Instant::now())),
482 self.shutdown.clone(),
483 );
484 let res = (|| {
485 write_json_line(
486 &mut io.writer,
487 &json!({"type": "request", "id": id, "method": method, "params": params}),
488 LIMIT,
489 )
490 .map_err(|e| match e {
491 FrameError::TooLarge { .. } => {
492 ServiceError::new("request_too_large", "request exceeds the frame limit")
493 }
494 _ => closed(),
495 })?;
496 io.writer.flush().map_err(|_| closed())?;
497 loop {
498 let Io { reader, buf, .. } = &mut *io;
499 match read_json_line::<_, Value>(reader, buf, LIMIT) {
500 Ok(Some(v)) if v["id"].as_u64() == Some(id) => match v["type"].as_str() {
501 Some("event") => on_event(&v["event"]),
502 Some("response") => {
503 return if v["ok"] == true {
504 Ok(v["result"].clone())
505 } else {
506 Err(ServiceError::from_value(&v["error"]))
507 };
508 }
509 _ => {}
510 },
511 Ok(Some(_)) => {} Ok(None) | Err(FrameError::Io(_)) => return Err(closed()),
513 Err(FrameError::TooLarge { .. }) => {
517 return Err(ServiceError::new(
518 ERROR_RESPONSE_TOO_LARGE,
519 "the host's reply exceeds the client's frame limit",
520 ));
521 }
522 Err(_) => {}
523 }
524 }
525 })();
526 if wd.fired() && res.is_err() {
527 self.dead.store(true, Ordering::SeqCst);
528 return Err(timed_out(&format!("request `{method}`"), total));
529 }
530 if matches!(&res, Err(e) if e.code == ERROR_CONNECTION_CLOSED) {
531 self.dead.store(true, Ordering::SeqCst);
532 }
533 res
534 }
535}
536
537#[derive(Debug, Clone, Copy)]
539pub struct RetryPolicy {
540 pub max_attempts: u32,
541 pub backoff: Duration,
542 pub give_up_after: Duration,
545 pub request_deadline: Option<Duration>,
547 pub idempotent: bool,
552}
553
554impl Default for RetryPolicy {
555 fn default() -> Self {
556 Self {
557 max_attempts: 8,
558 backoff: Duration::from_millis(150),
559 give_up_after: Duration::from_secs(120),
560 request_deadline: Some(DEFAULT_DEADLINE),
561 idempotent: false,
562 }
563 }
564}
565
566impl Client {
567 pub fn call_retrying(
571 suite: &Suite,
572 service: &str,
573 app: &str,
574 method: &str,
575 params: Value,
576 policy: RetryPolicy,
577 on_event: &mut dyn FnMut(&Value),
578 ) -> Result<Value, ServiceError> {
579 Self::retry_loop(
580 &|deadline| Self::connect_bounded(suite, service, app, deadline, None),
581 method,
582 params,
583 policy,
584 on_event,
585 )
586 }
587
588 pub fn call_retrying_expecting(
594 config: &ClientConfig,
595 expect: &HostExpectation,
596 method: &str,
597 params: Value,
598 policy: RetryPolicy,
599 on_event: &mut dyn FnMut(&Value),
600 ) -> Result<Value, ServiceError> {
601 let suite = config.resolve_suite().map_err(suite_error)?;
602 Self::retry_loop(
603 &|deadline| {
604 Self::connect_bounded(
605 &suite,
606 &config.service,
607 &config.app,
608 deadline,
609 Some(expect.clone()),
610 )
611 },
612 method,
613 params,
614 policy,
615 on_event,
616 )
617 }
618
619 fn retry_loop(
620 connect: &dyn Fn(Option<Duration>) -> Result<Client, ConnectError>,
621 method: &str,
622 params: Value,
623 policy: RetryPolicy,
624 on_event: &mut dyn FnMut(&Value),
625 ) -> Result<Value, ServiceError> {
626 let started = Instant::now();
627 let mut last = ServiceError::retryable(ERROR_UNAVAILABLE, "no attempt made");
628 for attempt in 0..policy.max_attempts.max(1) {
629 if attempt > 0 {
630 if started.elapsed() > policy.give_up_after {
631 break;
632 }
633 std::thread::sleep(policy.backoff * attempt.min(8));
634 }
635 let c = match connect(policy.request_deadline) {
636 Ok(c) => c,
637 Err(ConnectError::Service(e)) if e.retryable => {
638 last = e; continue;
640 }
641 Err(e) => return Err(e.into_service_error()),
642 };
643 match c.call_events(method, params.clone(), on_event) {
644 Ok(v) => return Ok(v),
645 Err(e) if e.retryable && (e.code == ERROR_HOST_HANDOFF || policy.idempotent) => {
647 last = e
648 }
649 Err(e) => return Err(e),
650 }
651 }
652 Err(last)
653 }
654}
655
656#[cfg(test)]
657mod tests {
658 use super::*;
659
660 #[test]
661 fn defaults_are_the_unhardened_resolve_with_the_default_deadline() {
662 let cfg = ClientConfig::new("svc", "app");
663 assert_eq!((cfg.service(), cfg.app()), ("svc", "app"));
664 assert!(!cfg.ignores_env_overrides());
665 assert_eq!(cfg.deadline, Some(DEFAULT_DEADLINE));
666 }
667
668 #[test]
669 fn an_explicit_suite_is_used_as_is_unless_it_came_from_the_environment_and_the_switch_is_on() {
670 let dir = tempfile::tempdir().unwrap();
671 let explicit = Suite::at(dir.path()).unwrap();
672 for ignore in [false, true] {
673 let cfg = ClientConfig::new("svc", "app")
674 .suite(explicit.clone())
675 .ignore_env_overrides(ignore);
676 assert_eq!(cfg.resolve_suite().unwrap().root(), explicit.root());
677 }
678 #[cfg(feature = "test-mode")]
680 if let Ok(from_env) = Suite::resolve() {
681 let lenient = ClientConfig::new("svc", "app").suite(from_env);
682 assert!(lenient.resolve_suite().is_ok());
683 let strict = lenient.ignore_env_overrides(true);
684 assert!(matches!(
685 strict.resolve_suite(),
686 Err(SuiteError::Insecure(_))
687 ));
688 }
689 }
690
691 #[test]
692 fn a_missing_record_is_a_retryable_service_error_not_an_identity_error() {
693 let dir = tempfile::tempdir().unwrap();
694 let cfg = ClientConfig::new("nobody-home", "app").suite(Suite::at(dir.path()).unwrap());
695 let err = Client::connect_expecting(&cfg, &HostExpectation::executable("/x"))
696 .err()
697 .expect("there is no host");
698 match err {
699 ConnectError::Service(e) => {
700 assert_eq!(e.code, ERROR_UNAVAILABLE);
701 assert!(e.retryable);
702 }
703 other => panic!("expected a service error, got {other:?}"),
704 }
705 }
706}