finance_query/streaming/
news.rs1use std::collections::{HashSet, VecDeque};
7use std::pin::Pin;
8use std::task::{Context, Poll};
9use std::time::Duration;
10
11use futures::stream::Stream;
12use tokio::sync::{broadcast, mpsc};
13use tracing::warn;
14
15use super::subscription::Subscription;
16use crate::feeds::{FeedEntry, FeedSource, fetch_all};
17
18const DEFAULT_POLL_INTERVAL: Duration = Duration::from_secs(300);
20
21const CHANNEL_CAPACITY: usize = 256;
23
24const SEEN_CAP: usize = 2000;
27
28enum FeedCommand {
29 AddSources(Vec<FeedSource>),
30 RemoveSources(Vec<String>),
31 Close,
32}
33
34pub struct NewsStream {
61 inner: Subscription<FeedEntry, FeedCommand>,
62}
63
64impl NewsStream {
65 pub async fn subscribe(sources: impl IntoIterator<Item = FeedSource>) -> Self {
69 NewsStreamBuilder::new().sources(sources).build().await
70 }
71
72 async fn start(sources: Vec<FeedSource>, poll_interval: Duration) -> Self {
73 let inner = Subscription::start(CHANNEL_CAPACITY, 32, move |broadcast_tx, command_rx| {
74 run_feed_loop(sources, poll_interval, broadcast_tx, command_rx)
75 });
76 NewsStream { inner }
77 }
78
79 pub fn resubscribe(&self) -> Self {
83 NewsStream {
84 inner: self.inner.resubscribe(),
85 }
86 }
87
88 pub async fn add_sources(&self, sources: impl IntoIterator<Item = FeedSource>) {
102 self.inner
103 .send(FeedCommand::AddSources(sources.into_iter().collect()))
104 .await;
105 }
106
107 pub async fn remove_sources(&self, sources: impl IntoIterator<Item = FeedSource>) {
121 let urls = sources.into_iter().map(|s| s.url()).collect();
122 self.inner.send(FeedCommand::RemoveSources(urls)).await;
123 }
124
125 pub async fn close(&self) {
127 self.inner.send(FeedCommand::Close).await;
128 }
129}
130
131impl Stream for NewsStream {
132 type Item = FeedEntry;
133
134 fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
135 Pin::new(&mut self.inner).poll_next(cx)
136 }
137}
138
139pub struct NewsStreamBuilder {
141 sources: Vec<FeedSource>,
142 poll_interval: Duration,
143}
144
145impl NewsStreamBuilder {
146 pub fn new() -> Self {
148 Self {
149 sources: Vec::new(),
150 poll_interval: DEFAULT_POLL_INTERVAL,
151 }
152 }
153
154 pub fn sources(mut self, sources: impl IntoIterator<Item = FeedSource>) -> Self {
156 self.sources.extend(sources);
157 self
158 }
159
160 pub fn poll_interval(mut self, interval: Duration) -> Self {
162 self.poll_interval = interval;
163 self
164 }
165
166 pub async fn build(self) -> NewsStream {
168 NewsStream::start(self.sources, self.poll_interval).await
169 }
170}
171
172impl Default for NewsStreamBuilder {
173 fn default() -> Self {
174 Self::new()
175 }
176}
177
178async fn run_feed_loop(
179 initial_sources: Vec<FeedSource>,
180 poll_interval: Duration,
181 broadcast_tx: broadcast::Sender<FeedEntry>,
182 mut command_rx: mpsc::Receiver<FeedCommand>,
183) {
184 let mut sources: Vec<(String, FeedSource)> =
185 initial_sources.into_iter().map(|s| (s.url(), s)).collect();
186 let mut seen: HashSet<String> = HashSet::new();
187 let mut seen_order: VecDeque<String> = VecDeque::new();
188
189 let mut ticker = tokio::time::interval(poll_interval);
190 ticker.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
191
192 loop {
193 tokio::select! {
194 _ = ticker.tick() => {
195 if sources.is_empty() {
196 continue;
197 }
198 let batch = sources.iter().map(|(_, s)| s.clone());
199 match fetch_all(batch).await {
200 Ok(entries) => {
201 for entry in entries {
202 if seen.insert(entry.url.clone()) {
203 seen_order.push_back(entry.url.clone());
204 if seen_order.len() > SEEN_CAP
205 && let Some(old) = seen_order.pop_front()
206 {
207 seen.remove(&old);
208 }
209 let _ = broadcast_tx.send(entry);
210 }
211 }
212 }
213 Err(e) => warn!("news stream poll failed: {e}"),
214 }
215 }
216 cmd = command_rx.recv() => {
217 match cmd {
218 Some(FeedCommand::AddSources(new_sources)) => {
219 for s in new_sources {
220 let url = s.url();
221 if let Some(existing) = sources.iter_mut().find(|(u, _)| *u == url) {
222 existing.1 = s;
223 } else {
224 sources.push((url, s));
225 }
226 }
227 }
228 Some(FeedCommand::RemoveSources(urls)) => {
229 sources.retain(|(u, _)| !urls.contains(u));
230 }
231 Some(FeedCommand::Close) | None => break,
232 }
233 }
234 }
235 }
236}
237
238#[cfg(test)]
239mod tests {
240 use super::*;
241 use futures::StreamExt;
242
243 #[tokio::test]
244 async fn add_and_remove_sources_are_accepted() {
245 let stream = NewsStream::subscribe([]).await;
248 stream.add_sources([FeedSource::Bloomberg]).await;
249 stream.remove_sources([FeedSource::Bloomberg]).await;
250 stream.close().await;
251 }
252
253 #[tokio::test]
254 async fn resubscribe_gives_an_independent_receiver() {
255 let stream = NewsStream::subscribe([]).await;
256 let other = stream.resubscribe();
257 drop(other);
258 stream.close().await;
259 }
260
261 #[tokio::test]
262 async fn close_ends_the_stream() {
263 let mut stream = NewsStreamBuilder::new()
264 .poll_interval(Duration::from_millis(20))
265 .build()
266 .await;
267 stream.close().await;
268 let timeout = tokio::time::timeout(Duration::from_secs(2), stream.next()).await;
269 assert!(matches!(timeout, Ok(None)));
270 }
271}