etcd_client/rpc/

watch.rs

//! Etcd Watch RPC.

use crate::caller::{ClientCaller, ClientCallerBuilder};
pub use crate::rpc::pb::mvccpb::event::EventType;

use crate::error::{Error, Result};
use crate::intercept::InterceptedChannel;
use crate::rpc::pb::etcdserverpb::watch_client::WatchClient as PbWatchClient;
use crate::rpc::pb::etcdserverpb::watch_request::RequestUnion as WatchRequestUnion;
use crate::rpc::pb::etcdserverpb::{
    WatchCancelRequest, WatchCreateRequest, WatchProgressRequest, WatchRequest,
    WatchResponse as PbWatchResponse,
};
use crate::rpc::pb::mvccpb::Event as PbEvent;
use crate::rpc::{KeyRange, KeyValue, ResponseHeader};
use std::pin::Pin;
use std::task::{Context, Poll};
use tokio::sync::mpsc::{channel, Sender};
use tokio_stream::{wrappers::ReceiverStream, Stream};
use tonic::Streaming;

type Client = PbWatchClient<InterceptedChannel>;

/// Client for watch operations.
#[repr(transparent)]
#[derive(Clone)]
pub struct WatchClient {
    inner: ClientCaller<Client>,
}

impl WatchClient {
    /// Creates a watch client.
    #[inline]
    pub(crate) fn new(builder: ClientCallerBuilder) -> Self {
        Self {
            inner: builder.build(Client::new),
        }
    }

    /// Limits the maximum size of a decoded message.
    ///
    /// Default: `4MB`
    pub fn max_decoding_message_size(mut self, limit: usize) -> Self {
        self.inner = self
            .inner
            .with(|client| client.max_decoding_message_size(limit));
        self
    }

    /// Watches for events happening or that have happened. Both input and output
    /// are streams; the input stream is for creating and canceling watchers and the output
    /// stream receives responses and events.
    ///
    /// One watch stream can watch on multiple key ranges, streaming events for several watches
    /// are grouped by watch ID. The entire event history can be watched starting from the
    /// last compaction revision.
    pub async fn watch(
        &mut self,
        key: impl Into<Vec<u8>>,
        options: Option<WatchOptions>,
    ) -> Result<WatchStream> {
        async fn watch_impl(client: &mut Client, options: WatchOptions) -> Result<WatchStream> {
            let (request_sender, request_receiver) = channel::<WatchRequest>(100);
            request_sender
                .send(options.into())
                .await
                .map_err(|e| Error::WatchError(e.to_string()))?;
            let request_stream = ReceiverStream::new(request_receiver);
            let stream = client.watch(request_stream).await?.into_inner();
            Ok(WatchStream::new(request_sender, stream))
        }
        self.inner
            .do_call(options.unwrap_or_default().with_key(key), watch_impl)
            .await
    }
}

/// Options for `Watch` operation.
#[derive(Debug, Default, Clone)]
pub struct WatchOptions {
    req: WatchCreateRequest,
    key_range: KeyRange,
}

impl WatchOptions {
    /// Sets key.
    #[inline]
    pub fn with_key(mut self, key: impl Into<Vec<u8>>) -> Self {
        self.key_range.with_key(key);
        self
    }

    /// Creates a new `WatchOptions`.
    #[inline]
    pub const fn new() -> Self {
        Self {
            req: WatchCreateRequest {
                key: Vec::new(),
                range_end: Vec::new(),
                start_revision: 0,
                progress_notify: false,
                filters: Vec::new(),
                prev_kv: false,
                watch_id: 0,
                fragment: false,
            },
            key_range: KeyRange::new(),
        }
    }

    /// Sets the end of the range `[key, end)` to watch.
    ///
    /// If `end` is not given, only the key argument is watched.
    ///
    /// If `end` is equal to `\0`, all keys greater than or equal to the key argument are watched.
    #[inline]
    pub fn with_range(mut self, end: impl Into<Vec<u8>>) -> Self {
        self.key_range.with_range(end);
        self
    }

    /// Watches all keys >= key.
    #[inline]
    pub fn with_from_key(mut self) -> Self {
        self.key_range.with_from_key();
        self
    }

    /// Watches all keys prefixed with key.
    #[inline]
    pub fn with_prefix(mut self) -> Self {
        self.key_range.with_prefix();
        self
    }

    /// Watches all keys.
    #[inline]
    pub fn with_all_keys(mut self) -> Self {
        self.key_range.with_all_keys();
        self
    }

    /// Sets the revision to watch from (inclusive). No `start_revision` is "now".
    #[inline]
    pub const fn with_start_revision(mut self, revision: i64) -> Self {
        self.req.start_revision = revision;
        self
    }

    /// `progress_notify` is set so that the etcd server will periodically send a `WatchResponse` with
    /// no events to the new watcher if there are no recent events. It is useful when clients
    /// wish to recover a disconnected watcher starting from a recent known revision.
    /// The etcd server may decide how often it will send notifications based on current load.
    #[inline]
    pub const fn with_progress_notify(mut self) -> Self {
        self.req.progress_notify = true;
        self
    }

    /// Filter the events at server side before it sends back to the watcher.
    #[inline]
    pub fn with_filters(mut self, filters: impl Into<Vec<WatchFilterType>>) -> Self {
        self.req.filters = filters.into().into_iter().map(|f| f as i32).collect();
        self
    }

    /// If `prev_kv` is set, created watcher gets the previous KV before the event happens.
    /// If the previous KV is already compacted, nothing will be returned.
    #[inline]
    pub const fn with_prev_key(mut self) -> Self {
        self.req.prev_kv = true;
        self
    }

    /// If `watch_id` is provided and non-zero, it will be assigned to this watcher.
    /// Since creating a watcher in etcd is not a synchronous operation,
    /// this can be used ensure that ordering is correct when creating multiple
    /// watchers on the same stream. Creating a watcher with an ID already in
    /// use on the stream will cause an error to be returned.
    #[inline]
    pub const fn with_watch_id(mut self, watch_id: i64) -> Self {
        self.req.watch_id = watch_id;
        self
    }

    /// Enables splitting large revisions into multiple watch responses.
    #[inline]
    pub const fn with_fragment(mut self) -> Self {
        self.req.fragment = true;
        self
    }
}

impl From<WatchOptions> for WatchCreateRequest {
    #[inline]
    fn from(mut options: WatchOptions) -> Self {
        let (key, range_end) = options.key_range.build();
        options.req.key = key;
        options.req.range_end = range_end;
        options.req
    }
}

impl From<WatchOptions> for WatchRequest {
    #[inline]
    fn from(options: WatchOptions) -> Self {
        Self {
            request_union: Some(WatchRequestUnion::CreateRequest(options.into())),
        }
    }
}

impl From<WatchCancelRequest> for WatchRequest {
    #[inline]
    fn from(req: WatchCancelRequest) -> Self {
        Self {
            request_union: Some(WatchRequestUnion::CancelRequest(req)),
        }
    }
}

impl From<WatchProgressRequest> for WatchRequest {
    #[inline]
    fn from(req: WatchProgressRequest) -> Self {
        Self {
            request_union: Some(WatchRequestUnion::ProgressRequest(req)),
        }
    }
}

/// Watch filter type.
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord, ::prost::Enumeration)]
#[repr(i32)]
pub enum WatchFilterType {
    /// Filter out put event.
    NoPut = 0,
    /// Filter out delete event.
    NoDelete = 1,
}

/// Response for `Watch` operation.
#[cfg_attr(feature = "pub-response-field", visible::StructFields(pub))]
#[derive(Debug, Clone)]
#[repr(transparent)]
pub struct WatchResponse(PbWatchResponse);

impl WatchResponse {
    /// Creates a new `WatchResponse`.
    #[inline]
    const fn new(resp: PbWatchResponse) -> Self {
        Self(resp)
    }

    /// Watch response header.
    #[inline]
    pub fn header(&self) -> Option<&ResponseHeader> {
        self.0.header.as_ref().map(From::from)
    }

    /// Takes the header out of the response, leaving a [`None`] in its place.
    #[inline]
    pub fn take_header(&mut self) -> Option<ResponseHeader> {
        self.0.header.take().map(ResponseHeader::new)
    }

    /// The ID of the watcher that corresponds to the response.
    #[inline]
    pub const fn watch_id(&self) -> i64 {
        self.0.watch_id
    }

    /// created is set to true if the response is for a create watch request.
    /// The client should record the watch_id and expect to receive events for
    /// the created watcher from the same stream.
    /// All events sent to the created watcher will attach with the same watch_id.
    #[inline]
    pub const fn created(&self) -> bool {
        self.0.created
    }

    /// `canceled` is set to true if the response is for a cancel watch request
    /// or if the start_revision has already been compacted.
    /// No further events will be sent to the canceled watcher.
    #[inline]
    pub const fn canceled(&self) -> bool {
        self.0.canceled
    }

    /// `compact_revision` is set to the minimum index if a watcher tries to watch
    /// at a compacted index.
    ///
    /// This happens when creating a watcher at a compacted revision or the watcher cannot
    /// catch up with the progress of the key-value store.
    ///
    /// The client should treat the watcher as canceled and should not try to create any
    /// watcher with the same start_revision again.
    #[inline]
    pub const fn compact_revision(&self) -> i64 {
        self.0.compact_revision
    }

    /// Indicates the reason for canceling the watcher.
    #[inline]
    pub fn cancel_reason(&self) -> &str {
        &self.0.cancel_reason
    }

    /// Events happened on the watched keys.
    #[inline]
    pub fn events(&self) -> &[Event] {
        unsafe { &*(self.0.events.as_slice() as *const _ as *const [Event]) }
    }
}

/// Watching event.
#[cfg_attr(feature = "pub-response-field", visible::StructFields(pub))]
#[derive(Debug, Clone)]
#[repr(transparent)]
pub struct Event(PbEvent);

impl Event {
    /// The kind of event. If type is a `Put`, it indicates
    /// new data has been stored to the key. If type is a `Delete`,
    /// it indicates the key was deleted.
    #[inline]
    pub fn event_type(&self) -> EventType {
        match self.0.r#type {
            0 => EventType::Put,
            1 => EventType::Delete,
            i => panic!("unknown event {i}"),
        }
    }

    /// The KeyValue for the event.
    /// A `Put` event contains current kv pair.
    /// A `Put` event with `kv.version()==1` indicates the creation of a key.
    /// A `Delete` event contains the deleted key with
    /// its modification revision set to the revision of deletion.
    #[inline]
    pub fn kv(&self) -> Option<&KeyValue> {
        self.0.kv.as_ref().map(From::from)
    }

    /// The key-value pair before the event happens.
    #[inline]
    pub fn prev_kv(&self) -> Option<&KeyValue> {
        self.0.prev_kv.as_ref().map(From::from)
    }
}

/// The watching handle.
#[cfg_attr(feature = "pub-response-field", visible::StructFields(pub))]
#[derive(Debug)]
pub struct WatchStream {
    request_sender: WatchRequestSender,
    response_stream: WatchResponseStream,
}

/// The sender for sending watch requests in the existing watch stream.
///
/// The watch request can be sending using the [`WatchStream`] or the [`WatchRequestSender`].
///
/// The [`WatchRequestSender`] can be obtained by splitting the [`WatchStream`] using the
/// [`WatchStream::split`] method.
#[cfg_attr(feature = "pub-response-field", visible::StructFields(pub))]
#[derive(Debug)]
pub struct WatchRequestSender(Sender<WatchRequest>);

/// The response stream for receiving watch responses in the existing watch stream.
///
/// The watch response can be receiving using the [`WatchStream`] or the [`WatchResponseStream`].
///
/// The [`WatchResponseStream`] can be obtained by splitting the [`WatchStream`] using the
/// [`WatchStream::split`] method.
#[cfg_attr(feature = "pub-response-field", visible::StructFields(pub))]
#[derive(Debug)]
pub struct WatchResponseStream(Streaming<PbWatchResponse>);

impl WatchResponseStream {
    /// Receive [`WatchResponse`] from this watch response stream.
    ///
    /// See also [`WatchStream::message`] for receiving watch response from the [`WatchStream`].
    #[inline]
    pub async fn message(&mut self) -> Result<Option<WatchResponse>> {
        self.0
            .message()
            .await
            .map(|resp| resp.map(WatchResponse::new))
            .map_err(From::from)
    }
}

impl WatchStream {
    /// Creates a new `WatchStream`.
    #[inline]
    const fn new(
        request_sender: Sender<WatchRequest>,
        response_stream: Streaming<PbWatchResponse>,
    ) -> Self {
        Self {
            request_sender: WatchRequestSender(request_sender),
            response_stream: WatchResponseStream(response_stream),
        }
    }

    /// Send watch request in the existing watch stream.
    #[inline]
    pub async fn watch(
        &mut self,
        key: impl Into<Vec<u8>>,
        options: Option<WatchOptions>,
    ) -> Result<()> {
        self.request_sender.watch(key, options).await
    }

    /// Cancels watch by specified `watch_id`.
    #[inline]
    pub async fn cancel(&mut self, watch_id: i64) -> Result<()> {
        self.request_sender.cancel(watch_id).await
    }

    /// Requests a watch stream progress status be sent in the watch response stream as soon as
    /// possible.
    #[inline]
    pub async fn request_progress(&mut self) -> Result<()> {
        self.request_sender.request_progress().await
    }

    /// Receive [`WatchResponse`] from this watch stream.
    #[inline]
    pub async fn message(&mut self) -> Result<Option<WatchResponse>> {
        self.response_stream.message().await
    }

    /// Splits the watch stream into a request sender and a response receiver (stream).
    pub fn split(self) -> (WatchRequestSender, WatchResponseStream) {
        (self.request_sender, self.response_stream)
    }
}

impl WatchRequestSender {
    /// Send watch request in the existing watch stream.
    #[inline]
    async fn send(&mut self, req: WatchRequest) -> Result<()> {
        self.0
            .send(req)
            .await
            .map_err(|e| Error::WatchError(e.to_string()))
    }

    /// Send watch request in the existing watch stream.
    ///
    /// See also [`WatchStream::watch`] for sending watch request using [`WatchStream`].
    #[inline]
    pub async fn watch(
        &mut self,
        key: impl Into<Vec<u8>>,
        options: Option<WatchOptions>,
    ) -> Result<()> {
        self.send(options.unwrap_or_default().with_key(key).into())
            .await
    }

    /// Cancels watch by specified `watch_id`.
    ///
    ///
    /// See also [`WatchStream::cancel`] for canceling watch using [`WatchStream`].
    #[inline]
    pub async fn cancel(&mut self, watch_id: i64) -> Result<()> {
        let req = WatchCancelRequest { watch_id };
        self.send(req.into()).await
    }

    /// Requests a watch stream progress status be sent in the watch response stream as soon as
    /// possible.
    ///
    /// See also [`WatchStream::request_progress`] for requesting watch stream progress status
    /// using [`WatchStream`].
    #[inline]
    pub async fn request_progress(&mut self) -> Result<()> {
        let req = WatchProgressRequest {};
        self.send(req.into()).await
    }
}

impl Stream for WatchResponseStream {
    type Item = Result<WatchResponse>;

    #[inline]
    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        Pin::new(&mut self.get_mut().0)
            .poll_next(cx)
            .map(|t| match t {
                Some(Ok(resp)) => Some(Ok(WatchResponse::new(resp))),
                Some(Err(e)) => Some(Err(From::from(e))),
                None => None,
            })
    }
}

impl Stream for WatchStream {
    type Item = Result<WatchResponse>;

    #[inline]
    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        Pin::new(&mut self.get_mut().response_stream).poll_next(cx)
    }
}