1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149

use std::{
    ops::Deref,
    pin::Pin,
};

use futures::{
    channel::mpsc::UnboundedSender,
    task,
    Stream,
    StreamExt,
};
use tokio_stream::wrappers::{
    errors::BroadcastStreamRecvError,
    BroadcastStream,
};

use crate::{
    base_client::{
        FunctionResult,
        QueryResults,
        SubscriberId,
    },
    client::worker::{
        ClientRequest,
        UnsubscribeRequest,
    },
};
#[cfg(doc)]
use crate::{
    ConvexClient,
    Value,
};

/// This structure represents a single subscription to a query with args.
/// For convenience, [`QuerySubscription`] also implements
/// [`Stream`]<[`FunctionResult`]>, giving a stream of results to the query.
///
/// It is returned by [`ConvexClient::subscribe`]. The subscription lives
/// in the active query set for as long as this token stays in scope.
///
/// For a consistent [`QueryResults`] of all your queries, use
/// [`ConvexClient::watch_all()`] instead.
pub struct QuerySubscription {
    pub(super) subscriber_id: SubscriberId,
    pub(super) request_sender: UnboundedSender<ClientRequest>,
    pub(super) watch: BroadcastStream<QueryResults>,
    pub(super) initial: Option<FunctionResult>,
}
impl QuerySubscription {
    /// Returns an identifier for this subscription based on its query and args.
    /// This identifier can be used to find the result within a
    /// [`QuerySetSubscription`] as returned by [`ConvexClient::watch_all()`]
    pub fn id(&self) -> &SubscriberId {
        &self.subscriber_id
    }
}
impl std::fmt::Debug for QuerySubscription {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("QuerySubscription")
            .field("subscriber_id", &self.subscriber_id)
            .finish()
    }
}
impl Deref for QuerySubscription {
    type Target = SubscriberId;

    fn deref(&self) -> &SubscriberId {
        &self.subscriber_id
    }
}
impl Drop for QuerySubscription {
    fn drop(&mut self) {
        let _ = self
            .request_sender
            .unbounded_send(ClientRequest::Unsubscribe(UnsubscribeRequest {
                subscriber_id: self.subscriber_id,
            }));
    }
}
impl Stream for QuerySubscription {
    type Item = FunctionResult;

    fn poll_next(
        mut self: Pin<&mut Self>,
        cx: &mut task::Context<'_>,
    ) -> task::Poll<Option<Self::Item>> {
        if let Some(initial) = self.initial.take() {
            return task::Poll::Ready(Some(initial));
        }
        loop {
            return match self.watch.poll_next_unpin(cx) {
                // Ok to be lagged (skip intermediate values) - since Convex
                // only guarantees a newer value than the previous value.
                task::Poll::Ready(Some(Err(BroadcastStreamRecvError::Lagged(_amt)))) => continue,
                task::Poll::Ready(Some(Ok(map))) => {
                    let Some(value) = map.get(self.id()) else {
                        // No result yet in the query result set. Keep polling.
                        continue;
                    };
                    task::Poll::Ready(Some(value.clone()))
                },
                task::Poll::Ready(None) => task::Poll::Ready(None),
                task::Poll::Pending => task::Poll::Pending,
            };
        }
    }
}

/// A subscription to a consistent view of multiple queries.
///
/// [`QuerySetSubscription`]
/// implements [`Stream`]<[`QueryResults`]>.
/// Each item in the stream contains a consistent view
/// of the results of all the queries in the query set.
///
/// Queries can be added to the query set via [`ConvexClient::subscribe`].
/// Queries can be removed from the query set via dropping the
/// [`QuerySubscription`] token returned by [`ConvexClient::subscribe`].
///
///
/// [`QueryResults`] is a copy-on-write mapping from [`SubscriberId`] to
/// its latest result [`Value`].
pub struct QuerySetSubscription {
    watch: BroadcastStream<QueryResults>,
}
impl QuerySetSubscription {
    pub(super) fn new(watch: BroadcastStream<QueryResults>) -> Self {
        Self { watch }
    }
}
impl Stream for QuerySetSubscription {
    type Item = QueryResults;

    fn poll_next(
        mut self: Pin<&mut Self>,
        cx: &mut task::Context<'_>,
    ) -> task::Poll<Option<Self::Item>> {
        loop {
            return match self.watch.poll_next_unpin(cx) {
                // Ok to be lagged (skip intermediate values) - since Convex
                // only guarantees a newer value than the previous value.
                task::Poll::Ready(Some(Err(BroadcastStreamRecvError::Lagged(_amt)))) => continue,
                task::Poll::Ready(Some(Ok(map))) => task::Poll::Ready(Some(map)),
                task::Poll::Ready(None) => task::Poll::Ready(None),
                task::Poll::Pending => task::Poll::Pending,
            };
        }
    }
}