brokapi_finam/stream.rs
1//! One supervised gRPC server-stream: open, read, reopen with backoff.
2//!
3//! A stream task only reads and forwards frames. Parsing into events happens in the pump.
4//! Closing a stream (dropping it) is the protocol-level unsubscribe for gRPC (RST_STREAM).
5//!
6//! What the supervisor does:
7//!
8//! - **Opens** the stream after taking quota for it (opening counts as a request) and within a deadline.
9//! - **Reopens** it after a break or a normal end, with a pause that doubles from 0.5 s to 30 s. A
10//! stream that lived at least 10 s counts as healthy and the pause starts over.
11//! - **Asks for a new token** when the broker answers `UNAUTHENTICATED`.
12//! - **Treats some refusals as permanent** (`INVALID_ARGUMENT`, `NOT_FOUND`, `PERMISSION_DENIED`,
13//! `UNIMPLEMENTED`, `FAILED_PRECONDITION`): it reports `Failed` once and keeps trying slowly (every 30 s),
14//! because a permission may be granted later.
15//! - **Reports** every step as a [`StreamEvent`] through the `wrap` function.
16//!
17//! # Examples
18//!
19//! ```
20//! use std::sync::Arc;
21//! use std::time::Duration;
22//! use brokapi_core::SubscriptionId;
23//! use brokapi_finam::core::StreamEvent;
24//! use brokapi_finam::limits::default_limits;
25//! use brokapi_finam::stream::{OpenFn, StreamParams, spawn_stream};
26//! use brokapi_runtime::BrokerLimiter;
27//! use brokapi_finam_protocol::pb::grpc::tradeapi::v1::marketdata as pm;
28//! use brokapi_finam_protocol::{ApiEndpoint, FinamApi, start_token_manager};
29//! use tokio::sync::mpsc;
30//! use tokio_util::sync::CancellationToken;
31//! # use brokapi_mocks::finam::{MockServer, SECRET};
32//!
33//! # #[tokio::main(flavor = "current_thread")]
34//! # async fn main() -> Result<(), Box<dyn std::error::Error>> {
35//! # let server = MockServer::start().await;
36//! # let endpoint = ApiEndpoint::new(server.url.clone());
37//! let tokens = start_token_manager(&endpoint, SECRET, "").await?;
38//! let api = FinamApi::connect(&endpoint, Arc::clone(&tokens), false)?;
39//!
40//! // How to open the stream. The supervisor calls it again after every reconnect.
41//! let open: OpenFn<pm::SubscribeOrderBookResponse> = Arc::new(move || {
42//! let api = api.clone();
43//! Box::pin(async move {
44//! let request = pm::SubscribeOrderBookRequest { symbol: "SBER@MISX".into() };
45//! api.market_data().subscribe_order_book(request).await.map(|r| r.into_inner())
46//! })
47//! });
48//!
49//! let (tx, mut frames) = mpsc::channel(16);
50//! let handle = spawn_stream(
51//! StreamParams {
52//! sub: SubscriptionId::new(1),
53//! method: "SubscribeOrderBook",
54//! limiter: BrokerLimiter::new(&default_limits()),
55//! tokens,
56//! open,
57//! open_timeout: Duration::from_secs(5),
58//! wrap: Arc::new(|event| event), // keep the events as they are
59//! tx,
60//! first: None,
61//! },
62//! CancellationToken::new(),
63//! );
64//!
65//! let first = tokio::time::timeout(Duration::from_secs(5), frames.recv()).await?.unwrap();
66//! assert!(matches!(first, StreamEvent::Connected));
67//!
68//! assert!(handle.stop(Duration::from_secs(2)).await);
69//! # Ok(())
70//! # }
71//! ```
72
73use std::sync::Arc;
74use std::time::Duration;
75
76use brokapi_core::SubscriptionId;
77use brokapi_finam_protocol::{FinamError, TokenManager};
78use brokapi_runtime::BrokerLimiter;
79use futures_util::future::BoxFuture;
80use tokio::sync::{Notify, mpsc, oneshot};
81use tokio::task::JoinHandle;
82use tokio_util::sync::CancellationToken;
83use tonic::{Code, Streaming};
84use tracing::{debug, warn};
85
86use crate::core::StreamEvent;
87
88const BACKOFF_MIN: Duration = Duration::from_millis(500);
89const BACKOFF_MAX: Duration = Duration::from_secs(30);
90/// A stream that lived this long counts as healthy: the next drop restarts the backoff.
91const HEALTHY_AFTER: Duration = Duration::from_secs(10);
92
93/// How to open a stream: a function that returns the future of a new stream. Called again after
94/// every drop, so it must build a fresh request each time. It reads the token through the
95/// [`FinamApi`](brokapi_finam_protocol::FinamApi) it captured.
96pub type OpenFn<M> =
97 Arc<dyn Fn() -> BoxFuture<'static, Result<Streaming<M>, tonic::Status>> + Send + Sync>;
98
99/// Result of the first open attempt, reported once to `subscribe_*`.
100pub type FirstOutcome = oneshot::Sender<Result<(), FinamError>>;
101
102/// Everything [`spawn_stream`] needs. `M` is the message type of the stream, `F` the frame type the
103/// connector wants.
104pub struct StreamParams<M, F> {
105 /// The subscription this stream serves. Only used for logging here.
106 pub sub: SubscriptionId,
107 /// Limiter key of the open call (`SubscribeOrders`, ...).
108 pub method: &'static str,
109 /// Rate limiter. A refused open waits for the time the limiter gives and tries again.
110 pub limiter: Arc<BrokerLimiter>,
111 /// Token manager. The supervisor asks it for a renewal after `UNAUTHENTICATED`.
112 pub tokens: Arc<TokenManager>,
113 /// How to open the stream.
114 pub open: OpenFn<M>,
115 /// How long one open attempt may take before it counts as a failure.
116 pub open_timeout: Duration,
117 /// Turns a [`StreamEvent`] into the connector's frame.
118 pub wrap: Arc<dyn Fn(StreamEvent<M>) -> F + Send + Sync>,
119 /// Where the frames go. When the receiver is gone, the supervisor ends.
120 pub tx: mpsc::Sender<F>,
121 /// Receives the outcome of the FIRST open attempt, once. Used by `subscribe_*` to return a refusal.
122 pub first: Option<FirstOutcome>,
123}
124
125/// Handle of a running stream supervisor: restart it, stop it, or abort it.
126///
127/// Dropping the handle does not stop the stream. Call [`stop`](Self::stop) or [`abort`](Self::abort),
128/// or cancel the token given to [`spawn_stream`].
129pub struct StreamHandle {
130 cancel: CancellationToken,
131 restart: Arc<Notify>,
132 join: JoinHandle<()>,
133}
134
135impl StreamHandle {
136 /// Drop the current stream and open a fresh one right away (book out of sync).
137 ///
138 /// The pause of the backoff starts over. The new stream is reported as `Connected` again.
139 pub fn restart(&self) {
140 self.restart.notify_one();
141 }
142
143 /// Cancels the stream (protocol-level unsubscribe) and waits for the task, at most `limit`.
144 /// `true` when the task ended in time.
145 ///
146 /// # Examples
147 ///
148 /// See the [module docs](self).
149 pub async fn stop(self, limit: Duration) -> bool {
150 self.cancel.cancel();
151 tokio::time::timeout(limit, self.join).await.is_ok()
152 }
153
154 /// Cancels the stream and kills the task at once, without waiting for it to finish its step.
155 pub fn abort(self) {
156 self.cancel.cancel();
157 self.join.abort();
158 }
159}
160
161/// Starts the supervisor task. `cancel` should be a child of the connector's root token.
162///
163/// Must be called inside a Tokio runtime. See the [module docs](self) for a complete example.
164pub fn spawn_stream<M, F>(params: StreamParams<M, F>, cancel: CancellationToken) -> StreamHandle
165where
166 M: Send + 'static,
167 F: Send + 'static,
168{
169 let restart = Arc::new(Notify::new());
170 let join = tokio::spawn(run(params, cancel.clone(), Arc::clone(&restart)));
171 StreamHandle {
172 cancel,
173 restart,
174 join,
175 }
176}
177
178fn is_permanent(code: Code) -> bool {
179 matches!(
180 code,
181 Code::InvalidArgument
182 | Code::NotFound
183 | Code::PermissionDenied
184 | Code::Unimplemented
185 | Code::FailedPrecondition
186 )
187}
188
189async fn forward<M, F>(
190 p: &StreamParams<M, F>,
191 cancel: &CancellationToken,
192 ev: StreamEvent<M>,
193) -> bool {
194 let frame = (p.wrap)(ev);
195 tokio::select! {
196 _ = cancel.cancelled() => false,
197 r = p.tx.send(frame) => r.is_ok(),
198 }
199}
200
201async fn sleep_or_cancel(d: Duration, cancel: &CancellationToken) -> bool {
202 tokio::select! {
203 _ = cancel.cancelled() => false,
204 _ = tokio::time::sleep(d) => true,
205 }
206}
207
208async fn run<M, F>(mut p: StreamParams<M, F>, cancel: CancellationToken, restart: Arc<Notify>)
209where
210 M: Send + 'static,
211 F: Send + 'static,
212{
213 let mut backoff = BACKOFF_MIN;
214 let mut reported_down = false;
215 loop {
216 if cancel.is_cancelled() {
217 return;
218 }
219 // Opening a stream spends quota like any call.
220 if let Err(e) = p.limiter.acquire(p.method, brokapi_runtime::now()) {
221 let wait = match e {
222 brokapi_core::Error::RateLimited { retry_after } => retry_after,
223 _ => BACKOFF_MAX,
224 };
225 if !sleep_or_cancel(wait, &cancel).await {
226 return;
227 }
228 continue;
229 }
230
231 let opened = tokio::select! {
232 _ = cancel.cancelled() => return,
233 r = tokio::time::timeout(p.open_timeout, (p.open)()) => r,
234 };
235 let status = match opened {
236 Ok(Ok(stream)) => {
237 reported_down = false;
238 if let Some(first) = p.first.take() {
239 let _ = first.send(Ok(()));
240 }
241 if !forward(&p, &cancel, StreamEvent::Connected).await {
242 return;
243 }
244 let started = brokapi_runtime::now();
245 match read(&mut p, stream, &cancel, &restart).await {
246 ReadEnd::Cancelled => return,
247 ReadEnd::Restart => {
248 backoff = BACKOFF_MIN;
249 continue;
250 }
251 ReadEnd::Dropped(why, unauth) => {
252 if started.elapsed() >= HEALTHY_AFTER {
253 backoff = BACKOFF_MIN;
254 }
255 if !forward(&p, &cancel, StreamEvent::Disconnected(why.clone())).await {
256 return;
257 }
258 reported_down = true;
259 debug!(sub = ?p.sub, reason = %why, "stream dropped, reopening");
260 // An unauthenticated drop usually means the JWT lapsed.
261 if unauth {
262 p.tokens.request_refresh();
263 }
264 None
265 }
266 }
267 }
268 Ok(Err(status)) => Some(status),
269 Err(_) => Some(tonic::Status::deadline_exceeded(format!(
270 "stream not opened within {:?}",
271 p.open_timeout
272 ))),
273 };
274
275 if let Some(status) = status {
276 let err: FinamError = status.clone().into();
277 if status.code() == Code::Unauthenticated {
278 p.tokens.request_refresh();
279 }
280 if is_permanent(status.code()) {
281 warn!(sub = ?p.sub, error = %err, "stream refused");
282 if let Some(first) = p.first.take() {
283 let _ = first.send(Err(err.clone()));
284 }
285 if !forward(&p, &cancel, StreamEvent::Failed(err.to_string())).await {
286 return;
287 }
288 // The refusal may be lifted later (permissions): keep trying, slowly.
289 backoff = BACKOFF_MAX;
290 } else if !reported_down {
291 reported_down = true;
292 if !forward(&p, &cancel, StreamEvent::Disconnected(err.to_string())).await {
293 return;
294 }
295 }
296 }
297
298 if !sleep_or_cancel(backoff, &cancel).await {
299 return;
300 }
301 backoff = (backoff * 2).min(BACKOFF_MAX);
302 }
303}
304
305enum ReadEnd {
306 Cancelled,
307 Restart,
308 Dropped(String, bool),
309}
310
311async fn read<M, F>(
312 p: &mut StreamParams<M, F>,
313 mut stream: Streaming<M>,
314 cancel: &CancellationToken,
315 restart: &Notify,
316) -> ReadEnd
317where
318 M: Send + 'static,
319 F: Send + 'static,
320{
321 loop {
322 tokio::select! {
323 _ = cancel.cancelled() => return ReadEnd::Cancelled,
324 _ = restart.notified() => return ReadEnd::Restart,
325 msg = stream.message() => match msg {
326 Ok(Some(m)) => {
327 if !forward(p, cancel, StreamEvent::Msg(m)).await {
328 return ReadEnd::Cancelled;
329 }
330 }
331 Ok(None) => return ReadEnd::Dropped("stream ended by server".into(), false),
332 Err(status) => {
333 let unauth = status.code() == Code::Unauthenticated;
334 let e: FinamError = status.into();
335 return ReadEnd::Dropped(e.to_string(), unauth);
336 }
337 },
338 }
339 }
340}