1use std::collections::HashMap;
2use std::collections::hash_map::Entry;
3use std::convert::Infallible;
4use std::fmt;
5use std::sync::{Arc, Mutex, MutexGuard, PoisonError};
6use std::time::Duration;
7
8use http::HeaderMap;
9use tokio::sync::watch;
10
11use crate::builder::{ClientBuilder, Config};
12use crate::error::Error;
13use crate::protocol::{Command, Incoming, Kind, Protocol, SUBPROTOCOL_UNSUPPORTED};
14use crate::subscription::{Decision, Receivers, Registration, SubscribeRequest, Subscription};
15use crate::transport::{Conn, DialOptions};
16
17#[derive(Clone)]
25pub struct Client {
26 inner: Arc<Inner>,
27}
28
29impl Client {
30 pub fn builder(url: impl Into<String>) -> ClientBuilder {
33 ClientBuilder::new(url.into())
34 }
35
36 pub(crate) fn new(config: Config) -> Client {
37 Client {
38 inner: Arc::new(Inner {
39 config,
40 state: Mutex::new(State::default()),
41 write_lock: tokio::sync::Mutex::new(()),
42 connected: watch::Sender::new(false),
43 stopping: watch::Sender::new(false),
44 done: watch::Sender::new(false),
45 }),
46 }
47 }
48
49 pub async fn connect(&self) -> Result<(), Error> {
62 self.inner.state().start()?;
63
64 tokio::spawn(Arc::clone(&self.inner).run());
65
66 let mut connected = self.inner.connected.subscribe();
67 let mut done = self.inner.done.subscribe();
68 tokio::select! {
69 biased;
70 _ = connected.wait_for(|connected| *connected) => Ok(()),
71 _ = done.wait_for(|done| *done) => Err(self.inner.state().failure_or_closed()),
72 }
73 }
74
75 pub fn connected(&self) -> bool {
77 let state = self.inner.state();
78 state.welcomed && state.conn.is_some()
79 }
80
81 pub async fn done(&self) {
87 let mut done = self.inner.done.subscribe();
88 let _stopped = done.wait_for(|done| *done).await;
89 }
90
91 pub fn err(&self) -> Option<Error> {
94 let state = self.inner.state();
95 if state.stopped {
96 Some(state.failure_or_closed())
97 } else {
98 None
99 }
100 }
101
102 pub fn last_error(&self) -> Option<Error> {
106 self.inner.state().last_error.clone()
107 }
108
109 pub async fn subscribe(
125 &self,
126 request: impl Into<SubscribeRequest>,
127 ) -> Result<Subscription, Error> {
128 let (identifier, callbacks) = request.into().into_parts();
129 let key = identifier.key();
130 let (registration, mut receivers, joined) = self
131 .inner
132 .state()
133 .register(key.clone(), self.inner.config.message_buffer)?;
134 let mut forgetting = Forgetting::new(Arc::clone(&self.inner), Arc::clone(®istration));
135 let running = callbacks.spawn(registration.watch_events());
136
137 match joined {
138 Joined::New => self.inner.send_subscribe(&key).await,
139 Joined::Pending => {}
140 Joined::Confirmed { reconnected } => registration.confirm(reconnected),
141 }
142
143 let mut done = self.inner.done.subscribe();
144 let verdict = tokio::select! {
145 biased;
146 decision = &mut receivers.decision => decision.ok(),
147 _ = done.wait_for(|done| *done) => None,
148 };
149
150 match verdict {
151 Some(Decision::Confirmed) => {
152 forgetting.disarm();
153 Ok(Subscription::new(
154 Arc::downgrade(&self.inner),
155 registration,
156 receivers,
157 running,
158 ))
159 }
160 Some(Decision::Rejected) => {
161 forgetting.disarm();
162 Err(Error::Rejected { key })
163 }
164 None => Err(self.inner.state().failure_or_closed()),
165 }
166 }
167
168 pub async fn close(&self) {
171 let started = {
172 let mut state = self.inner.state();
173 state.stopped = true;
174 state.failure.get_or_insert(Error::Closed);
175 state.started
176 };
177
178 if started {
179 self.inner.stopping.send_replace(true);
180 } else {
181 self.inner.done.send_replace(true);
182 }
183 self.done().await;
184 }
185}
186
187impl fmt::Debug for Client {
188 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
189 f.debug_struct("Client")
190 .field("url", &self.inner.config.url)
191 .field("connected", &self.connected())
192 .finish()
193 }
194}
195
196struct Forgetting {
201 inner: Arc<Inner>,
202 registration: Arc<Registration>,
203 armed: bool,
204}
205
206impl Forgetting {
207 fn new(inner: Arc<Inner>, registration: Arc<Registration>) -> Forgetting {
208 Forgetting {
209 inner,
210 registration,
211 armed: true,
212 }
213 }
214
215 fn disarm(&mut self) {
216 self.armed = false;
217 }
218}
219
220impl Drop for Forgetting {
221 fn drop(&mut self) {
222 if self.armed && self.inner.forget(&self.registration, &Error::Unsubscribed) {
223 self.inner.unsubscribe_later(&self.registration.key);
224 }
225 }
226}
227
228pub(crate) struct Inner {
229 config: Config,
230 state: Mutex<State>,
231 write_lock: tokio::sync::Mutex<()>,
234 connected: watch::Sender<bool>,
236 stopping: watch::Sender<bool>,
238 done: watch::Sender<bool>,
240}
241
242#[derive(Default)]
243#[allow(
244 clippy::struct_excessive_bools,
245 reason = "a connection is a handful of independent facts, and an enum of their combinations would be longer than the struct"
246)]
247struct State {
248 conn: Option<Arc<dyn Conn>>,
249 protocol: Option<Arc<dyn Protocol>>,
251 channels: HashMap<String, Channel>,
252 attempts: u32,
253 last_error: Option<Error>,
256 reconnected: bool,
257 welcomed: bool,
258 ever_welcomed: bool,
259 started: bool,
260 stopped: bool,
261 failure: Option<Error>,
262 stop_announced: bool,
265}
266
267struct Channel {
271 confirmed: bool,
272 holders: Vec<Arc<Registration>>,
273}
274
275type Writer = (Arc<dyn Conn>, Arc<dyn Protocol>);
278
279enum Joined {
282 New,
283 Pending,
284 Confirmed { reconnected: bool },
285}
286
287impl State {
288 fn start(&mut self) -> Result<(), Error> {
289 if self.stopped {
290 Err(self.failure_or_closed())
291 } else if self.started {
292 Err(Error::AlreadyConnected)
293 } else {
294 self.started = true;
295 Ok(())
296 }
297 }
298
299 fn register(
300 &mut self,
301 key: String,
302 buffer: usize,
303 ) -> Result<(Arc<Registration>, Receivers, Joined), Error> {
304 if self.stopped {
305 Err(self.failure_or_closed())
306 } else if self.started {
307 let (registration, receivers) = Registration::register(key.clone(), buffer);
308 let reconnected = self.reconnected;
309 let joined = match self.channels.entry(key) {
310 Entry::Occupied(mut entry) => {
311 let channel = entry.get_mut();
312 channel.holders.push(Arc::clone(®istration));
313 if channel.confirmed {
314 Joined::Confirmed { reconnected }
315 } else {
316 Joined::Pending
317 }
318 }
319 Entry::Vacant(entry) => {
320 entry.insert(Channel {
321 confirmed: false,
322 holders: vec![Arc::clone(®istration)],
323 });
324 Joined::New
325 }
326 };
327 Ok((registration, receivers, joined))
328 } else {
329 Err(Error::NotConnected)
330 }
331 }
332
333 fn writer(&self) -> Result<Writer, Error> {
336 match (&self.conn, &self.protocol) {
337 (Some(conn), Some(protocol)) if self.welcomed => {
338 Ok((Arc::clone(conn), Arc::clone(protocol)))
339 }
340 _ => Err(Error::NotConnected),
341 }
342 }
343
344 fn failure_or_closed(&self) -> Error {
346 self.failure.clone().unwrap_or(Error::Closed)
347 }
348
349 fn registrations(&self) -> Vec<Arc<Registration>> {
350 self.channels
351 .values()
352 .flat_map(|channel| channel.holders.iter().cloned())
353 .collect()
354 }
355
356 fn holders(&self, key: &str) -> Vec<Arc<Registration>> {
357 self.channels
358 .get(key)
359 .map(|channel| channel.holders.clone())
360 .unwrap_or_default()
361 }
362
363 fn pending_keys(&self) -> Vec<String> {
365 self.channels
366 .iter()
367 .filter(|(_, channel)| !channel.confirmed)
368 .map(|(key, _)| key.clone())
369 .collect()
370 }
371
372 fn unconfirm_channels(&mut self) {
373 for channel in self.channels.values_mut() {
374 channel.confirmed = false;
375 }
376 }
377}
378
379impl Inner {
380 pub(crate) async fn send(&self, command: Command) -> Result<(), Error> {
384 let _writing = self.write_lock.lock().await;
385 self.write(command).await
386 }
387
388 async fn write(&self, command: Command) -> Result<(), Error> {
390 let (conn, protocol) = self.state().writer()?;
391 let payload = protocol.encode(&command)?;
392
393 let written = conn.write(&payload).await;
394 if let Err(error) = &written {
395 tracing::debug!(%error, "write failed, hanging up for a redial");
396 conn.close().await;
397 }
398 written
399 }
400
401 pub(crate) fn forget(&self, registration: &Arc<Registration>, reason: &Error) -> bool {
404 let last = {
405 let mut state = self.state();
406 match state.channels.get_mut(®istration.key) {
407 Some(channel) => {
408 channel
409 .holders
410 .retain(|holder| !Arc::ptr_eq(holder, registration));
411 let last = channel.holders.is_empty();
412 if last {
413 state.channels.remove(®istration.key);
414 }
415 last
416 }
417 None => false,
418 }
419 };
420
421 registration.close(reason);
422 last
423 }
424
425 pub(crate) fn unsubscribe_later(self: &Arc<Inner>, key: &str) {
429 match tokio::runtime::Handle::try_current() {
430 Ok(runtime) => {
431 let inner = Arc::clone(self);
432 let key = key.to_string();
433 runtime.spawn(async move {
434 if let Err(error) = inner.send(Command::unsubscribe(&key)).await {
435 tracing::debug!(%key, %error, "unsubscribing after a drop");
436 }
437 });
438 }
439 Err(_no_runtime) => {
440 tracing::warn!(%key, "dropped outside a tokio runtime; the server keeps the subscription until the connection ends");
441 }
442 }
443 }
444
445 async fn run(self: Arc<Inner>) {
446 let mut stopping = self.stopping.subscribe();
447
448 loop {
449 let ended = tokio::select! {
450 biased;
451 _ = stopping.wait_for(|stopping| *stopping) => None,
452 ended = self.session() => Some(ended),
453 };
454 self.disconnect().await;
455
456 if let Some(error) = ended
457 && !self.is_stopped()
458 {
459 tracing::warn!(url = %self.config.url, %error, "connection ended");
460 }
461 if self.is_stopped() {
462 break;
463 }
464
465 tokio::select! {
466 biased;
467 _ = stopping.wait_for(|stopping| *stopping) => break,
468 () = tokio::time::sleep(self.reconnect_delay()) => {}
469 }
470 }
471
472 self.announce_stop();
473 self.close_subscriptions();
474 self.done.send_replace(true);
475 }
476
477 async fn session(&self) -> Error {
481 if self.config.protocols.is_empty() {
482 return self.stop(Error::NoProtocols);
483 }
484
485 let headers = match self.dial_headers().await {
486 Ok(headers) => headers,
487 Err(error) => return self.failed(error),
488 };
489
490 let dialing = self.config.transport.dial(
491 &self.config.url,
492 DialOptions {
493 subprotocols: self.offered_subprotocols(),
494 headers,
495 },
496 );
497 let conn: Arc<dyn Conn> = match dialing.await {
498 Ok(conn) => conn.into(),
499 Err(error) => return self.failed(error),
500 };
501
502 let protocol = match self.negotiated(conn.subprotocol()) {
503 Ok(protocol) => protocol,
504 Err(error) => {
505 conn.close().await;
506 return self.stop(error);
507 }
508 };
509
510 {
511 let mut state = self.state();
512 state.conn = Some(Arc::clone(&conn));
513 state.protocol = Some(Arc::clone(&protocol));
514 }
515
516 let ended = tokio::select! {
517 ended = self.receive(&conn, &protocol) => ended,
518 never = self.guarantee_subscriptions() => match never {},
519 };
520 self.failed(ended)
521 }
522
523 fn failed(&self, error: Error) -> Error {
527 if self.is_stopped() {
528 return error;
529 }
530
531 if let Some(terminal) = &self.config.stop_on_error
532 && terminal(&error)
533 {
534 return self.stop(error);
535 }
536
537 let attempts = self.count_attempt(&error);
538 if attempts == self.config.max_attempts {
539 self.stop(Error::GaveUp {
540 attempts,
541 last: Box::new(error.clone()),
542 });
543 }
544 error
545 }
546
547 async fn dial_headers(&self) -> Result<HeaderMap, Error> {
551 let mut headers = self.config.headers.clone();
552 if let Some(provider) = &self.config.header_provider {
553 headers.extend(provider.headers().await?);
554 }
555 Ok(headers)
556 }
557
558 fn offered_subprotocols(&self) -> Vec<String> {
560 let mut offered = self.subprotocols();
561 offered.push(SUBPROTOCOL_UNSUPPORTED.to_string());
562 offered
563 }
564
565 fn subprotocols(&self) -> Vec<String> {
566 self.config
567 .protocols
568 .iter()
569 .map(|protocol| protocol.subprotocol().to_string())
570 .collect()
571 }
572
573 fn negotiated(&self, subprotocol: Option<&str>) -> Result<Arc<dyn Protocol>, Error> {
577 let negotiated = subprotocol.unwrap_or_default();
578 match self
579 .config
580 .protocols
581 .iter()
582 .find(|protocol| protocol.subprotocol() == negotiated)
583 {
584 Some(protocol) => Ok(Arc::clone(protocol)),
585 None => Err(Error::UnsupportedSubprotocol {
586 negotiated: negotiated.to_string(),
587 offered: self.subprotocols(),
588 }),
589 }
590 }
591
592 async fn receive(&self, conn: &Arc<dyn Conn>, protocol: &Arc<dyn Protocol>) -> Error {
595 loop {
596 let read = tokio::time::timeout(self.config.stale_after, conn.read()).await;
597 let payload = match read {
598 Ok(Ok(payload)) => payload,
599 Ok(Err(error)) => return error,
600 Err(_elapsed) => {
601 return Error::Stale {
602 after: self.config.stale_after,
603 };
604 }
605 };
606
607 match protocol.decode(&payload) {
608 Ok(incoming) => {
609 if let Some(ended) = self.dispatch(incoming).await {
610 return ended;
611 }
612 }
613 Err(error) => {
614 tracing::warn!(%error, "dropping undecodable frame");
615 }
616 }
617 }
618 }
619
620 async fn dispatch(&self, incoming: Incoming) -> Option<Error> {
622 match incoming.kind {
623 Kind::Welcome => {
624 self.welcome().await;
625 None
626 }
627 Kind::Ping => None,
630 Kind::Disconnect => Some(self.hang_up(incoming)),
631 Kind::Confirmation => {
632 self.confirm(&incoming.identifier);
633 None
634 }
635 Kind::Rejection => {
636 self.reject(&incoming.identifier);
637 None
638 }
639 Kind::Message => {
640 self.deliver(incoming);
641 None
642 }
643 }
644 }
645
646 async fn welcome(&self) {
649 let _writing = self.write_lock.lock().await;
650 let keys = {
651 let mut state = self.state();
652 state.attempts = 0;
653 state.welcomed = true;
654 state.reconnected = state.ever_welcomed;
655 state.ever_welcomed = true;
656 state.unconfirm_channels();
657 state.pending_keys()
658 };
659
660 self.connected.send_replace(true);
661
662 self.write_subscribes(keys).await;
663 }
664
665 async fn guarantee_subscriptions(&self) -> Infallible {
669 loop {
670 tokio::time::sleep(self.config.subscribe_retry).await;
671
672 let _writing = self.write_lock.lock().await;
673 let pending = self.state().pending_keys();
674 self.write_subscribes(pending).await;
675 }
676 }
677
678 async fn write_subscribes(&self, keys: Vec<String>) {
685 for key in keys {
686 if let Err(error) = self.write(Command::subscribe(&key)).await {
687 tracing::debug!(%key, %error, "resubscribing; the next welcome will send it again");
688 }
689 }
690 }
691
692 async fn send_subscribe(&self, key: &str) {
694 if let Err(error) = self.send(Command::subscribe(key)).await {
695 tracing::debug!(%key, %error, "subscribing; the next welcome will send it again");
696 }
697 }
698
699 fn confirm(&self, key: &str) {
702 let (holders, reconnected) = {
703 let mut state = self.state();
704 let reconnected = state.reconnected;
705 match state.channels.get_mut(key) {
706 Some(channel) if !channel.confirmed => {
707 channel.confirmed = true;
708 (channel.holders.clone(), reconnected)
709 }
710 _ => (Vec::new(), reconnected),
711 }
712 };
713
714 for holder in holders {
715 holder.confirm(reconnected);
716 }
717 }
718
719 fn reject(&self, key: &str) {
720 let holders = self
721 .state()
722 .channels
723 .remove(key)
724 .map(|channel| channel.holders)
725 .unwrap_or_default();
726
727 for holder in holders {
728 holder.reject();
729 }
730 }
731
732 fn deliver(&self, incoming: Incoming) {
733 let holders = self.state().holders(&incoming.identifier);
734
735 match incoming.message {
736 Some(message) if holders.is_empty() => {
737 tracing::debug!(identifier = %incoming.identifier, %message, "no subscription, dropping message");
738 }
739 Some(message) => {
740 for holder in holders {
741 if !holder.deliver(message.clone()) {
742 tracing::warn!(identifier = %incoming.identifier, "message buffer full, dropping message");
743 }
744 }
745 }
746 None => {
747 tracing::debug!(identifier = %incoming.identifier, "frame without a message, dropping it");
748 }
749 }
750 }
751
752 fn hang_up(&self, incoming: Incoming) -> Error {
753 let disconnect = Error::Disconnected {
754 reason: incoming.reason,
755 reconnect: incoming.reconnect,
756 };
757 if incoming.reconnect {
758 disconnect
759 } else {
760 self.stop(disconnect)
761 }
762 }
763
764 async fn disconnect(&self) {
766 let conn = {
767 let mut state = self.state();
768 state.protocol = None;
769 state.welcomed = false;
770 state.unconfirm_channels();
771 state.conn.take()
772 };
773
774 if let Some(conn) = conn {
775 conn.close().await;
776 self.announce_disconnected();
777 }
778 }
779
780 fn announce_disconnected(&self) {
783 let (registrations, will_reconnect) = {
784 let mut state = self.state();
785 let will_reconnect = !state.stopped;
786 if !will_reconnect {
787 state.stop_announced = true;
788 }
789 (state.registrations(), will_reconnect)
790 };
791
792 for registration in registrations {
793 registration.disconnect(will_reconnect);
794 }
795 }
796
797 fn count_attempt(&self, error: &Error) -> u32 {
798 let mut state = self.state();
799 state.attempts += 1;
800 state.last_error = Some(error.clone());
801 state.attempts
802 }
803
804 fn reconnect_delay(&self) -> Duration {
808 let attempts = self.state().attempts;
809 let doublings = attempts.saturating_sub(1).min(16);
810 let delay = self
811 .config
812 .initial_backoff
813 .saturating_mul(1 << doublings)
814 .min(self.config.longest_backoff);
815 let half = delay / 2;
816 half + jitter(half)
817 }
818
819 fn announce_stop(&self) {
823 if !self.state().stop_announced {
824 self.announce_disconnected();
825 }
826 }
827
828 fn close_subscriptions(&self) {
829 let (registrations, failure) = {
830 let mut state = self.state();
831 let registrations = state.registrations();
832 state.channels.clear();
833 (registrations, state.failure_or_closed())
834 };
835
836 for registration in registrations {
837 registration.close(&failure);
838 }
839 }
840
841 fn stop(&self, error: Error) -> Error {
843 {
844 let mut state = self.state();
845 state.stopped = true;
846 state.failure.get_or_insert(error.clone());
847 }
848 self.stopping.send_replace(true);
849 error
850 }
851
852 fn is_stopped(&self) -> bool {
853 self.state().stopped
854 }
855
856 fn state(&self) -> MutexGuard<'_, State> {
857 self.state.lock().unwrap_or_else(PoisonError::into_inner)
858 }
859}
860
861fn jitter(up_to: Duration) -> Duration {
865 let nanos = u64::try_from(up_to.as_nanos()).unwrap_or(u64::MAX);
866 let random = getrandom::u64().unwrap_or(0);
867 Duration::from_nanos(random % nanos.saturating_add(1))
868}
869
870#[cfg(test)]
871#[allow(
872 clippy::unwrap_used,
873 clippy::expect_used,
874 reason = "a test that can't have what it asked for has nothing left to assert"
875)]
876mod tests {
877 use super::*;
878
879 #[test]
880 fn jitter_stays_within_the_interval() {
881 for _ in 0..100 {
882 assert!(jitter(Duration::from_millis(10)) <= Duration::from_millis(10));
883 }
884 assert_eq!(Duration::ZERO, jitter(Duration::ZERO));
885 assert!(jitter(Duration::MAX) <= Duration::from_nanos(u64::MAX));
886 }
887}