diff --git a/src/libraries/System.Diagnostics.DiagnosticSource/ref/System.Diagnostics.DiagnosticSourceActivity.cs b/src/libraries/System.Diagnostics.DiagnosticSource/ref/System.Diagnostics.DiagnosticSourceActivity.cs index 9af996309bb10e..7b4aa10de2bf56 100644 --- a/src/libraries/System.Diagnostics.DiagnosticSource/ref/System.Diagnostics.DiagnosticSourceActivity.cs +++ b/src/libraries/System.Diagnostics.DiagnosticSource/ref/System.Diagnostics.DiagnosticSourceActivity.cs @@ -57,6 +57,7 @@ public string? Id public System.Diagnostics.Activity AddTag(string key, string? value) { throw null; } public System.Diagnostics.Activity AddTag(string key, object? value) { throw null; } public System.Diagnostics.Activity SetTag(string key, object? value) { throw null; } + public System.Diagnostics.Activity SetBaggage(string key, string? value) { throw null; } public string? GetBaggageItem(string key) { throw null; } public System.Diagnostics.Activity SetEndTime(System.DateTime endTimeUtc) { throw null; } public System.Diagnostics.Activity SetIdFormat(System.Diagnostics.ActivityIdFormat format) { throw null; } diff --git a/src/libraries/System.Diagnostics.DiagnosticSource/src/System/Diagnostics/Activity.cs b/src/libraries/System.Diagnostics.DiagnosticSource/src/System/Diagnostics/Activity.cs index a1c6830a77e4c5..e867f367963c3c 100644 --- a/src/libraries/System.Diagnostics.DiagnosticSource/src/System/Diagnostics/Activity.cs +++ b/src/libraries/System.Diagnostics.DiagnosticSource/src/System/Diagnostics/Activity.cs @@ -77,7 +77,7 @@ public partial class Activity : IDisposable private byte _w3CIdFlags; private TagsLinkedList? _tags; - private LinkedList>? _baggage; + private BaggageLinkedList? _baggage; private LinkedList? _links; private LinkedList? _events; private Dictionary? _customProperties; @@ -432,9 +432,33 @@ public Activity AddBaggage(string key, string? value) { KeyValuePair kvp = new KeyValuePair(key, value); - if (_baggage != null || Interlocked.CompareExchange(ref _baggage, new LinkedList>(kvp), null) != null) + if (_baggage != null || Interlocked.CompareExchange(ref _baggage, new BaggageLinkedList(kvp), null) != null) { - _baggage.AddFront(kvp); + _baggage.Add(kvp); + } + + return this; + } + + /// + /// Add or update the Activity baggage with the input key and value. + /// If the input value is null + /// - if the collection has any baggage with the same key, then this baggage will get removed from the collection. + /// - otherwise, nothing will happen and the collection will not change. + /// If the input value is not null + /// - if the collection has any baggage with the same key, then the value mapped to this key will get updated with the new input value. + /// - otherwise, the key and value will get added as a new baggage to the collection. + /// + /// The baggage key name + /// The baggage value mapped to the input key + /// 'this' for convenient chaining + public Activity SetBaggage(string key, string? value) + { + KeyValuePair kvp = new KeyValuePair(key, value); + + if (_baggage != null || Interlocked.CompareExchange(ref _baggage, new BaggageLinkedList(kvp, set: true), null) != null) + { + _baggage.Set(kvp); } return this; @@ -1307,12 +1331,94 @@ public void AddFront(T value) } } - // Note: Some customers use this GetEnumerator dynamically to avoid allocations. + // Note: Some consumers use this GetEnumerator dynamically to avoid allocations. public Enumerator GetEnumerator() => new Enumerator(_first); IEnumerator IEnumerable.GetEnumerator() => GetEnumerator(); IEnumerator IEnumerable.GetEnumerator() => GetEnumerator(); } + private class BaggageLinkedList : IEnumerable> + { + private LinkedListNode>? _first; + + public BaggageLinkedList(KeyValuePair firstValue, bool set = false) => _first = ((set && firstValue.Value == null) ? null : new LinkedListNode>(firstValue)); + + public LinkedListNode>? First => _first; + + public void Add(KeyValuePair value) + { + LinkedListNode> newNode = new LinkedListNode>(value); + + lock (this) + { + newNode.Next = _first; + _first = newNode; + } + } + + public void Set(KeyValuePair value) + { + if (value.Value == null) + { + Remove(value.Key); + return; + } + + lock (this) + { + LinkedListNode>? current = _first; + while (current != null) + { + if (current.Value.Key == value.Key) + { + current.Value = value; + return; + } + + current = current.Next; + } + + LinkedListNode> newNode = new LinkedListNode>(value); + newNode.Next = _first; + _first = newNode; + } + } + + public void Remove(string key) + { + lock (this) + { + if (_first == null) + { + return; + } + + if (_first.Value.Key == key) + { + _first = _first.Next; + return; + } + + LinkedListNode> previous = _first; + + while (previous.Next != null) + { + if (previous.Next.Value.Key == key) + { + previous.Next = previous.Next.Next; + return; + } + previous = previous.Next; + } + } + } + + // Note: Some consumers use this GetEnumerator dynamically to avoid allocations. + public Enumerator> GetEnumerator() => new Enumerator>(_first); + IEnumerator> IEnumerable>.GetEnumerator() => GetEnumerator(); + IEnumerator IEnumerable.GetEnumerator() => GetEnumerator(); + } + private class TagsLinkedList : IEnumerable> { private LinkedListNode>? _first; @@ -1380,13 +1486,13 @@ public void Add(KeyValuePair value) public void Remove(string key) { - if (_first == null) - { - return; - } lock (this) { + if (_first == null) + { + return; + } if (_first.Value.Key == key) { _first = _first.Next; @@ -1443,7 +1549,7 @@ public void Set(KeyValuePair value) } } - // Note: Some customers use this GetEnumerator dynamically to avoid allocations. + // Note: Some consumers use this GetEnumerator dynamically to avoid allocations. public Enumerator> GetEnumerator() => new Enumerator>(_first); IEnumerator> IEnumerable>.GetEnumerator() => GetEnumerator(); IEnumerator IEnumerable.GetEnumerator() => GetEnumerator(); @@ -1464,7 +1570,7 @@ public void Set(KeyValuePair value) } } - // Note: Some customers use this Enumerator dynamically to avoid allocations. + // Note: Some consumers use this Enumerator dynamically to avoid allocations. private struct Enumerator : IEnumerator { private LinkedListNode? _nextNode; diff --git a/src/libraries/System.Diagnostics.DiagnosticSource/tests/ActivityTests.cs b/src/libraries/System.Diagnostics.DiagnosticSource/tests/ActivityTests.cs index f753dc85a7f832..531c443de567c1 100644 --- a/src/libraries/System.Diagnostics.DiagnosticSource/tests/ActivityTests.cs +++ b/src/libraries/System.Diagnostics.DiagnosticSource/tests/ActivityTests.cs @@ -120,6 +120,55 @@ public void TestBaggageOrderAndDuplicateKeys() Assert.Equal("4", a.GetBaggageItem("1")); } + [Fact] + public void TestSetBaggage() + { + Activity a = new Activity("SetBaggage"); + Assert.Equal(0, a.Baggage.Count()); + + a.SetBaggage("1", "1"); + a.SetBaggage("2", "2"); + a.SetBaggage("3", "3"); + a.SetBaggage("4", "4"); + a.SetBaggage("5", "5"); + + Assert.Equal(5, a.Baggage.Count()); + Assert.Equal("1", a.GetBaggageItem("1")); + Assert.Equal("2", a.GetBaggageItem("2")); + Assert.Equal("3", a.GetBaggageItem("3")); + Assert.Equal("4", a.GetBaggageItem("4")); + Assert.Equal("5", a.GetBaggageItem("5")); + + // Check not added item + Assert.Null(a.GetBaggageItem("6")); + + // Adding none existing key with null value is no-op + a.SetBaggage("6", null); + Assert.Equal(5, a.Baggage.Count()); + Assert.Null(a.GetBaggageItem("6")); + + // Check updated item + a.SetBaggage("5", "5.1"); + Assert.Equal(5, a.Baggage.Count()); + Assert.Equal("5.1", a.GetBaggageItem("5")); + + // Add() always add new entry even we have matching key in the list. + // Baggage always return last entered value + a.AddBaggage("5", "5.2"); + Assert.Equal(6, a.Baggage.Count()); + Assert.Equal("5.2", a.GetBaggageItem("5")); + + // Now Remove first duplicated item + a.SetBaggage("5", null); + Assert.Equal(5, a.Baggage.Count()); + Assert.Equal("5.1", a.GetBaggageItem("5")); + + // Now Remove second item + a.SetBaggage("5", null); + Assert.Equal(4, a.Baggage.Count()); + Assert.Null(a.GetBaggageItem("5")); + } + [ConditionalFact(typeof(RemoteExecutor), nameof(RemoteExecutor.IsSupported))] public void TestBaggageWithChainedActivities() {