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
21pub 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#[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 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}