Skip to main content

rust_etcd_utils/
watcher.rs

1use {
2    super::{Revision, retry::retry_etcd_legacy},
3    crate::retry::is_transient,
4    etcd_client::{
5        Error, EventType, WatchClient, WatchFilterType, WatchOptions, WatchResponse,
6        WatchStream as EtcdClientWatchStream,
7    },
8    futures::{Future, Stream},
9    pin_project::pin_project,
10    retry::delay::Exponential,
11    serde::de::DeserializeOwned,
12    std::{
13        collections::VecDeque,
14        marker::PhantomData,
15        pin::{Pin, pin},
16        task::{Context, Poll, ready},
17    },
18    tracing::{error, info},
19};
20
21///
22/// Custom types for watch events.
23///
24/// Unwrap the etcd watch event to a more user-friendly event.
25///
26pub enum WatchEvent<V> {
27    Put {
28        key: Vec<u8>,
29        value: V,
30        revision: Revision,
31    },
32    Delete {
33        key: Vec<u8>,
34        prev_value: Option<V>,
35        revision: Revision,
36    },
37}
38
39enum ReconnectState<S> {
40    Disconnected,
41    Connecting(Pin<Box<dyn Future<Output = Result<S, Error>> + Send>>),
42    Streaming { stream: S },
43    Terminated,
44}
45
46pub trait EtcdConnector {
47    type WatchStream: Stream<Item = Result<WatchResponse, Error>> + Unpin + Send + 'static;
48    type ConnectFut: Future<Output = Result<Self::WatchStream, Error>> + Send + 'static;
49
50    fn connect_watch(&mut self, last_revision: Option<Revision>) -> Self::ConnectFut;
51}
52
53pub struct GrpcEtcdConenctor {
54    watch_client: WatchClient,
55    key: Vec<u8>,
56    watch_options_prototype: WatchOptions,
57}
58
59impl GrpcEtcdConenctor {
60    pub fn new(
61        watch_client: WatchClient,
62        key: Vec<u8>,
63        watch_options_prototype: WatchOptions,
64    ) -> Self {
65        Self {
66            watch_client,
67            key,
68            watch_options_prototype,
69        }
70    }
71}
72
73impl EtcdConnector for GrpcEtcdConenctor {
74    type WatchStream = EtcdClientWatchStream;
75    type ConnectFut = Pin<Box<dyn Future<Output = Result<Self::WatchStream, Error>> + Send>>;
76
77    fn connect_watch(&mut self, last_revision: Option<Revision>) -> Self::ConnectFut {
78        let wc = self.watch_client.clone();
79        let key = self.key.clone();
80        let mut wopts = self.watch_options_prototype.clone();
81        if let Some(rev) = last_revision {
82            wopts = wopts.with_start_revision(rev);
83        }
84
85        Box::pin(async move {
86            let retry_strategy = Exponential::from_millis_with_factor(10, 10.0).take(3);
87            retry_etcd_legacy(retry_strategy, move || {
88                let mut wc = wc.clone();
89                let key = key.clone();
90                let wopts = wopts.clone();
91                async move { wc.watch(key.clone(), Some(wopts)).await }
92            })
93            .await
94        })
95    }
96}
97
98pub struct AutoReconnectWatchStream<C>
99where
100    C: EtcdConnector,
101{
102    connector: C,
103    state: ReconnectState<C::WatchStream>,
104    last_revision: Option<Revision>,
105}
106
107impl<C> AutoReconnectWatchStream<C>
108where
109    C: EtcdConnector,
110{
111    pub fn new(connector: C) -> Self {
112        Self {
113            connector,
114            state: ReconnectState::Disconnected,
115            last_revision: None,
116        }
117    }
118}
119
120impl<C> Stream for AutoReconnectWatchStream<C>
121where
122    C: EtcdConnector + Unpin,
123{
124    type Item = WatchResponse;
125
126    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
127        let this = self.as_mut().get_mut();
128
129        loop {
130            match &mut this.state {
131                ReconnectState::Disconnected => {
132                    let fut = this.connector.connect_watch(this.last_revision);
133                    this.state = ReconnectState::Connecting(Box::pin(fut));
134                }
135                ReconnectState::Connecting(fut) => {
136                    let stream = match ready!(fut.as_mut().poll(cx)) {
137                        Ok(v) => v,
138                        Err(e) => {
139                            error!("watch reconnect failed: {e}");
140                            this.state = ReconnectState::Terminated;
141                            continue;
142                        }
143                    };
144                    this.state = ReconnectState::Streaming { stream };
145                }
146                ReconnectState::Streaming { stream } => match Pin::new(stream).poll_next(cx) {
147                    Poll::Ready(Some(Ok(watch_resp))) => {
148                        if watch_resp.canceled() {
149                            error!("watch cancelled: {watch_resp:?}");
150                            this.state = ReconnectState::Terminated;
151                            continue;
152                        }
153
154                        if let Some(revision) = watch_resp
155                            .events()
156                            .iter()
157                            .filter_map(|ev| ev.kv())
158                            .max_by_key(|kv| kv.mod_revision())
159                            .map(|kv| kv.mod_revision())
160                        {
161                            this.last_revision.replace(revision);
162                        }
163
164                        return Poll::Ready(Some(watch_resp));
165                    }
166                    Poll::Ready(Some(Err(e))) => {
167                        if is_transient(&e) {
168                            this.state = ReconnectState::Disconnected;
169                            continue;
170                        }
171                        error!("watch stream failed with non-transient error: {e}");
172                        this.state = ReconnectState::Terminated;
173                    }
174                    Poll::Ready(None) => {
175                        this.state = ReconnectState::Disconnected;
176                    }
177                    Poll::Pending => return Poll::Pending,
178                },
179                ReconnectState::Terminated => return Poll::Ready(None),
180            }
181        }
182    }
183}
184
185pub trait WatchStreamValueDecoder {
186    type Item;
187    type Error: std::error::Error + Send + 'static;
188
189    fn decode_watch_response(
190        &mut self,
191        key: &[u8],
192        value: &[u8],
193    ) -> Result<Self::Item, Self::Error>;
194}
195
196pub struct JsonDecoder<V> {
197    _phantom: PhantomData<V>,
198}
199
200impl<V> Default for JsonDecoder<V> {
201    fn default() -> Self {
202        Self {
203            _phantom: PhantomData,
204        }
205    }
206}
207
208impl<V> WatchStreamValueDecoder for JsonDecoder<V>
209where
210    V: DeserializeOwned,
211{
212    type Item = V;
213    type Error = serde_json::Error;
214
215    fn decode_watch_response(
216        &mut self,
217        _key: &[u8],
218        value: &[u8],
219    ) -> Result<Self::Item, Self::Error> {
220        serde_json::from_slice(value)
221    }
222}
223
224#[pin_project]
225pub struct ValueWatchStream<Source, Value, Decoder> {
226    #[pin]
227    inner: Source,
228    pending: VecDeque<WatchEvent<Value>>,
229    decoder: Decoder,
230}
231
232impl<S, V, D> ValueWatchStream<S, V, D> {
233    pub fn new(inner: S, decoder: D) -> Self {
234        Self {
235            inner,
236            pending: VecDeque::new(),
237            decoder,
238        }
239    }
240}
241
242impl<S, V, D> Stream for ValueWatchStream<S, V, D>
243where
244    D: WatchStreamValueDecoder<Item = V>,
245    S: Stream<Item = WatchResponse> + Unpin,
246    V: DeserializeOwned,
247{
248    type Item = WatchEvent<V>;
249
250    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
251        let mut this = self.project();
252        loop {
253            if let Some(next) = this.pending.pop_front() {
254                return Poll::Ready(Some(next));
255            }
256            let watch_resp = match ready!(this.inner.as_mut().poll_next(cx)) {
257                Some(v) => v,
258                None => return Poll::Ready(None),
259            };
260
261            for event in watch_resp.events() {
262                let parsed = match event.event_type() {
263                    EventType::Put => {
264                        let kv = event.kv().expect("put event with no kv");
265                        let key = Vec::from(kv.key());
266                        let value = this
267                            .decoder
268                            .decode_watch_response(&key, kv.value())
269                            .expect("failed to deserialize controller state");
270                        WatchEvent::Put {
271                            key,
272                            value,
273                            revision: kv.mod_revision(),
274                        }
275                    }
276                    EventType::Delete => {
277                        let kv = event.kv().expect("delete event with no kv");
278                        let prev_value = event
279                            .prev_kv()
280                            .map(|prev_kv| {
281                                this.decoder
282                                    .decode_watch_response(prev_kv.key(), prev_kv.value())
283                            })
284                            .transpose()
285                            .expect("failed to deserialize prev controller state");
286                        let key = Vec::from(kv.key());
287                        WatchEvent::Delete {
288                            key,
289                            prev_value,
290                            revision: kv.mod_revision(),
291                        }
292                    }
293                };
294
295                this.pending.push_back(parsed);
296            }
297        }
298    }
299}
300
301#[pin_project]
302pub struct PutWatchStream<Source, Value, Decoder> {
303    #[pin]
304    inner: Source,
305    _phantom: PhantomData<Value>,
306    decoder: Decoder,
307}
308
309impl<S, T, D> PutWatchStream<S, T, D> {
310    pub fn new(inner: S, decoder: D) -> Self {
311        Self {
312            inner,
313            _phantom: PhantomData,
314            decoder,
315        }
316    }
317}
318
319impl<S, T, D> Stream for PutWatchStream<S, T, D>
320where
321    S: Stream<Item = WatchResponse> + Unpin,
322    D: WatchStreamValueDecoder<Item = T>,
323{
324    type Item = (Revision, T);
325
326    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
327        let mut this = self.project();
328        loop {
329            let watch_resp = match ready!(Pin::new(&mut this.inner).poll_next(cx)) {
330                Some(v) => v,
331                None => return Poll::Ready(None),
332            };
333
334            let max_kv = watch_resp
335                .events()
336                .iter()
337                .filter_map(|ev| ev.kv())
338                .max_by_key(|kv| kv.mod_revision());
339
340            if let Some(kv) = max_kv {
341                let revision = kv.mod_revision();
342                let state = this
343                    .decoder
344                    .decode_watch_response(kv.key(), kv.value())
345                    .expect("failed to deserialize kv value");
346                return Poll::Ready(Some((revision, state)));
347            }
348        }
349    }
350}
351
352pub struct LockKeyChangeStream<S> {
353    inner: S,
354    key: Vec<u8>,
355    key_mod_revision: Revision,
356    done: bool,
357}
358
359impl<S> LockKeyChangeStream<S> {
360    pub fn new(inner: S, key: Vec<u8>, key_mod_revision: Revision) -> Self {
361        Self {
362            inner,
363            key,
364            key_mod_revision,
365            done: false,
366        }
367    }
368}
369
370impl<S> Stream for LockKeyChangeStream<S>
371where
372    S: Stream<Item = WatchResponse> + Unpin,
373{
374    type Item = Revision;
375
376    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
377        let this = self.as_mut().get_mut();
378
379        if this.done {
380            return Poll::Ready(None);
381        }
382
383        loop {
384            let watch_resp = match ready!(Pin::new(&mut this.inner).poll_next(cx)) {
385                Some(v) => v,
386                None => {
387                    this.done = true;
388                    return Poll::Ready(None);
389                }
390            };
391
392            for event in watch_resp.events() {
393                match event.event_type() {
394                    EventType::Put => {
395                        let kv = event.kv().expect("put event with no kv");
396                        if kv.key() == this.key {
397                            continue;
398                        }
399                        let revision = kv.mod_revision();
400                        if revision <= this.key_mod_revision {
401                            continue;
402                        }
403                        info!(
404                            "watcher detected put event on key {:?} with revision {} > {}",
405                            this.key, revision, this.key_mod_revision
406                        );
407                        this.done = true;
408                        return Poll::Ready(Some(revision));
409                    }
410                    EventType::Delete => {
411                        let kv = event.kv().expect("delete event with no kv");
412                        let revision = kv.mod_revision();
413                        if revision < this.key_mod_revision {
414                            continue;
415                        }
416                        if kv.key() == this.key {
417                            let key_label = String::from_utf8_lossy(&this.key);
418                            info!(
419                                "watcher detected delete event on key {:?} with revision {} >= {}",
420                                key_label, revision, this.key_mod_revision
421                            );
422                            this.done = true;
423                            return Poll::Ready(Some(revision));
424                        }
425                    }
426                }
427            }
428        }
429    }
430}
431
432pub type EtcdReconnectWatchStream = AutoReconnectWatchStream<GrpcEtcdConenctor>;
433pub type EtcdJsonWatchStream<V> = ValueWatchStream<EtcdReconnectWatchStream, V, JsonDecoder<V>>;
434pub type EtcdJsonPutWatchStream<T> = PutWatchStream<EtcdReconnectWatchStream, T, JsonDecoder<T>>;
435pub type EtcdLockKeyChangeStream = LockKeyChangeStream<EtcdReconnectWatchStream>;
436
437///
438/// Extension trait for [`WatchClient`].
439///
440/// This trait provides utility methods for working with [`WatchClient`].
441///
442/// This extension trait provides rust channel of watch stream and more reliability in case of transient errors.
443///
444/// On transient errors, the watch stream will be retried and resume where you left off.
445///
446#[async_trait::async_trait]
447pub trait WatchClientExt {
448    fn get_watch_client(&self) -> WatchClient;
449
450    fn json_watch_stream<V>(
451        &self,
452        key: impl Into<Vec<u8>>,
453        watch_options: Option<WatchOptions>,
454    ) -> EtcdJsonWatchStream<V>
455    where
456        V: DeserializeOwned + Send + 'static,
457    {
458        self.value_watch_stream(key, watch_options, JsonDecoder::default())
459    }
460
461    fn value_watch_stream<D>(
462        &self,
463        key: impl Into<Vec<u8>>,
464        watch_options: Option<WatchOptions>,
465        decoder: D,
466    ) -> ValueWatchStream<EtcdReconnectWatchStream, D::Item, D>
467    where
468        D: WatchStreamValueDecoder,
469    {
470        let wc = self.get_watch_client();
471        let key: Vec<u8> = key.into();
472        let wopts_prototype = watch_options.unwrap_or_default().with_prev_key();
473
474        let connector = GrpcEtcdConenctor::new(wc, key, wopts_prototype);
475        let reconnecting_stream = AutoReconnectWatchStream::new(connector);
476
477        ValueWatchStream::new(reconnecting_stream, decoder)
478    }
479
480    fn put_watch_stream<D>(
481        &self,
482        key: impl Into<Vec<u8>>,
483        watch_options: Option<WatchOptions>,
484        decoder: D,
485    ) -> PutWatchStream<EtcdReconnectWatchStream, D::Item, D>
486    where
487        D: WatchStreamValueDecoder,
488    {
489        let wc = self.get_watch_client();
490        let key: Vec<u8> = key.into();
491        let wopts_prototype = watch_options
492            .unwrap_or_default()
493            .with_filters(vec![WatchFilterType::NoDelete]);
494
495        let connector = GrpcEtcdConenctor::new(wc, key, wopts_prototype);
496        let reconnecting_stream = AutoReconnectWatchStream::new(connector);
497
498        PutWatchStream::new(reconnecting_stream, decoder)
499    }
500
501    fn json_put_watch_stream<T>(
502        &self,
503        key: impl Into<Vec<u8>>,
504        watch_options: Option<WatchOptions>,
505    ) -> EtcdJsonPutWatchStream<T>
506    where
507        T: DeserializeOwned + Send + 'static,
508    {
509        self.put_watch_stream(key, watch_options, JsonDecoder::default())
510    }
511
512    ///
513    /// Creates a channel that watches for changes to a key in etcd.
514    ///
515    /// The channel will send a [`WatchEvent`] for each change to the key.
516    /// The channel will be retried on transient errors.
517    ///
518    /// The channel will be closed if the watch is cancelled or if the stream is closed.
519    ///
520    /// The watch expect value to be JSON encoded.
521    fn json_watch_channel<V>(
522        &self,
523        key: impl Into<Vec<u8>>,
524        watch_options: Option<WatchOptions>,
525    ) -> EtcdJsonWatchStream<V>
526    where
527        V: DeserializeOwned + Send + 'static,
528    {
529        self.json_watch_stream::<V>(key, watch_options)
530    }
531
532    fn json_put_watch_channel<T>(
533        &self,
534        key: impl Into<Vec<u8>>,
535        watch_options: Option<WatchOptions>,
536    ) -> EtcdJsonPutWatchStream<T>
537    where
538        T: DeserializeOwned + Send + 'static,
539    {
540        self.json_put_watch_stream::<T>(key, watch_options)
541    }
542
543    fn watch_lock_key_change_stream(
544        &self,
545        key: impl Into<Vec<u8>>,
546        key_mod_revision: Revision,
547    ) -> EtcdLockKeyChangeStream {
548        let wc = self.get_watch_client();
549        let key: Vec<u8> = key.into();
550        let wopts_prototype = WatchOptions::new().with_start_revision(key_mod_revision);
551
552        let connector = GrpcEtcdConenctor::new(wc, key.clone(), wopts_prototype);
553        let reconnecting_stream = AutoReconnectWatchStream::new(connector);
554
555        LockKeyChangeStream::new(reconnecting_stream, key, key_mod_revision)
556    }
557}
558
559impl WatchClientExt for WatchClient {
560    fn get_watch_client(&self) -> WatchClient {
561        self.clone()
562    }
563}
564
565#[cfg(test)]
566mod tests {
567    use super::*;
568    use futures::{StreamExt, executor::block_on, future, stream};
569    use std::{
570        collections::VecDeque,
571        sync::{Arc, Mutex},
572    };
573
574    fn make_watch_response(canceled: bool, revision: Option<Revision>) -> WatchResponse {
575        let events = revision
576            .map(|rev| {
577                vec![etcd_client::proto::PbEvent {
578                    r#type: 0,
579                    kv: Some(etcd_client::proto::PbKeyValue {
580                        mod_revision: rev,
581                        ..Default::default()
582                    }),
583                    ..Default::default()
584                }]
585            })
586            .unwrap_or_default();
587
588        WatchResponse(etcd_client::proto::PbWatchResponse {
589            canceled,
590            events,
591            ..Default::default()
592        })
593    }
594
595    fn transient_error() -> Error {
596        Error::GRpcStatus(tonic::Status::new(tonic::Code::Unavailable, "transient"))
597    }
598
599    fn non_transient_error() -> Error {
600        Error::GRpcStatus(tonic::Status::new(
601            tonic::Code::InvalidArgument,
602            "non-transient",
603        ))
604    }
605
606    struct TestConnector {
607        connect_called: Arc<Mutex<bool>>,
608    }
609
610    impl EtcdConnector for TestConnector {
611        type WatchStream = stream::Empty<Result<WatchResponse, Error>>;
612        type ConnectFut = future::Ready<Result<Self::WatchStream, Error>>;
613
614        fn connect_watch(&mut self, _last_revision: Option<Revision>) -> Self::ConnectFut {
615            *self.connect_called.lock().expect("mutex poisoned") = true;
616            future::ready(Err(Error::WatchError("connect failed".to_string())))
617        }
618    }
619
620    struct MockConnector {
621        calls: Arc<Mutex<Vec<Option<Revision>>>>,
622        plans: VecDeque<Result<Vec<Result<WatchResponse, Error>>, Error>>,
623    }
624
625    impl MockConnector {
626        fn new(plans: Vec<Result<Vec<Result<WatchResponse, Error>>, Error>>) -> Self {
627            Self {
628                calls: Arc::new(Mutex::new(Vec::new())),
629                plans: plans.into(),
630            }
631        }
632    }
633
634    impl EtcdConnector for MockConnector {
635        type WatchStream = stream::Iter<std::vec::IntoIter<Result<WatchResponse, Error>>>;
636        type ConnectFut = future::Ready<Result<Self::WatchStream, Error>>;
637
638        fn connect_watch(&mut self, last_revision: Option<Revision>) -> Self::ConnectFut {
639            self.calls
640                .lock()
641                .expect("mutex poisoned")
642                .push(last_revision);
643            let planned = self.plans.pop_front().expect("missing test plan");
644            match planned {
645                Ok(items) => future::ready(Ok(stream::iter(items))),
646                Err(e) => future::ready(Err(e)),
647            }
648        }
649    }
650
651    #[test]
652    fn reconnecting_stream_terminates_when_connect_fails() {
653        let connect_called = Arc::new(Mutex::new(false));
654        let mut stream = AutoReconnectWatchStream::new(TestConnector {
655            connect_called: Arc::clone(&connect_called),
656        });
657
658        let next = block_on(stream.next());
659        assert!(next.is_none());
660        assert!(*connect_called.lock().expect("mutex poisoned"));
661    }
662
663    #[test]
664    fn json_watch_stream_returns_none_when_inner_is_empty() {
665        let mut stream =
666            ValueWatchStream::<_, serde_json::Value, JsonDecoder<serde_json::Value>>::new(
667                stream::empty(),
668                JsonDecoder::default(),
669            );
670        let next = block_on(stream.next());
671        assert!(next.is_none());
672    }
673
674    #[test]
675    fn json_put_watch_stream_returns_none_when_inner_is_empty() {
676        let mut stream =
677            PutWatchStream::<_, serde_json::Value, JsonDecoder<serde_json::Value>>::new(
678                stream::empty(),
679                JsonDecoder::default(),
680            );
681        let next = block_on(stream.next());
682        assert!(next.is_none());
683    }
684
685    #[test]
686    fn lock_key_change_stream_returns_none_when_inner_is_empty() {
687        let mut stream = LockKeyChangeStream::new(stream::empty(), b"/lock/key".to_vec(), 42);
688        let next = block_on(stream.next());
689        assert!(next.is_none());
690    }
691
692    #[test]
693    fn reconnecting_stream_terminates_on_non_transient_stream_error() {
694        let connector = MockConnector::new(vec![Ok(vec![Err(non_transient_error())])]);
695        let calls = Arc::clone(&connector.calls);
696        let mut stream = AutoReconnectWatchStream::new(connector);
697
698        let next = block_on(stream.next());
699        assert!(next.is_none());
700        assert_eq!(*calls.lock().expect("mutex poisoned"), vec![None]);
701    }
702
703    #[test]
704    fn reconnecting_stream_reconnects_on_transient_stream_error() {
705        let connector = MockConnector::new(vec![
706            Ok(vec![Err(transient_error())]),
707            Ok(vec![Ok(make_watch_response(false, Some(3)))]),
708        ]);
709        let calls = Arc::clone(&connector.calls);
710        let mut stream = AutoReconnectWatchStream::new(connector);
711
712        let next = block_on(stream.next());
713        assert!(next.is_some());
714        assert_eq!(*calls.lock().expect("mutex poisoned"), vec![None, None]);
715    }
716
717    #[test]
718    fn reconnecting_stream_reconnects_when_inner_stream_ends() {
719        let connector = MockConnector::new(vec![
720            Ok(vec![]),
721            Ok(vec![Ok(make_watch_response(false, Some(7)))]),
722        ]);
723        let calls = Arc::clone(&connector.calls);
724        let mut stream = AutoReconnectWatchStream::new(connector);
725
726        let next = block_on(stream.next());
727        assert!(next.is_some());
728        assert_eq!(*calls.lock().expect("mutex poisoned"), vec![None, None]);
729    }
730
731    #[test]
732    fn reconnecting_stream_stops_on_canceled_watch_response() {
733        let connector = MockConnector::new(vec![Ok(vec![Ok(make_watch_response(true, None))])]);
734        let mut stream = AutoReconnectWatchStream::new(connector);
735
736        let next = block_on(stream.next());
737        assert!(next.is_none());
738    }
739
740    #[test]
741    fn reconnecting_stream_tracks_last_revision_for_reconnect() {
742        let connector = MockConnector::new(vec![
743            Ok(vec![Ok(make_watch_response(false, Some(11)))]),
744            Ok(vec![]),
745            Err(non_transient_error()),
746        ]);
747        let calls = Arc::clone(&connector.calls);
748        let mut stream = AutoReconnectWatchStream::new(connector);
749
750        let first = block_on(stream.next());
751        assert!(first.is_some());
752        let second = block_on(stream.next());
753        assert!(second.is_none());
754
755        let calls = calls.lock().expect("mutex poisoned");
756        assert_eq!(calls[0], None);
757        assert!(calls.iter().skip(1).all(|v| *v == Some(11)));
758    }
759}