Skip to content

Commit 62f085b

Browse files
authored
Fix missed events race when creating subscriptions (#1023)
Previously, we fetched the initial state *before* registering the new subscription. Any events emitted after the DB read but before the subscription was installed were dropped—most visible under low-resource conditions (e.g., CI). Change: - Register the subscription first, then asynchronously fetch and send the initial state (spawned task). This eliminates the window where events could be missed. - Require `F: Send + Sync` and store `on_new_subscription` as `Arc<F>` so it can be safely used from the spawned task. Result: - No gap between “subscribe” and “start receiving,” avoiding lost events. - Initial state still delivered, now via a background task.
1 parent 6067242 commit 62f085b

1 file changed

Lines changed: 39 additions & 26 deletions

File tree

crates/cdk/src/pub_sub.rs

Lines changed: 39 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -41,10 +41,10 @@ pub struct Manager<T, I, F>
4141
where
4242
T: Indexable<Type = I> + Clone + Send + Sync + 'static,
4343
I: PartialOrd + Clone + Debug + Ord + Send + Sync + 'static,
44-
F: OnNewSubscription<Index = I, Event = T> + 'static,
44+
F: OnNewSubscription<Index = I, Event = T> + Send + Sync + 'static,
4545
{
4646
indexes: IndexTree<T, I>,
47-
on_new_subscription: Option<F>,
47+
on_new_subscription: Option<Arc<F>>,
4848
unsubscription_sender: mpsc::Sender<(SubId, Vec<Index<I>>)>,
4949
active_subscriptions: Arc<AtomicUsize>,
5050
background_subscription_remover: Option<JoinHandle<()>>,
@@ -54,7 +54,7 @@ impl<T, I, F> Default for Manager<T, I, F>
5454
where
5555
T: Indexable<Type = I> + Clone + Send + Sync + 'static,
5656
I: PartialOrd + Clone + Debug + Ord + Send + Sync + 'static,
57-
F: OnNewSubscription<Index = I, Event = T> + 'static,
57+
F: OnNewSubscription<Index = I, Event = T> + Send + Sync + 'static,
5858
{
5959
fn default() -> Self {
6060
let (sender, receiver) = mpsc::channel(DEFAULT_REMOVE_SIZE);
@@ -79,11 +79,11 @@ impl<T, I, F> From<F> for Manager<T, I, F>
7979
where
8080
T: Indexable<Type = I> + Clone + Send + Sync + 'static,
8181
I: PartialOrd + Clone + Debug + Ord + Send + Sync + 'static,
82-
F: OnNewSubscription<Index = I, Event = T> + 'static,
82+
F: OnNewSubscription<Index = I, Event = T> + Send + Sync + 'static,
8383
{
8484
fn from(value: F) -> Self {
8585
let mut manager: Self = Default::default();
86-
manager.on_new_subscription = Some(value);
86+
manager.on_new_subscription = Some(Arc::new(value));
8787
manager
8888
}
8989
}
@@ -92,7 +92,7 @@ impl<T, I, F> Manager<T, I, F>
9292
where
9393
T: Indexable<Type = I> + Clone + Send + Sync + 'static,
9494
I: PartialOrd + Clone + Debug + Ord + Send + Sync + 'static,
95-
F: OnNewSubscription<Index = I, Event = T> + 'static,
95+
F: OnNewSubscription<Index = I, Event = T> + Send + Sync + 'static,
9696
{
9797
#[inline]
9898
/// Broadcast an event to all listeners
@@ -143,32 +143,45 @@ where
143143
indexes: Vec<Index<I>>,
144144
) -> ActiveSubscription<T, I> {
145145
let (sender, receiver) = mpsc::channel(10);
146-
if let Some(on_new_subscription) = self.on_new_subscription.as_ref() {
147-
match on_new_subscription
148-
.on_new_subscription(&indexes.iter().map(|x| x.deref()).collect::<Vec<_>>())
149-
.await
150-
{
151-
Ok(events) => {
152-
for event in events {
153-
let _ = sender.try_send((sub_id.clone(), event));
154-
}
155-
}
156-
Err(err) => {
157-
tracing::info!(
158-
"Failed to get initial state for subscription: {:?}, {}",
159-
sub_id,
160-
err
161-
);
162-
}
163-
}
164-
}
165146

166147
let mut index_storage = self.indexes.write().await;
148+
// Subscribe to events as soon as possible
167149
for index in indexes.clone() {
168150
index_storage.insert(index, sender.clone());
169151
}
170152
drop(index_storage);
171153

154+
if let Some(on_new_subscription) = self.on_new_subscription.clone() {
155+
// After we're subscribed already, fetch the current status of matching events. It is
156+
// down in another thread to return right away
157+
let indexes_for_worker = indexes.clone();
158+
let sub_id_for_worker = sub_id.clone();
159+
tokio::spawn(async move {
160+
match on_new_subscription
161+
.on_new_subscription(
162+
&indexes_for_worker
163+
.iter()
164+
.map(|x| x.deref())
165+
.collect::<Vec<_>>(),
166+
)
167+
.await
168+
{
169+
Ok(events) => {
170+
for event in events {
171+
let _ = sender.try_send((sub_id_for_worker.clone(), event));
172+
}
173+
}
174+
Err(err) => {
175+
tracing::info!(
176+
"Failed to get initial state for subscription: {:?}, {}",
177+
sub_id_for_worker,
178+
err
179+
);
180+
}
181+
}
182+
});
183+
}
184+
172185
self.active_subscriptions
173186
.fetch_add(1, atomic::Ordering::Relaxed);
174187

@@ -232,7 +245,7 @@ impl<T, I, F> Drop for Manager<T, I, F>
232245
where
233246
T: Indexable<Type = I> + Clone + Send + Sync + 'static,
234247
I: Clone + Debug + PartialOrd + Ord + Send + Sync + 'static,
235-
F: OnNewSubscription<Index = I, Event = T> + 'static,
248+
F: OnNewSubscription<Index = I, Event = T> + Send + Sync + 'static,
236249
{
237250
fn drop(&mut self) {
238251
if let Some(handler) = self.background_subscription_remover.take() {

0 commit comments

Comments
 (0)