Revert "Try alternative SpanCollection implementation"

This reverts commit 9b542fd6b89bfa0b31025d9587fd6f64df8aad03.

Author: andrewlock

tracer/src/Datadog.Trace/Agent/SpanCollection.cs

@@ -10,6 +10,9 @@
 using System.Collections.Generic;
 using System.Runtime.CompilerServices;
 using Datadog.Trace.SourceGenerators;
+#if NETFRAMEWORK
+using Datadog.Trace.VendoredMicrosoftCode.System.Runtime.CompilerServices.Unsafe;
+#endif
 
 namespace Datadog.Trace.Agent;
 
@@ -18,8 +21,7 @@ namespace Datadog.Trace.Agent;
 /// </summary>
 internal readonly struct SpanCollection : IEnumerable<Span>
 {
-    private readonly Span? _span;
-    private readonly Span[]? _spans;
+    private readonly object? _values;
     public readonly int Count;
 
     /// <summary>
@@ -28,7 +30,7 @@ namespace Datadog.Trace.Agent;
     /// <param name="value">The span to include in the collection.</param>
     public SpanCollection(Span value)
     {
-        _span = value;
+        _values = value;
         Count = 1;
     }
 
@@ -38,7 +40,7 @@ public SpanCollection(Span value)
     /// <param name="arrayBuilderCapacity">The value to initializer</param>
     public SpanCollection(int arrayBuilderCapacity)
     {
-        _spans = new Span[arrayBuilderCapacity];
+        _values = new Span[arrayBuilderCapacity];
         Count = 0;
     }
 
@@ -49,7 +51,7 @@ public SpanCollection(Span[] values, int count)
     {
         // We assume that the caller is "sensible" here, and doesn't set count > values.Length,
         // but that will get hit "safely" elsewhere if it happens
-        _spans = values;
+        _values = values;
         Count = count;
     }
 
@@ -67,18 +69,20 @@ public Span? RootSpan
         [MethodImpl(MethodImplOptions.AggressiveInlining)]
         get
         {
-            if (_span is not null)
+            // Take local copy of _values so type checks remain valid even if the SpanCollection is overwritten in memory
+            var value = _values;
+            if (value is Span span)
             {
-                return _span;
+                return span;
             }
 
-            if (_spans is null)
+            if (value is null)
             {
                 return null;
             }
 
             // Not Span, not null, can only be SpanArray
-            return _spans[0];
+            return Unsafe.As<Span[]>(value)[0];
         }
     }
 
@@ -92,16 +96,19 @@ public Span this[int index]
         [MethodImpl(MethodImplOptions.AggressiveInlining)]
         get
         {
-            if (_span is not null)
+            // Take local copy of _values so type checks remain valid even if the SpanCollection is overwritten in memory
+            object? value = _values;
+            if (index < Count)
             {
-                if (index == 0)
+                if (value is Span str)
                 {
-                    return _span;
+                    return str;
+                }
+                else if (value != null)
+                {
+                    // Not Span, not null, can only be Span[]
+                    return Unsafe.As<Span[]>(value)[index];
                 }
-            }
-            else if (_spans is not null && index < Count)
-            {
-                return _spans[index];
             }
 
             return OutOfBounds(); // throws
@@ -116,19 +123,23 @@ public Span this[int index]
     /// <returns>The concatenation of <paramref name="values"/> and <paramref name="value"/>.</returns>
     public static SpanCollection Append(in SpanCollection values, Span value)
     {
-        if (values._span is not null)
+        // Take local copy of _values so type checks remain valid even if the SpanCollection is overwritten in memory
+        var current = values._values;
+        if (current is null)
         {
-            // We use a default capacity of 4 spans
-            // 2 Spans would cover 25% not covered by single span case, 4 covers ~ 70%, 8 covers ~92%
-            return new SpanCollection([values._span, value, null!, null!], 2);
+            return new SpanCollection(value);
         }
 
-        if (values._spans is null)
+        if (current is Span span)
         {
-            return new SpanCollection(value);
+            // We use a default capacity of 4 spans
+            // 2 Spans would cover 25% not covered by single span case, 4 covers ~ 70%, 8 covers ~92%
+            return new SpanCollection([span, value, null!, null!], 2);
         }
 
-        var array = GrowIfNeeded(values._spans, values.Count);
+        // Not Span, not null, can only be Span[], so add the span
+        var array = Unsafe.As<Span[]>(current);
+        array = GrowIfNeeded(array, values.Count);
         array[values.Count] = value;
         return new SpanCollection(array, values.Count + 1);
     }
@@ -150,17 +161,20 @@ private static Span OutOfBounds()
     /// </remarks>
     public ArraySegment<Span> ToArray()
     {
-        if (_spans is not null)
+        // Take local copy of _values so type checks remain valid even if the SpanCollection is overwritten in memory
+        object? value = _values;
+        if (value is Span[] values)
         {
-            return new ArraySegment<Span>(_spans, 0, Count);
+            return new ArraySegment<Span>(values, 0, Count);
         }
-        else if (_span is not null)
+        else if (value != null)
         {
-            return new ArraySegment<Span>([_span], 0, 1);
+            // value not array, can only be Span
+            return new ArraySegment<Span>([Unsafe.As<Span>(value)], 0, 1);
         }
         else
         {
-            return new ArraySegment<Span>([], 0, 0);
+            return new ArraySegment<Span>(Array.Empty<Span>(), 0, 0);
         }
     }
 
@@ -183,7 +197,7 @@ private static Span[] GrowIfNeeded(Span[] array, int currentCount)
     /// <returns>An enumerator that can be used to iterate through the <see cref="SpanCollection" />.</returns>
     public Enumerator GetEnumerator()
     {
-        return new Enumerator(_span, _spans, Count);
+        return new Enumerator(_values, Count);
     }
 
     /// <inheritdoc cref="GetEnumerator()" />
@@ -208,9 +222,9 @@ public struct Enumerator : IEnumerator<Span>
         private int _index;
         private Span? _current;
 
-        internal Enumerator(Span? span, Span[]? spans, int count)
+        internal Enumerator(object? value, int count)
         {
-            if (span is not null)
+            if (value is Span span)
             {
                 _values = null;
                 _current = span;
@@ -219,13 +233,22 @@ internal Enumerator(Span? span, Span[]? spans, int count)
             else
             {
                 _current = null;
-                _values = spans;
+                _values = Unsafe.As<Span[]>(value);
                 _count = count;
             }
 
             _index = 0;
         }
 
+        /// <summary>
+        /// Initializes a new instance of the <see cref="Enumerator"/> struct.
+        /// </summary>
+        /// <param name="values">The <see cref="SpanCollection"/> to enumerate.</param>
+        public Enumerator(ref SpanCollection values)
+            : this(values._values, values.Count)
+        {
+        }
+
         /// <summary>
         /// Gets the element at the current position of the enumerator.
         /// </summary>