Added wrapper for Marten repository instead of the internal logic for tracing

Core.EventStoreDB/Repository/Config.cs

@@ -1,4 +1,5 @@
 using Core.Aggregates;
+using Core.OpenTelemetry;
 using Core.OptimisticConcurrency;
 using Microsoft.Extensions.DependencyInjection;
 
@@ -8,26 +9,33 @@ public static class Config
 {
     public static IServiceCollection AddEventStoreDBRepository<T>(
         this IServiceCollection services,
-        bool withAppendScope = true
+        bool withAppendScope = true,
+        bool withTelemetry = true
     ) where T : class, IAggregate
     {
-        services.AddScoped<EventStoreDBRepository<T>, EventStoreDBRepository<T>>();
+        services.AddScoped<IEventStoreDBRepository<T>, EventStoreDBRepository<T>>();
 
-        if (!withAppendScope)
+        if (withAppendScope)
         {
-            services.AddScoped<IEventStoreDBRepository<T>, EventStoreDBRepository<T>>();
-        }
-        else
-        {
-            services.AddScoped<IEventStoreDBRepository<T>, EventStoreDBRepositoryWithETagDecorator<T>>(
-                sp => new EventStoreDBRepositoryWithETagDecorator<T>(
-                    sp.GetRequiredService<EventStoreDBRepository<T>>(),
+            services.Decorate<IEventStoreDBRepository<T>>(
+                (inner, sp) => new EventStoreDBRepositoryWithETagDecorator<T>(
+                    inner,
                     sp.GetRequiredService<IExpectedResourceVersionProvider>(),
                     sp.GetRequiredService<INextResourceVersionProvider>()
                 )
             );
         }
 
+        if (withTelemetry)
+        {
+            services.Decorate<IEventStoreDBRepository<T>>(
+                (inner, sp) => new EventStoreDBRepositoryWithTelemetryDecorator<T>(
+                    inner,
+                    sp.GetRequiredService<IActivityScope>()
+                )
+            );
+        }
+
         return services;
     }
 }

Core.EventStoreDB/Repository/EventStoreDBRepository.cs

@@ -36,63 +36,42 @@ IActivityScope activityScope
             cancellationToken
         );
 
-    public Task<ulong> Add(T aggregate, CancellationToken token = default) =>
-        activityScope.Run($"{typeof(EventStoreDBRepository<T>).Name}/{nameof(Add)}",
-            async (activity, ct) =>
-            {
-                var result = await eventStore.AppendToStreamAsync(
-                    StreamNameMapper.ToStreamId<T>(aggregate.Id),
-                    StreamState.NoStream,
-                    GetEventsToStore(aggregate, TelemetryPropagator.GetPropagationContext(activity)),
-                    cancellationToken: ct
-                ).ConfigureAwait(false);
-                return result.NextExpectedStreamRevision.ToUInt64();
-            },
-            token
-        );
+    public async Task<ulong> Add(T aggregate, CancellationToken ct = default)
+    {
+        var result = await eventStore.AppendToStreamAsync(
+            StreamNameMapper.ToStreamId<T>(aggregate.Id),
+            StreamState.NoStream,
+            GetEventsToStore(aggregate),
+            cancellationToken: ct
+        ).ConfigureAwait(false);
 
-    public Task<ulong> Update(T aggregate, ulong? expectedRevision = null, CancellationToken token = default) =>
-        activityScope.Run($"{typeof(EventStoreDBRepository<T>).Name}/{nameof(Update)}",
-            async (activity, ct) =>
-            {
-                var eventsToAppend = GetEventsToStore(aggregate, TelemetryPropagator.GetPropagationContext(activity));
-                var nextVersion = expectedRevision ?? (ulong)(aggregate.Version - eventsToAppend.Count);
+        return result.NextExpectedStreamRevision.ToUInt64();
+    }
 
-                var result = await eventStore.AppendToStreamAsync(
-                    StreamNameMapper.ToStreamId<T>(aggregate.Id),
-                    nextVersion,
-                    eventsToAppend,
-                    cancellationToken: ct
-                ).ConfigureAwait(false);
-                return result.NextExpectedStreamRevision.ToUInt64();
-            },
-            token
-        );
+    public async Task<ulong> Update(T aggregate, ulong? expectedRevision = null, CancellationToken ct = default)
+    {
+        var eventsToAppend = GetEventsToStore(aggregate);
+        var nextVersion = expectedRevision ?? (ulong)(aggregate.Version - eventsToAppend.Count);
 
-    public Task<ulong> Delete(T aggregate, ulong? expectedRevision = null, CancellationToken token = default) =>
-        activityScope.Run($"{typeof(EventStoreDBRepository<T>).Name}/{nameof(Delete)}",
-            async (activity, ct) =>
-            {
-                var eventsToAppend = GetEventsToStore(aggregate, TelemetryPropagator.GetPropagationContext(activity));
-                var nextVersion = expectedRevision ?? (ulong)(aggregate.Version - eventsToAppend.Count);
+        var result = await eventStore.AppendToStreamAsync(
+            StreamNameMapper.ToStreamId<T>(aggregate.Id),
+            nextVersion,
+            eventsToAppend,
+            cancellationToken: ct
+        ).ConfigureAwait(false);
 
-                var result = await eventStore.AppendToStreamAsync(
-                    StreamNameMapper.ToStreamId<T>(aggregate.Id),
-                    nextVersion,
-                    eventsToAppend,
-                    cancellationToken: ct
-                ).ConfigureAwait(false);
-                return result.NextExpectedStreamRevision.ToUInt64();
-            },
-            token
-        );
+        return result.NextExpectedStreamRevision.ToUInt64();
+    }
+
+    public Task<ulong> Delete(T aggregate, ulong? expectedRevision = null, CancellationToken ct = default) =>
+        Update(aggregate, expectedRevision, ct);
 
-    private static List<EventData> GetEventsToStore(T aggregate, PropagationContext? propagationContext)
+    private static List<EventData> GetEventsToStore(T aggregate)
     {
         var events = aggregate.DequeueUncommittedEvents();
 
         return events
-            .Select(@event => @event.ToJsonEventData(propagationContext))
+            .Select(@event => @event.ToJsonEventData(TelemetryPropagator.GetPropagationContext()))
             .ToList();
     }
 }

Core.EventStoreDB/Repository/EventStoreDBRepositoryWithTelemetryDecorator.cs

@@ -0,0 +1,45 @@
+using Core.Aggregates;
+using Core.OpenTelemetry;
+using Microsoft.Extensions.Logging;
+
+namespace Core.EventStoreDB.Repository;
+
+public class EventStoreDBRepositoryWithTelemetryDecorator<T>: IEventStoreDBRepository<T>
+    where T : class, IAggregate
+{
+    private readonly IEventStoreDBRepository<T> inner;
+    private readonly IActivityScope activityScope;
+
+    public EventStoreDBRepositoryWithTelemetryDecorator(
+        IEventStoreDBRepository<T> inner,
+        IActivityScope activityScope
+    )
+    {
+        this.inner = inner;
+        this.activityScope = activityScope;
+    }
+
+    public Task<T?> Find(Guid id, CancellationToken cancellationToken) =>
+        inner.Find(id, cancellationToken);
+
+    public Task<ulong> Add(T aggregate, CancellationToken cancellationToken = default) =>
+        activityScope.Run($"EventStoreDBRepository/{nameof(Add)}",
+            (_, ct) => inner.Add(aggregate, ct),
+            new StartActivityOptions { Tags = { { TelemetryTags.Logic.Entity, typeof(T).Name } } },
+            cancellationToken
+        );
+
+    public Task<ulong> Update(T aggregate, ulong? expectedVersion = null, CancellationToken token = default) =>
+        activityScope.Run($"EventStoreDBRepository/{nameof(Update)}",
+            (_, ct) => inner.Update(aggregate, expectedVersion, ct),
+            new StartActivityOptions { Tags = { { TelemetryTags.Logic.Entity, typeof(T).Name } } },
+            token
+        );
+
+    public Task<ulong> Delete(T aggregate, ulong? expectedVersion = null, CancellationToken token = default) =>
+        activityScope.Run($"EventStoreDBRepository/{nameof(Delete)}",
+            (_, ct) => inner.Delete(aggregate, expectedVersion, ct),
+            new StartActivityOptions { Tags = { { TelemetryTags.Logic.Entity, typeof(T).Name } } },
+            token
+        );
+}

Core.Marten/Repository/Config.cs

@@ -1,31 +1,41 @@
 using Core.Aggregates;
+using Core.OpenTelemetry;
 using Core.OptimisticConcurrency;
+using Marten;
 using Microsoft.Extensions.DependencyInjection;
+using Microsoft.Extensions.Logging;
 
 namespace Core.Marten.Repository;
 
 public static class Config
 {
     public static IServiceCollection AddMartenRepository<T>(
         this IServiceCollection services,
-        bool withAppendScope = true
+        bool withAppendScope = true,
+        bool withTelemetry = true
     ) where T : class, IAggregate
     {
-        services.AddScoped<MartenRepository<T>, MartenRepository<T>>();
+        services.AddScoped<IMartenRepository<T>, MartenRepository<T>>();
 
-        if (!withAppendScope)
-        {
-            services.AddScoped<IMartenRepository<T>, MartenRepository<T>>();
-        }
-        else
-        {
-            services.AddScoped<IMartenRepository<T>, MartenRepositoryWithETagDecorator<T>>(
-                sp => new MartenRepositoryWithETagDecorator<T>(
-                    sp.GetRequiredService<MartenRepository<T>>(),
+        if (withAppendScope)
+            services.Decorate<IMartenRepository<T>>(
+                (inner, sp) => new MartenRepositoryWithETagDecorator<T>(
+                    inner,
                     sp.GetRequiredService<IExpectedResourceVersionProvider>(),
                     sp.GetRequiredService<INextResourceVersionProvider>()
                 )
             );
+
+        if (withTelemetry)
+        {
+            services.Decorate<IMartenRepository<T>>(
+                (inner, sp) => new MartenRepositoryWithTracingDecorator<T>(
+                    inner,
+                    sp.GetRequiredService<IDocumentSession>(),
+                    sp.GetRequiredService<IActivityScope>(),
+                    sp.GetRequiredService<ILogger<MartenRepositoryWithTracingDecorator<T>>>()
+                )
+            );
         }
 
         return services;

Core.Marten/Repository/MartenRepository.cs

@@ -17,111 +17,44 @@ public interface IMartenRepository<T> where T : class, IAggregate
 public class MartenRepository<T>: IMartenRepository<T> where T : class, IAggregate
 {
     private readonly IDocumentSession documentSession;
-    private readonly IActivityScope activityScope;
-    private readonly ILogger<MartenRepository<T>> logger;
 
-    public MartenRepository(
-        IDocumentSession documentSession,
-        IActivityScope activityScope,
-        ILogger<MartenRepository<T>> logger
-    )
-    {
+    public MartenRepository(IDocumentSession documentSession) =>
         this.documentSession = documentSession;
-        this.activityScope = activityScope;
-        this.logger = logger;
-    }
-
-    public Task<T?> Find(Guid id, CancellationToken cancellationToken) =>
-        documentSession.Events.AggregateStreamAsync<T>(id, token: cancellationToken);
-
-    public Task<long> Add(T aggregate, CancellationToken cancellationToken = default) =>
-        activityScope.Run($"{typeof(MartenRepository<T>).Name}/{nameof(Add)}",
-            async (activity, ct) =>
-            {
-                PropagateTelemetry(activity);
-
-                var events = aggregate.DequeueUncommittedEvents();
-
-                documentSession.Events.StartStream<Aggregate>(
-                    aggregate.Id,
-                    events
-                );
-
-                await documentSession.SaveChangesAsync(ct).ConfigureAwait(false);
-
-                return (long)events.Length;
-            },
-            new StartActivityOptions { Tags = { { TelemetryTags.Logic.Entity, typeof(T).Name } } },
-            cancellationToken
-        );
-
-    public Task<long> Update(T aggregate, long? expectedVersion = null, CancellationToken token = default) =>
-        activityScope.Run($"MartenRepository/{nameof(Update)}",
-            async (activity, ct) =>
-            {
-                PropagateTelemetry(activity);
 
-                var events = aggregate.DequeueUncommittedEvents();
+    public Task<T?> Find(Guid id, CancellationToken ct) =>
+        documentSession.Events.AggregateStreamAsync<T>(id, token: ct);
 
-                var nextVersion = (expectedVersion ?? aggregate.Version) + events.Length;
-
-                documentSession.Events.Append(
-                    aggregate.Id,
-                    nextVersion,
-                    events
-                );
-
-                await documentSession.SaveChangesAsync(ct).ConfigureAwait(false);
+    public async Task<long> Add(T aggregate, CancellationToken ct = default)
+    {
+        var events = aggregate.DequeueUncommittedEvents();
 
-                return nextVersion;
-            },
-            new StartActivityOptions { Tags = { { TelemetryTags.Logic.Entity, typeof(T).Name } } },
-            token
+        documentSession.Events.StartStream<Aggregate>(
+            aggregate.Id,
+            events
         );
 
-    public Task<long> Delete(T aggregate, long? expectedVersion = null, CancellationToken token = default) =>
-        activityScope.Run($"MartenRepository/{nameof(Delete)}",
-            async (activity, ct) =>
-            {
-                PropagateTelemetry(activity);
+        await documentSession.SaveChangesAsync(ct).ConfigureAwait(false);
 
-                var events = aggregate.DequeueUncommittedEvents();
-
-                var nextVersion = (expectedVersion ?? aggregate.Version) + events.Length;
+        return events.Length;
+    }
 
-                documentSession.Events.Append(
-                    aggregate.Id,
-                    nextVersion,
-                    events
-                );
+    public async Task<long> Update(T aggregate, long? expectedVersion = null, CancellationToken ct = default)
+    {
+        var events = aggregate.DequeueUncommittedEvents();
 
-                await documentSession.SaveChangesAsync(ct).ConfigureAwait(false);
+        var nextVersion = (expectedVersion ?? aggregate.Version) + events.Length;
 
-                return nextVersion;
-            },
-            new StartActivityOptions { Tags = { { TelemetryTags.Logic.Entity, typeof(T).Name } } },
-            token
+        documentSession.Events.Append(
+            aggregate.Id,
+            nextVersion,
+            events
         );
 
-    private void PropagateTelemetry(Activity? activity)
-    {
-        var propagationContext = activity.Propagate(documentSession, InjectTelemetryIntoDocumentSession);
-
-        if (!propagationContext.HasValue) return;
+        await documentSession.SaveChangesAsync(ct).ConfigureAwait(false);
 
-        documentSession.CorrelationId = propagationContext.Value.ActivityContext.TraceId.ToHexString();
-        documentSession.CausationId = propagationContext.Value.ActivityContext.SpanId.ToHexString();
+        return nextVersion;
     }
 
-    private void InjectTelemetryIntoDocumentSession(IDocumentSession session, string key, string value)
-    {
-        try
-        {
-            session.SetHeader(key, value);
-        }
-        catch (Exception ex)
-        {
-            logger.LogError(ex, "Failed to inject trace context");
-        }
-    }
+    public Task<long> Delete(T aggregate, long? expectedVersion = null, CancellationToken ct = default) =>
+        Update(aggregate, expectedVersion, ct);
 }