Use context cancellation tokens in topic actor (#2246)

rogeralsing · web-flow · commit ac5dde9c8e1b · 2025-08-06T16:03:51.000+02:00
diff --git a/src/Proto.Cluster/PubSub/TopicActor.cs b/src/Proto.Cluster/PubSub/TopicActor.cs
@@ -46,7 +46,7 @@ public Task ReceiveAsync(IContext context) =>
             PubSubBatch batch                        => OnPubSubBatch(context, batch),
             NotifyAboutFailingSubscribersRequest msg => OnNotifyAboutFailingSubscribers(context, msg),
             ClusterTopology msg                      => OnClusterTopologyChanged(context, msg),
-            DeadLetterResponse msg                   => OnDeadLetterResponse(msg),
+            DeadLetterResponse msg                   => OnDeadLetterResponse(context, msg),
             _                                        => Task.CompletedTask
         };
 
@@ -55,7 +55,7 @@ private async Task OnStarted(IContext context)
         _topic = context.Get<ClusterIdentity>()!.Identity;
         _topologySubscription = context.System.EventStream.Subscribe<ClusterTopology>(context, context.Self);
 
-        var subs = await LoadSubscriptions(_topic).ConfigureAwait(false);
+        var subs = await LoadSubscriptions(_topic, context.CancellationToken).ConfigureAwait(false);
 
         if (subs.Subscribers_ is not null)
         {
@@ -169,7 +169,7 @@ private static Subscribers GetSubscribersForAddress(
 
     private async Task OnNotifyAboutFailingSubscribers(IContext context, NotifyAboutFailingSubscribersRequest msg)
     {
-        await UnsubscribeUnreachablePidSubscribers(msg.InvalidDeliveries).ConfigureAwait(false);
+        await UnsubscribeUnreachablePidSubscribers(msg.InvalidDeliveries, context.CancellationToken).ConfigureAwait(false);
         LogDeliveryErrors(msg.InvalidDeliveries);
 
         context.Respond(new NotifyAboutFailingSubscribersResponse());
@@ -188,7 +188,7 @@ private void LogDeliveryErrors(IReadOnlyCollection<SubscriberDeliveryReport> all
     }
 
     private async Task UnsubscribeUnreachablePidSubscribers(
-        IReadOnlyCollection<SubscriberDeliveryReport> allInvalidDeliveryReports)
+        IReadOnlyCollection<SubscriberDeliveryReport> allInvalidDeliveryReports, CancellationToken ct)
     {
         var allUnreachable = allInvalidDeliveryReports
             .Where(r => r is
@@ -200,7 +200,7 @@ private async Task UnsubscribeUnreachablePidSubscribers(
             .Select(s => s.Subscriber)
             .ToList();
 
-        await RemoveSubscribers(allUnreachable).ConfigureAwait(false);
+        await RemoveSubscribers(allUnreachable, ct).ConfigureAwait(false);
     }
 
     private async Task OnClusterTopologyChanged(IContext context, ClusterTopology topology)
@@ -214,7 +214,7 @@ private async Task OnClusterTopologyChanged(IContext context, ClusterTopology to
                             addressesThatLeft.Contains(s.Pid.Address))
                 .ToList();
 
-            await RemoveSubscribers(subscribersThatLeft).ConfigureAwait(false);
+            await RemoveSubscribers(subscribersThatLeft, context.CancellationToken).ConfigureAwait(false);
         }
     }
 
@@ -227,10 +227,10 @@ private async Task UnsubscribeSubscribersOnMembersThatLeft(IContext ctx)
                         !activeMemberAddresses.Contains(s.Pid.Address))
             .ToList();
 
-        await RemoveSubscribers(subscribersThatLeft).ConfigureAwait(false);
+        await RemoveSubscribers(subscribersThatLeft, ctx.CancellationToken).ConfigureAwait(false);
     }
 
-    private async Task RemoveSubscribers(IReadOnlyCollection<SubscriberIdentity> subscribersThatLeft)
+    private async Task RemoveSubscribers(IReadOnlyCollection<SubscriberIdentity> subscribersThatLeft, CancellationToken ct)
     {
         if (subscribersThatLeft.Count > 0)
         {
@@ -247,16 +247,15 @@ private async Task RemoveSubscribers(IReadOnlyCollection<SubscriberIdentity> sub
                     string.Join(", ", subscribersThatLeft));
             }
 
-            await SaveSubscriptions(_topic, new Subscribers { Subscribers_ = { _subscribers } }).ConfigureAwait(false);
+            await SaveSubscriptions(_topic, new Subscribers { Subscribers_ = { _subscribers } }, ct).ConfigureAwait(false);
         }
     }
 
-    private async Task<Subscribers> LoadSubscriptions(string topic)
+    private async Task<Subscribers> LoadSubscriptions(string topic, CancellationToken ct)
     {
         try
         {
-            //TODO: cancellation token config?
-            var state = await _subscriptionStore.GetAsync(topic, CancellationToken.None).ConfigureAwait(false);
+            var state = await _subscriptionStore.GetAsync(topic, ct).ConfigureAwait(false);
             Logger.LogDebug("Topic {Topic} loaded subscriptions {Subscriptions}", _topic, state);
 
             return state ?? new Subscribers();
@@ -272,13 +271,12 @@ private async Task<Subscribers> LoadSubscriptions(string topic)
         }
     }
 
-    private async Task SaveSubscriptions(string topic, Subscribers subs)
+    private async Task SaveSubscriptions(string topic, Subscribers subs, CancellationToken ct)
     {
         try
         {
-            //TODO: cancellation token config?
             Logger.LogDebug("Topic {Topic} saved subscriptions {Subscriptions}", _topic, subs);
-            await _subscriptionStore.SetAsync(topic, subs, CancellationToken.None).ConfigureAwait(false);
+            await _subscriptionStore.SetAsync(topic, subs, ct).ConfigureAwait(false);
         }
         catch (Exception e)
         {
@@ -293,26 +291,31 @@ private async Task OnUnsubscribe(IContext context, UnsubscribeRequest unsub)
     {
         _subscribers = _subscribers.Remove(unsub.Subscriber);
         Logger.LogDebug("Topic {Topic} - {Subscriber} unsubscribed", _topic, unsub);
-        await SaveSubscriptions(_topic, new Subscribers { Subscribers_ = { _subscribers } }).ConfigureAwait(false);
+        await SaveSubscriptions(_topic, new Subscribers { Subscribers_ = { _subscribers } }, context.CancellationToken)
+            .ConfigureAwait(false);
         context.Respond(new UnsubscribeResponse());
     }
 
     private async Task OnSubscribe(IContext context, SubscribeRequest sub)
     {
         _subscribers = _subscribers.Add(sub.Subscriber);
         Logger.LogDebug("Topic {Topic} - {Subscriber} subscribed", _topic, sub);
-        await SaveSubscriptions(_topic, new Subscribers { Subscribers_ = { _subscribers } }).ConfigureAwait(false);
+        await SaveSubscriptions(_topic, new Subscribers { Subscribers_ = { _subscribers } }, context.CancellationToken)
+            .ConfigureAwait(false);
         context.Respond(new SubscribeResponse());
     }
 
-    private async Task OnDeadLetterResponse(DeadLetterResponse msg)
+    private async Task OnDeadLetterResponse(IContext context, DeadLetterResponse msg)
     {
-        var deadLetterSub = msg.Target == null ? null : _subscribers.FirstOrDefault(s => s.Pid.Address == msg.Target.Address);
+        var deadLetterSub = msg.Target == null
+            ? null
+            : _subscribers.FirstOrDefault(s => s.Pid.Address == msg.Target.Address);
         if (deadLetterSub != null)
         {
             _subscribers = _subscribers.Remove(deadLetterSub);
             Logger.LogDebug("Topic {Topic} - {Subscriber} unsubscribed due to dead letter", _topic, deadLetterSub);
-            await SaveSubscriptions(_topic, new Subscribers { Subscribers_ = { _subscribers } }).ConfigureAwait(false);
+            await SaveSubscriptions(_topic, new Subscribers { Subscribers_ = { _subscribers } }, context.CancellationToken)
+                .ConfigureAwait(false);
         }
     }
 }

Original file line number	Diff line number	Diff line change
`@@ -46,7 +46,7 @@ public Task ReceiveAsync(IContext context) =>`
`46`	`46`	`PubSubBatch batch => OnPubSubBatch(context, batch),`
`47`	`47`	`NotifyAboutFailingSubscribersRequest msg => OnNotifyAboutFailingSubscribers(context, msg),`
`48`	`48`	`ClusterTopology msg => OnClusterTopologyChanged(context, msg),`
`49`		`- DeadLetterResponse msg => OnDeadLetterResponse(msg),`
	`49`	`+ DeadLetterResponse msg => OnDeadLetterResponse(context, msg),`
`50`	`50`	`_ => Task.CompletedTask`
`51`	`51`	`};`
`52`	`52`
`@@ -55,7 +55,7 @@ private async Task OnStarted(IContext context)`
`55`	`55`	`_topic = context.Get<ClusterIdentity>()!.Identity;`
`56`	`56`	`_topologySubscription = context.System.EventStream.Subscribe<ClusterTopology>(context, context.Self);`
`57`	`57`
`58`		`- var subs = await LoadSubscriptions(_topic).ConfigureAwait(false);`
	`58`	`+ var subs = await LoadSubscriptions(_topic, context.CancellationToken).ConfigureAwait(false);`
`59`	`59`
`60`	`60`	`if (subs.Subscribers_ is not null)`
`61`	`61`	`{`
`@@ -169,7 +169,7 @@ private static Subscribers GetSubscribersForAddress(`
`169`	`169`
`170`	`170`	`private async Task OnNotifyAboutFailingSubscribers(IContext context, NotifyAboutFailingSubscribersRequest msg)`
`171`	`171`	`{`
`172`		`- await UnsubscribeUnreachablePidSubscribers(msg.InvalidDeliveries).ConfigureAwait(false);`
	`172`	`+ await UnsubscribeUnreachablePidSubscribers(msg.InvalidDeliveries, context.CancellationToken).ConfigureAwait(false);`
`173`	`173`	`LogDeliveryErrors(msg.InvalidDeliveries);`
`174`	`174`
`175`	`175`	`context.Respond(new NotifyAboutFailingSubscribersResponse());`
`@@ -188,7 +188,7 @@ private void LogDeliveryErrors(IReadOnlyCollection<SubscriberDeliveryReport> all`
`188`	`188`	`}`
`189`	`189`
`190`	`190`	`private async Task UnsubscribeUnreachablePidSubscribers(`
`191`		`- IReadOnlyCollection<SubscriberDeliveryReport> allInvalidDeliveryReports)`
	`191`	`+ IReadOnlyCollection<SubscriberDeliveryReport> allInvalidDeliveryReports, CancellationToken ct)`
`192`	`192`	`{`
`193`	`193`	`var allUnreachable = allInvalidDeliveryReports`
`194`	`194`	`.Where(r => r is`
`@@ -200,7 +200,7 @@ private async Task UnsubscribeUnreachablePidSubscribers(`
`200`	`200`	`.Select(s => s.Subscriber)`
`201`	`201`	`.ToList();`
`202`	`202`
`203`		`- await RemoveSubscribers(allUnreachable).ConfigureAwait(false);`
	`203`	`+ await RemoveSubscribers(allUnreachable, ct).ConfigureAwait(false);`
`204`	`204`	`}`
`205`	`205`
`206`	`206`	`private async Task OnClusterTopologyChanged(IContext context, ClusterTopology topology)`
`@@ -214,7 +214,7 @@ private async Task OnClusterTopologyChanged(IContext context, ClusterTopology to`
`214`	`214`	`addressesThatLeft.Contains(s.Pid.Address))`
`215`	`215`	`.ToList();`
`216`	`216`
`217`		`- await RemoveSubscribers(subscribersThatLeft).ConfigureAwait(false);`
	`217`	`+ await RemoveSubscribers(subscribersThatLeft, context.CancellationToken).ConfigureAwait(false);`
`218`	`218`	`}`
`219`	`219`	`}`
`220`	`220`
`@@ -227,10 +227,10 @@ private async Task UnsubscribeSubscribersOnMembersThatLeft(IContext ctx)`
`227`	`227`	`!activeMemberAddresses.Contains(s.Pid.Address))`
`228`	`228`	`.ToList();`
`229`	`229`
`230`		`- await RemoveSubscribers(subscribersThatLeft).ConfigureAwait(false);`
	`230`	`+ await RemoveSubscribers(subscribersThatLeft, ctx.CancellationToken).ConfigureAwait(false);`
`231`	`231`	`}`
`232`	`232`
`233`		`- private async Task RemoveSubscribers(IReadOnlyCollection<SubscriberIdentity> subscribersThatLeft)`
	`233`	`+ private async Task RemoveSubscribers(IReadOnlyCollection<SubscriberIdentity> subscribersThatLeft, CancellationToken ct)`
`234`	`234`	`{`
`235`	`235`	`if (subscribersThatLeft.Count > 0)`
`236`	`236`	`{`
`@@ -247,16 +247,15 @@ private async Task RemoveSubscribers(IReadOnlyCollection<SubscriberIdentity> sub`
`247`	`247`	`string.Join(", ", subscribersThatLeft));`
`248`	`248`	`}`
`249`	`249`
`250`		`- await SaveSubscriptions(_topic, new Subscribers { Subscribers_ = { _subscribers } }).ConfigureAwait(false);`
	`250`	`+ await SaveSubscriptions(_topic, new Subscribers { Subscribers_ = { _subscribers } }, ct).ConfigureAwait(false);`
`251`	`251`	`}`
`252`	`252`	`}`
`253`	`253`
`254`		`- private async Task<Subscribers> LoadSubscriptions(string topic)`
	`254`	`+ private async Task<Subscribers> LoadSubscriptions(string topic, CancellationToken ct)`
`255`	`255`	`{`
`256`	`256`	`try`
`257`	`257`	`{`
`258`		`- //TODO: cancellation token config?`
`259`		`- var state = await _subscriptionStore.GetAsync(topic, CancellationToken.None).ConfigureAwait(false);`
	`258`	`+ var state = await _subscriptionStore.GetAsync(topic, ct).ConfigureAwait(false);`
`260`	`259`	`Logger.LogDebug("Topic {Topic} loaded subscriptions {Subscriptions}", _topic, state);`
`261`	`260`
`262`	`261`	`return state ?? new Subscribers();`
`@@ -272,13 +271,12 @@ private async Task<Subscribers> LoadSubscriptions(string topic)`
`272`	`271`	`}`
`273`	`272`	`}`
`274`	`273`
`275`		`- private async Task SaveSubscriptions(string topic, Subscribers subs)`
	`274`	`+ private async Task SaveSubscriptions(string topic, Subscribers subs, CancellationToken ct)`
`276`	`275`	`{`
`277`	`276`	`try`
`278`	`277`	`{`
`279`		`- //TODO: cancellation token config?`
`280`	`278`	`Logger.LogDebug("Topic {Topic} saved subscriptions {Subscriptions}", _topic, subs);`
`281`		`- await _subscriptionStore.SetAsync(topic, subs, CancellationToken.None).ConfigureAwait(false);`
	`279`	`+ await _subscriptionStore.SetAsync(topic, subs, ct).ConfigureAwait(false);`
`282`	`280`	`}`
`283`	`281`	`catch (Exception e)`
`284`	`282`	`{`
`@@ -293,26 +291,31 @@ private async Task OnUnsubscribe(IContext context, UnsubscribeRequest unsub)`
`293`	`291`	`{`
`294`	`292`	`_subscribers = _subscribers.Remove(unsub.Subscriber);`
`295`	`293`	`Logger.LogDebug("Topic {Topic} - {Subscriber} unsubscribed", _topic, unsub);`
`296`		`- await SaveSubscriptions(_topic, new Subscribers { Subscribers_ = { _subscribers } }).ConfigureAwait(false);`
	`294`	`+ await SaveSubscriptions(_topic, new Subscribers { Subscribers_ = { _subscribers } }, context.CancellationToken)`
	`295`	`+ .ConfigureAwait(false);`
`297`	`296`	`context.Respond(new UnsubscribeResponse());`
`298`	`297`	`}`
`299`	`298`
`300`	`299`	`private async Task OnSubscribe(IContext context, SubscribeRequest sub)`
`301`	`300`	`{`
`302`	`301`	`_subscribers = _subscribers.Add(sub.Subscriber);`
`303`	`302`	`Logger.LogDebug("Topic {Topic} - {Subscriber} subscribed", _topic, sub);`
`304`		`- await SaveSubscriptions(_topic, new Subscribers { Subscribers_ = { _subscribers } }).ConfigureAwait(false);`
	`303`	`+ await SaveSubscriptions(_topic, new Subscribers { Subscribers_ = { _subscribers } }, context.CancellationToken)`
	`304`	`+ .ConfigureAwait(false);`
`305`	`305`	`context.Respond(new SubscribeResponse());`
`306`	`306`	`}`
`307`	`307`
`308`		`- private async Task OnDeadLetterResponse(DeadLetterResponse msg)`
	`308`	`+ private async Task OnDeadLetterResponse(IContext context, DeadLetterResponse msg)`
`309`	`309`	`{`
`310`		`- var deadLetterSub = msg.Target == null ? null : _subscribers.FirstOrDefault(s => s.Pid.Address == msg.Target.Address);`
	`310`	`+ var deadLetterSub = msg.Target == null`
	`311`	`+ ? null`
	`312`	`+ : _subscribers.FirstOrDefault(s => s.Pid.Address == msg.Target.Address);`
`311`	`313`	`if (deadLetterSub != null)`
`312`	`314`	`{`
`313`	`315`	`_subscribers = _subscribers.Remove(deadLetterSub);`
`314`	`316`	`Logger.LogDebug("Topic {Topic} - {Subscriber} unsubscribed due to dead letter", _topic, deadLetterSub);`
`315`		`- await SaveSubscriptions(_topic, new Subscribers { Subscribers_ = { _subscribers } }).ConfigureAwait(false);`
	`317`	`+ await SaveSubscriptions(_topic, new Subscribers { Subscribers_ = { _subscribers } }, context.CancellationToken)`
	`318`	`+ .ConfigureAwait(false);`
`316`	`319`	`}`
`317`	`320`	`}`
`318`	`321`	`}`