lmxopcua/src/Drivers/ZB.MOM.WW.OtOpcUa.Driver.Galaxy/Runtime/TracedGalaxySubscriber.cs

using System.Runtime.CompilerServices;
using ZB.MOM.WW.MxGateway.Contracts.Proto;

namespace ZB.MOM.WW.OtOpcUa.Driver.Galaxy.Runtime;

/// <summary>
///     PR 6.1 — Decorator that emits one <see cref="System.Diagnostics.Activity"/> span
///     per gw subscription RPC. Wraps the production <see cref="GatewayGalaxySubscriber"/>;
///     tests substitute a fake at the same seam without taking the tracing overhead.
/// </summary>
internal sealed class TracedGalaxySubscriber(IGalaxySubscriber inner, string clientName) : IGalaxySubscriber
{
    /// <summary>Subscribes to multiple Galaxy tags in bulk with tracing.</summary>
    /// <param name="fullReferences">The full tag references to subscribe to.</param>
    /// <param name="bufferedUpdateIntervalMs">The buffered update interval in milliseconds.</param>
    /// <param name="cancellationToken">The cancellation token.</param>
    public async Task<IReadOnlyList<SubscribeResult>> SubscribeBulkAsync(
        IReadOnlyList<string> fullReferences, int bufferedUpdateIntervalMs, CancellationToken cancellationToken)
    {
        using var activity = GalaxyTelemetry.ActivitySource.StartActivity("galaxy.subscribe_bulk");
        activity?.SetTag("galaxy.client", clientName);
        activity?.SetTag("galaxy.tag_count", fullReferences.Count);
        activity?.SetTag("galaxy.buffered_interval_ms", bufferedUpdateIntervalMs);
        try
        {
            var results = await inner.SubscribeBulkAsync(fullReferences, bufferedUpdateIntervalMs, cancellationToken)
                .ConfigureAwait(false);
            activity?.SetTag("galaxy.success_count", results.Count(r => r.WasSuccessful));
            return results;
        }
        catch (Exception ex)
        {
            activity.RecordError(ex);
            throw;
        }
    }

    /// <summary>Unsubscribes from multiple Galaxy tags in bulk with tracing.</summary>
    /// <param name="itemHandles">The item handles to unsubscribe from.</param>
    /// <param name="cancellationToken">The cancellation token.</param>
    public async Task UnsubscribeBulkAsync(IReadOnlyList<int> itemHandles, CancellationToken cancellationToken)
    {
        using var activity = GalaxyTelemetry.ActivitySource.StartActivity("galaxy.unsubscribe_bulk");
        activity?.SetTag("galaxy.client", clientName);
        activity?.SetTag("galaxy.tag_count", itemHandles.Count);
        try
        {
            await inner.UnsubscribeBulkAsync(itemHandles, cancellationToken).ConfigureAwait(false);
        }
        catch (Exception ex)
        {
            activity.RecordError(ex);
            throw;
        }
    }

    /// <summary>
    ///     Streaming RPC — one parent span covers the entire stream lifetime. Per-event
    ///     spans would dominate the trace volume at 50k tags / 1Hz; ops gets per-event
    ///     visibility through <see cref="EventPump"/>'s metrics in PR 6.2 instead.
    /// </summary>
    /// <param name="cancellationToken">The cancellation token.</param>
    public async IAsyncEnumerable<MxEvent> StreamEventsAsync(
        [EnumeratorCancellation] CancellationToken cancellationToken)
    {
        using var activity = GalaxyTelemetry.ActivitySource.StartActivity("galaxy.stream_events");
        activity?.SetTag("galaxy.client", clientName);

        IAsyncEnumerator<MxEvent>? enumerator = null;
        try
        {
            enumerator = inner.StreamEventsAsync(cancellationToken).GetAsyncEnumerator(cancellationToken);
            var eventCount = 0L;
            while (true)
            {
                bool moveNext;
                try
                {
                    moveNext = await enumerator.MoveNextAsync().ConfigureAwait(false);
                }
                catch (Exception ex)
                {
                    activity.RecordError(ex);
                    activity?.SetTag("galaxy.event_count", eventCount);
                    throw;
                }

                if (!moveNext) break;
                eventCount++;
                yield return enumerator.Current;
            }
            activity?.SetTag("galaxy.event_count", eventCount);
        }
        finally
        {
            if (enumerator is not null) await enumerator.DisposeAsync().ConfigureAwait(false);
        }
    }
}