feat: improve sequence activity handling

Author: BEagle1984

docs/changelog/500.md

@@ -100,7 +100,8 @@ In the following sections, you'll find more a more detailed and canonical list o
 * Add `IgnoreNoMatchingSubscribersError` setting to MQTT producer endpoints, to prevent throwing when no subscriber is consuming the produced message
 * Basic support for MQTT 5 request-response
 * Add <xref:Silverback.Messaging.Subscribers.ConsumerNameFilterAttribute> to filter the subscribers by consumer name (replaces <xref:Silverback.Messaging.Subscribers.KafkaGroupIdFilterAttribute> and <xref:Silverback.Messaging.Subscribers.MqttClientIdFilterAttribute>)
-* Add new built-in (de-)serializers to support raw `string`, `Stream`, `byte[]` and `JsonDocument` (see <xref:serialization> and <xref:deserialization>)
+* Add new built-in (de-)serializers to support raw `string`, `Stream` and `byte[]` (see <xref:serialization> and <xref:deserialization>)
+* Improve [Activity](https://docs.microsoft.com/en-us/dotnet/api/system.diagnostics.activity) handling for sequences (batch processing, chunking, etc.)
 * Upgrade to [Confluent.Kafka 2.8.0](https://github.com/confluentinc/confluent-kafka-dotnet/releases/tag/v2.2.0)
 * Upgrade to [MQTTnet 5.0.1.1416](https://github.com/chkr1011/MQTTnet/releases/tag/v5.0.1.1416)
 

src/Silverback.Core/Messaging/Messages/IMessageStreamEnumerable.cs

@@ -13,25 +13,30 @@ namespace Silverback.Messaging.Messages;
 internal interface IMessageStreamEnumerable
 {
     /// <summary>
-    ///     Adds the specified message to the stream. The returned <see cref="ValueTask" /> will complete only when the
-    ///     message has actually been pulled and processed.
+    ///     Adds the specified message to the stream. The returned <see cref="ValueTask" /> will complete only when the message has actually
+    ///     been pulled and processed.
     /// </summary>
     /// <param name="message">
     ///     The message to be added.
     /// </param>
+    /// <param name="onPullAction">
+    ///     An action to be executed when the message is pulled.
+    /// </param>
+    /// <param name="onPullActionArgument">
+    ///     The argument to be passed to the <paramref name="onPullAction" />.
+    /// </param>
     /// <param name="cancellationToken">
     ///     A <see cref="CancellationToken" /> used to cancel the operation.
     /// </param>
     /// <returns>
     ///     A <see cref="Task" /> representing the asynchronous operation. The  <see cref="Task" /> will complete
     ///     only when the message has actually been pulled and processed.
     /// </returns>
-    Task PushAsync(object message, CancellationToken cancellationToken = default);
+    Task PushAsync(object message, Action<object?>? onPullAction, object? onPullActionArgument, CancellationToken cancellationToken);
 
     /// <summary>
-    ///     Aborts the ongoing enumeration and the pending calls to <see cref="PushAsync" />, then marks the
-    ///     stream as complete. Calling this method will cause an <see cref="OperationCanceledException" /> to be
-    ///     thrown by the enumerator and the <see cref="PushAsync" /> method.
+    ///     Aborts the ongoing enumeration and the pending calls to <see cref="PushAsync" />, then marks the stream as complete. Calling this
+    ///     method will cause an <see cref="OperationCanceledException" /> to be thrown by the enumerator and the <see cref="PushAsync" /> method.
     /// </summary>
     void Abort();
 

src/Silverback.Core/Messaging/Messages/MessageStreamEnumerable`1.cs

@@ -37,12 +37,16 @@ internal sealed class MessageStreamEnumerable<TMessage> : IMessageStreamEnumerab
 
     private bool _isComplete;
 
-    Task IMessageStreamEnumerable.PushAsync(object message, CancellationToken cancellationToken) =>
-        PushAsync((TMessage)message, cancellationToken);
+    private Action<object?>? _onPullAction;
 
-    /// <inheritdoc cref="IMessageStreamEnumerable.PushAsync(object,System.Threading.CancellationToken)" />
+    private object? _onPullActionArgument;
+
+    Task IMessageStreamEnumerable.PushAsync(object message, Action<object?>? onPullAction, object? onPullActionArgument, CancellationToken cancellationToken) =>
+        PushAsync((TMessage)message, onPullAction, onPullActionArgument, cancellationToken);
+
+    /// <inheritdoc cref="IMessageStreamEnumerable.PushAsync(object,Action{object},object,CancellationToken)" />
     [SuppressMessage("ReSharper", "InconsistentlySynchronizedField", Justification = "The lock is important to avoid multiple complete/abort, here is not important")]
-    public async Task PushAsync(TMessage message, CancellationToken cancellationToken = default)
+    public async Task PushAsync(TMessage message, Action<object?>? onPullAction = null, object? onPullActionArgument = null, CancellationToken cancellationToken = default)
     {
         Check.NotNull(message, nameof(message));
 
@@ -55,6 +59,8 @@ public async Task PushAsync(TMessage message, CancellationToken cancellationToke
 
         _current = message;
         _hasCurrent = true;
+        _onPullAction = onPullAction;
+        _onPullActionArgument = onPullActionArgument;
         SafelyRelease(_readSemaphore);
 
         await _processedSemaphore.WaitAsync(linkedTokenSource.Token).ConfigureAwait(false);
@@ -137,6 +143,8 @@ private IEnumerable<TMessage> GetEnumerable()
             if (!_hasCurrent || _current == null)
                 continue;
 
+            _onPullAction?.Invoke(_onPullActionArgument);
+
             yield return _current;
         }
     }
@@ -149,6 +157,8 @@ private async IAsyncEnumerable<TMessage> GetAsyncEnumerable([EnumeratorCancellat
             if (!_hasCurrent || _current == null)
                 continue;
 
+            _onPullAction?.Invoke(_onPullActionArgument);
+
             yield return _current;
         }
     }

src/Silverback.Core/Messaging/Messages/MessageStreamProvider.cs

@@ -40,10 +40,9 @@ public abstract ILazyMessageStreamEnumerable<object> CreateLazyStream(
     public abstract void AbortIfPending();
 
     /// <summary>
-    ///     Aborts the ongoing enumerations and the pending calls to
-    ///     <see cref="MessageStreamProvider{TMessage}.PushAsync(TMessage,System.Threading.CancellationToken)" />, then marks the
-    ///     stream as complete. Calling this method will cause an <see cref="OperationCanceledException" /> to be
-    ///     thrown by the enumerators and the <see cref="MessageStreamProvider{TMessage}.PushAsync(TMessage,System.Threading.CancellationToken)" /> method.
+    ///     Aborts the ongoing enumerations and the pending calls to <see cref="MessageStreamProvider{TMessage}.PushAsync" />, then marks the
+    ///     stream as complete. Calling this method will cause an <see cref="OperationCanceledException" /> to be thrown by the enumerators
+    ///     and the <see cref="MessageStreamProvider{TMessage}.PushAsync" /> method.
     /// </summary>
     public abstract void Abort();
 

src/Silverback.Core/Messaging/Messages/MessageStreamProvider`1.cs

@@ -37,33 +37,20 @@ internal sealed class MessageStreamProvider<TMessage> : MessageStreamProvider
     public override int StreamsCount => _lazyStreams.Count;
 
     /// <summary>
-    ///     Adds the specified message to the stream. The returned <see cref="Task" /> will complete only when the
-    ///     message has actually been pulled and processed.
+    ///     Adds the specified message to the stream. The returned <see cref="Task" /> will complete only when the message has actually been
+    ///     pulled and processed.
     /// </summary>
     /// <param name="message">
     ///     The message to be added.
     /// </param>
-    /// <param name="cancellationToken">
-    ///     A <see cref="CancellationToken" /> used to cancel the operation.
+    /// <param name="throwIfUnhandled">
+    ///     A boolean value indicating whether an exception must be thrown if no subscriber is handling the message.
     /// </param>
-    /// <returns>
-    ///     A <see cref="Task{TResult}" /> representing the asynchronous operation. The task will complete only
-    ///     when the message has actually been pulled and processed and its result contains the number of
-    ///     <see cref="IMessageStreamEnumerable{TMessage}" /> that have been pushed.
-    /// </returns>
-    public Task<int> PushAsync(TMessage message, CancellationToken cancellationToken = default) =>
-        PushAsync(message, true, cancellationToken);
-
-    /// <summary>
-    ///     Adds the specified message to the stream. The returned <see cref="Task" /> will complete only when the
-    ///     message has actually been pulled and processed.
-    /// </summary>
-    /// <param name="message">
-    ///     The message to be added.
+    /// <param name="onPullAction">
+    ///     An action to be executed when the message is pulled.
     /// </param>
-    /// <param name="throwIfUnhandled">
-    ///     A boolean value indicating whether an exception must be thrown if no subscriber is handling the
-    ///     message.
+    /// <param name="onPullActionArgument">
+    ///     The argument to be passed to the <paramref name="onPullAction" />.
     /// </param>
     /// <param name="cancellationToken">
     ///     A <see cref="CancellationToken" /> used to cancel the operation.
@@ -74,14 +61,19 @@ public Task<int> PushAsync(TMessage message, CancellationToken cancellationToken
     ///     <see cref="IMessageStreamEnumerable{TMessage}" /> that have been pushed.
     /// </returns>
     [SuppressMessage("ReSharper", "ParameterOnlyUsedForPreconditionCheck.Global", Justification = "False positive")]
-    public async Task<int> PushAsync(TMessage message, bool throwIfUnhandled, CancellationToken cancellationToken = default)
+    public async Task<int> PushAsync(
+        TMessage message,
+        bool throwIfUnhandled = true,
+        Action<object?>? onPullAction = null,
+        object? onPullActionArgument = null,
+        CancellationToken cancellationToken = default)
     {
         Check.NotNull(message, nameof(message));
 
         if (_completed || _aborted)
             throw new InvalidOperationException("The streams are already completed or aborted.");
 
-        List<Task> processingTasks = PushToCompatibleStreams(message, cancellationToken).ToList();
+        List<Task> processingTasks = PushToCompatibleStreams(message, onPullAction, onPullActionArgument, cancellationToken).ToList();
 
         if (processingTasks.Count > 0)
             await Task.WhenAll(processingTasks).ConfigureAwait(false);
@@ -199,6 +191,8 @@ private static LazyMessageStreamEnumerable<TMessageLinked> CreateLazyStreamCore<
     private static bool PushIfCompatibleType(
         ILazyMessageStreamEnumerable lazyStream,
         TMessage message,
+        Action<object?>? onPullAction,
+        object? onPullActionArgument,
         CancellationToken cancellationToken,
         out Task messageProcessingTask)
     {
@@ -210,14 +204,14 @@ private static bool PushIfCompatibleType(
 
         if (lazyStream.MessageType.IsInstanceOfType(message))
         {
-            messageProcessingTask = lazyStream.GetOrCreateStream().PushAsync(message, cancellationToken);
+            messageProcessingTask = lazyStream.GetOrCreateStream().PushAsync(message, onPullAction, onPullActionArgument, cancellationToken);
             return true;
         }
 
         if (message is IEnvelope { Message: not null } envelope &&
             lazyStream.MessageType.IsInstanceOfType(envelope.Message))
         {
-            messageProcessingTask = lazyStream.GetOrCreateStream().PushAsync(envelope.Message, cancellationToken);
+            messageProcessingTask = lazyStream.GetOrCreateStream().PushAsync(envelope.Message, onPullAction, onPullActionArgument, cancellationToken);
             return true;
         }
 
@@ -228,7 +222,7 @@ private static bool PushIfCompatibleType(
     private static bool IsFiltered(IReadOnlyCollection<IMessageFilter>? filters, object message) =>
         filters != null && filters.Count != 0 && !filters.All(filter => filter.MustProcess(message));
 
-    private IEnumerable<Task> PushToCompatibleStreams(TMessage message, CancellationToken cancellationToken)
+    private IEnumerable<Task> PushToCompatibleStreams(TMessage message, Action<object?>? onPullAction, object? onPullActionArgument, CancellationToken cancellationToken)
     {
         foreach (ILazyMessageStreamEnumerable? lazyStream in _lazyStreams)
         {
@@ -237,7 +231,7 @@ private IEnumerable<Task> PushToCompatibleStreams(TMessage message, Cancellation
             if (cancellationToken.IsCancellationRequested)
                 yield break;
 
-            if (PushIfCompatibleType(lazyStream, message, cancellationToken, out Task processingTask))
+            if (PushIfCompatibleType(lazyStream, message, onPullAction, onPullActionArgument, cancellationToken, out Task processingTask))
             {
                 yield return processingTask;
             }

src/Silverback.Integration/Messaging/Diagnostics/ActivityExtensions.cs

@@ -12,9 +12,7 @@ namespace Silverback.Messaging.Diagnostics;
 
 internal static class ActivityExtensions
 {
-    public static void AddBaggageRange(
-        this Activity activity,
-        IEnumerable<KeyValuePair<string, string>> baggageItems)
+    public static void AddBaggageRange(this Activity activity, IEnumerable<KeyValuePair<string, string>> baggageItems)
     {
         Check.NotNull(activity, nameof(activity));
         Check.NotNull(baggageItems, nameof(baggageItems));
@@ -25,6 +23,16 @@ public static void AddBaggageRange(
         }
     }
 
+    public static void ClearBaggage(this Activity activity)
+    {
+        Check.NotNull(activity, nameof(activity));
+
+        foreach (KeyValuePair<string, string?> pair in activity.Baggage)
+        {
+            activity.SetBaggage(pair.Key, null);
+        }
+    }
+
     public static void SetMessageHeaders(this Activity activity, MessageHeaderCollection headers)
     {
         if (activity.Id == null)
@@ -36,16 +44,10 @@ public static void SetMessageHeaders(this Activity activity, MessageHeaderCollec
 
         string? traceState = activity.TraceStateString;
         if (traceState != null)
-        {
             headers.Add(DefaultMessageHeaders.TraceState, traceState);
-        }
 
         if (activity.Baggage.Any())
-        {
-            headers.Add(
-                DefaultMessageHeaders.TraceBaggage,
-                ActivityBaggageSerializer.Serialize(activity.Baggage));
-        }
+            headers.Add(DefaultMessageHeaders.TraceBaggage, ActivityBaggageSerializer.Serialize(activity.Baggage));
     }
 
     public static void SetTraceIdAndState(this Activity activity, string? traceId, string? traceState)
@@ -61,7 +63,7 @@ public static void SetTraceIdAndState(this Activity activity, string? traceId, s
         }
     }
 
-    public static void AddEndpointName(this Activity activity, string endpointName) =>
+    public static void SetEndpointName(this Activity activity, string endpointName) =>
         activity.SetTag(ActivityTagNames.MessageDestination, endpointName);
 
     public static Activity? StartWithTraceId(
@@ -72,15 +74,10 @@ public static void AddEndpointName(this Activity activity, string endpointName)
         string? traceState)
     {
         if (!activitySource.HasListeners())
-        {
             return null;
-        }
-
-        if (traceId != null && ActivityContext.TryParse(traceId, traceState, out ActivityContext context))
-        {
-            return activitySource.StartActivity(name, activityKind, context);
-        }
 
-        return activitySource.StartActivity(name, activityKind);
+        return traceId != null && ActivityContext.TryParse(traceId, traceState, out ActivityContext context)
+            ? activitySource.StartActivity(name, activityKind, context)
+            : activitySource.StartActivity(name, activityKind);
     }
 }