IActivitySourceFactory.cs

using System.Runtime.InteropServices;
using System.Collections.Concurrent;
using System.Diagnostics;

namespace Utilities.Common.Telemetry;
using Microsoft.Extensions.DependencyInjection;
using Microsoft.Extensions.DependencyInjection.Extensions;

/// <summary>
/// Provides extension methods for IServiceCollection to add an activity source factory.
/// </summary>
public static class ActivitySourceFactoryExtensions
{
    /// <summary>
    /// Adds the default implementation of <see cref="IActivitySourceFactory"/> to the service collection.
    /// </summary>
    /// <param name="serviceCollection">The <see cref="IServiceCollection"/> to add the <see cref="IActivitySourceFactory"/> implementation to.</param>
    /// <returns>The modified <see cref="IServiceCollection"/>.</returns>
    public static IServiceCollection AddActivitySourceFactory(this IServiceCollection serviceCollection)
    {
        serviceCollection.TryAddSingleton<IActivitySourceFactory, DefaultActivitySourceFactory>();
        return serviceCollection;
    }
}

/// <summary>
/// Represents a factory for creating instances of the IActivitySource interface.
/// </summary>
public interface IActivitySourceFactory : IDisposable
{
    /// <summary>Creates a new instance of <see cref="ActivitySource"/> based on the provided options.</summary>
    /// <param name="name">The name of the ActivitySource <see cref="ActivitySource"/> object.</param>
    /// <param name="version">The version of the component publishing the tracing info..</param>
    /// <returns>The newly created <see cref="ActivitySource"/> instance.</returns>
    ActivitySource Create(string name, string? version);
}

/// <summary>
/// The default implementation of the <see cref="IActivitySourceFactory"/> interface.
/// </summary>
internal sealed class DefaultActivitySourceFactory : IActivitySourceFactory
{
    /// <summary>
    /// The dictionary that caches activity sources.
    /// </summary>
    private readonly ConcurrentDictionary<string, List<ActivitySource>> _cachedActivitySources = new();
    /// <summary>
    /// Indicates whether the object has been disposed.
    /// </summary>
    private volatile bool _disposed;

    /// <inheritdoc />
    /// <exception cref="ArgumentNullException">Thrown when the options parameter is null.</exception>
    /// <exception cref="ObjectDisposedException">Thrown when the DefaultActivitySourceFactory has been disposed.</exception>
    public ActivitySource Create(string name, string? version)
    {
        ArgumentException.ThrowIfNullOrWhiteSpace(name);
        ObjectDisposedException.ThrowIf(_disposed, typeof(DefaultActivitySourceFactory));

        if (_cachedActivitySources.TryGetValue(name, out var activitySourceCollection))
        {
            var activitySources = CollectionsMarshal.AsSpan(activitySourceCollection);
            var activitySourceReference = MemoryMarshal.GetReference(activitySources);
            for ( var index = 0; index < activitySources.Length; index++ )
            {
                ref var activitySourceRef = ref Unsafe.Add(ref activitySourceReference, index);
                if (activitySourceRef.Version == version)
                    return activitySourceRef;
            }
        }
        else
        {
            activitySourceCollection = [];
            _cachedActivitySources.TryAdd(name, activitySourceCollection);
        }

        var activitySource = new ActivitySource(name, version);

        activitySourceCollection.Add(activitySource);
        
        return activitySource;
    }
    
    public void Dispose()
    {
        if (_disposed) return;

        _disposed = true;

        if (_cachedActivitySources.IsEmpty) return;
        
        var cacheActivitySource = CollectionsMarshal.AsSpan(_cachedActivitySources.Values.ToList());
        ref var activitySourcesReference = ref MemoryMarshal.GetReference(cacheActivitySource);

        for ( var cacheActivitySourceIndex = 0; cacheActivitySourceIndex < cacheActivitySource.Length; cacheActivitySourceIndex++ )
        {
            var activitySources = CollectionsMarshal.AsSpan(Unsafe.Add(ref activitySourcesReference, cacheActivitySourceIndex));
            var activitySourceReference = MemoryMarshal.GetReference(activitySources);
            for ( var activitySourceIndex = 0; activitySourceIndex < activitySources.Length; activitySourceIndex++ )
            {
                Unsafe.Add(ref activitySourceReference, activitySourceIndex).Dispose();
            }
        }

        _cachedActivitySources.Clear();
    }
}