Skip to content

Commit cbf75a4

Browse files
authored
Merge pull request #950 from JonkaSusaki/tracing/subsegment-wrapper
feat(tracing): add Subsegment wrapper to prevent exposing Amazon.XRayRecorder.Core.Internal
2 parents 7f033da + 07ac2e4 commit cbf75a4

File tree

3 files changed

+260
-7
lines changed

3 files changed

+260
-7
lines changed
Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,17 @@
1+
using Amazon.XRay.Recorder.Core.Internal.Entities;
2+
3+
namespace AWS.Lambda.Powertools.Tracing.Internal;
4+
5+
/// <summary>
6+
/// Class TracingSubsegment.
7+
/// It's a wrapper for Subsegment from Amazon.XRay.Recorder.Core.Internal.
8+
/// </summary>
9+
/// <seealso cref="Subsegment" />
10+
public class TracingSubsegment : Subsegment
11+
{
12+
/// <summary>
13+
/// Wrapper constructor
14+
/// </summary>
15+
/// <param name="name"></param>
16+
public TracingSubsegment(string name) : base(name) { }
17+
}

libraries/src/AWS.Lambda.Powertools.Tracing/Tracing.cs

Lines changed: 8 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -130,7 +130,7 @@ public static void AddHttpInformation(string key, object value)
130130
/// <param name="name">The name of the subsegment.</param>
131131
/// <param name="subsegment">The AWS X-Ray subsegment for the wrapped consumer.</param>
132132
/// <exception cref="ArgumentNullException">Thrown when the name is not provided.</exception>
133-
public static void WithSubsegment(string name, Action<Subsegment> subsegment)
133+
public static void WithSubsegment(string name, Action<TracingSubsegment> subsegment)
134134
{
135135
WithSubsegment(null, name, subsegment);
136136
}
@@ -145,7 +145,7 @@ public static void WithSubsegment(string name, Action<Subsegment> subsegment)
145145
/// <param name="name">The name of the subsegment.</param>
146146
/// <param name="subsegment">The AWS X-Ray subsegment for the wrapped consumer.</param>
147147
/// <exception cref="System.ArgumentNullException">name</exception>
148-
public static void WithSubsegment(string nameSpace, string name, Action<Subsegment> subsegment)
148+
public static void WithSubsegment(string nameSpace, string name, Action<TracingSubsegment> subsegment)
149149
{
150150
if (string.IsNullOrWhiteSpace(name))
151151
throw new ArgumentNullException(nameof(name));
@@ -154,7 +154,8 @@ public static void WithSubsegment(string nameSpace, string name, Action<Subsegme
154154
XRayRecorder.Instance.SetNamespace(GetNamespaceOrDefault(nameSpace));
155155
try
156156
{
157-
subsegment?.Invoke((Subsegment) XRayRecorder.Instance.GetEntity());
157+
var entity = XRayRecorder.Instance.GetEntity() as TracingSubsegment;
158+
subsegment?.Invoke(entity);
158159
}
159160
finally
160161
{
@@ -174,7 +175,7 @@ public static void WithSubsegment(string nameSpace, string name, Action<Subsegme
174175
/// <param name="subsegment">The AWS X-Ray subsegment for the wrapped consumer.</param>
175176
/// <exception cref="ArgumentNullException">Thrown when the name is not provided.</exception>
176177
/// <exception cref="ArgumentNullException">Thrown when the entity is not provided.</exception>
177-
public static void WithSubsegment(string name, Entity entity, Action<Subsegment> subsegment)
178+
public static void WithSubsegment(string name, Entity entity, Action<TracingSubsegment> subsegment)
178179
{
179180
WithSubsegment(null, name, subsegment);
180181
}
@@ -191,15 +192,15 @@ public static void WithSubsegment(string name, Entity entity, Action<Subsegment>
191192
/// <param name="subsegment">The AWS X-Ray subsegment for the wrapped consumer.</param>
192193
/// <exception cref="System.ArgumentNullException">name</exception>
193194
/// <exception cref="System.ArgumentNullException">entity</exception>
194-
public static void WithSubsegment(string nameSpace, string name, Entity entity, Action<Subsegment> subsegment)
195+
public static void WithSubsegment(string nameSpace, string name, Entity entity, Action<TracingSubsegment> subsegment)
195196
{
196197
if (string.IsNullOrWhiteSpace(name))
197198
throw new ArgumentNullException(nameof(name));
198199

199200
if (entity is null)
200201
throw new ArgumentNullException(nameof(entity));
201202

202-
var childSubsegment = new Subsegment($"## {name}");
203+
var childSubsegment = new TracingSubsegment($"## {name}");
203204
entity.AddSubsegment(childSubsegment);
204205
childSubsegment.Sampled = entity.Sampled;
205206
childSubsegment.SetStartTimeToNow();
@@ -240,7 +241,7 @@ private static string GetNamespaceOrDefault(string nameSpace)
240241

241242
return PowertoolsConfigurations.Instance.Service;
242243
}
243-
244+
244245
/// <summary>
245246
/// Registers X-Ray for all instances of <see cref="Amazon.Runtime.AmazonServiceClient"/>.
246247
/// </summary>
Lines changed: 235 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,235 @@
1+
using AWS.Lambda.Powertools.Tracing.Internal;
2+
using Xunit;
3+
using Amazon.XRay.Recorder.Core.Internal.Entities;
4+
using System;
5+
6+
namespace AWS.Lambda.Powertools.Tracing.Tests;
7+
8+
[Collection("Sequential")]
9+
public class TracingSubsegmentTests
10+
{
11+
12+
[Fact]
13+
public void TracingSubsegment_Constructor_Should_Set_Name()
14+
{
15+
// Arrange
16+
var name = "test-segment";
17+
18+
// Act
19+
var subsegment = new TracingSubsegment(name);
20+
21+
// Assert
22+
Assert.Equal("test-segment", subsegment.Name);
23+
Assert.True(Entity.IsIdValid(subsegment.Id));
24+
Assert.Null(subsegment.TraceId);
25+
Assert.Null(subsegment.ParentId);
26+
Assert.False(subsegment.IsSubsegmentsAdded);
27+
}
28+
29+
[Fact]
30+
public void Test_Add_Ref_And_Release_With_TracingSubsegment()
31+
{
32+
// Arrange
33+
var parent = new Segment("parent", TraceId.NewId());
34+
var child = new TracingSubsegment("child");
35+
36+
// Act
37+
parent.AddSubsegment(child);
38+
39+
// Assert
40+
Assert.Equal(2, parent.Reference);
41+
Assert.Equal(1, child.Reference);
42+
43+
child.Release();
44+
Assert.Equal(1, parent.Reference);
45+
Assert.Equal(0, child.Reference);
46+
Assert.False(parent.IsEmittable());
47+
Assert.False(child.IsEmittable());
48+
49+
parent.Release();
50+
Assert.Equal(0, parent.Reference);
51+
Assert.True(parent.IsEmittable());
52+
Assert.True(child.IsEmittable());
53+
}
54+
55+
[Fact]
56+
public void IsEmittable_Returns_False_Without_Parent()
57+
{
58+
var subsegment = new TracingSubsegment("segment");
59+
Assert.False(subsegment.IsEmittable());
60+
}
61+
62+
[Fact]
63+
public void TracingSubsegment_Is_Assignable_BaseClass()
64+
{
65+
// Arrange
66+
var subsegment = new TracingSubsegment("segment");
67+
68+
// Act & Assert
69+
Assert.IsAssignableFrom<Subsegment>(subsegment);
70+
}
71+
72+
[Fact]
73+
public void Tracing_WithSubsegment_Invokes_Delegate_With_TracingSubsegment()
74+
{
75+
// Arrange
76+
bool delegateInvoked = false;
77+
78+
void TracingSubsegmentDelegate(TracingSubsegment subsegment)
79+
{
80+
delegateInvoked = true;
81+
}
82+
83+
// Act
84+
Tracing.WithSubsegment("namespace", "test", TracingSubsegmentDelegate);
85+
86+
// Assert
87+
Assert.True(delegateInvoked);
88+
}
89+
90+
[Fact]
91+
public void WithSubsegment_WithEntity_ThrowsArgumentNullException_WhenNameIsNull()
92+
{
93+
// Arrange
94+
var parent = new Segment("parent", TraceId.NewId());
95+
96+
// Act & Assert
97+
Assert.Throws<ArgumentNullException>(() =>
98+
Tracing.WithSubsegment(null, null, parent, _ => { }));
99+
}
100+
101+
[Fact]
102+
public void WithSubsegment_WithEntity_ThrowsArgumentNullException_WhenNameIsEmpty()
103+
{
104+
// Arrange
105+
var parent = new Segment("parent", TraceId.NewId());
106+
107+
// Act & Assert
108+
Assert.Throws<ArgumentNullException>(() =>
109+
Tracing.WithSubsegment(null, "", parent, _ => { }));
110+
}
111+
112+
[Fact]
113+
public void WithSubsegment_WithEntity_ThrowsArgumentNullException_WhenEntityIsNull()
114+
{
115+
// Act & Assert
116+
Assert.Throws<ArgumentNullException>(() =>
117+
Tracing.WithSubsegment(null, "test", null, _ => { }));
118+
}
119+
120+
[Fact]
121+
public void WithSubsegment_WithEntity_CreatesSubsegmentWithCorrectName()
122+
{
123+
// Arrange
124+
var parent = new Segment("parent", TraceId.NewId());
125+
TracingSubsegment capturedSubsegment = null;
126+
127+
// Act
128+
Tracing.WithSubsegment("test-namespace", "test-name", parent, subsegment =>
129+
{
130+
capturedSubsegment = subsegment;
131+
});
132+
133+
// Assert
134+
Assert.NotNull(capturedSubsegment);
135+
Assert.Equal("## test-name", capturedSubsegment.Name);
136+
Assert.Equal("test-namespace", capturedSubsegment.Namespace);
137+
}
138+
139+
[Fact]
140+
public void WithSubsegment_WithEntity_SetsSubsegmentProperties()
141+
{
142+
// Arrange
143+
var parent = new Segment("parent", TraceId.NewId());
144+
TracingSubsegment capturedSubsegment = null;
145+
146+
// Act
147+
Tracing.WithSubsegment("test-namespace", "test-name", parent, subsegment =>
148+
{
149+
capturedSubsegment = subsegment;
150+
});
151+
152+
// Assert
153+
Assert.NotNull(capturedSubsegment);
154+
Assert.Equal(parent.Sampled, capturedSubsegment.Sampled);
155+
Assert.False(capturedSubsegment.IsInProgress);
156+
Assert.True(capturedSubsegment.StartTime > 0);
157+
Assert.True(capturedSubsegment.EndTime > 0);
158+
}
159+
160+
[Fact]
161+
public void WithSubsegment_WithEntity_AddsSubsegmentToParent()
162+
{
163+
// Arrange
164+
var parent = new Segment("parent", TraceId.NewId());
165+
var initialSubsegmentCount = parent.Subsegments?.Count ?? 0;
166+
167+
// Act
168+
Tracing.WithSubsegment("test-namespace", "test-name", parent, _ => { });
169+
170+
// Assert
171+
Assert.True(parent.IsSubsegmentsAdded);
172+
Assert.Equal(initialSubsegmentCount + 1, parent.Subsegments.Count);
173+
}
174+
175+
[Fact]
176+
public void WithSubsegment_WithEntity_InvokesActionWithSubsegment()
177+
{
178+
// Arrange
179+
var parent = new Segment("parent", TraceId.NewId());
180+
bool actionInvoked = false;
181+
TracingSubsegment passedSubsegment = null;
182+
183+
// Act
184+
Tracing.WithSubsegment("test-namespace", "test-name", parent, subsegment =>
185+
{
186+
actionInvoked = true;
187+
passedSubsegment = subsegment;
188+
});
189+
190+
// Assert
191+
Assert.True(actionInvoked);
192+
Assert.NotNull(passedSubsegment);
193+
Assert.IsType<TracingSubsegment>(passedSubsegment);
194+
}
195+
196+
197+
[Fact]
198+
public void WithSubsegment_WithEntity_UsesDefaultNamespaceWhenNull()
199+
{
200+
// Arrange
201+
var parent = new Segment("parent", TraceId.NewId());
202+
TracingSubsegment capturedSubsegment = null;
203+
204+
// Act
205+
Tracing.WithSubsegment(null, "test-name", parent, subsegment =>
206+
{
207+
capturedSubsegment = subsegment;
208+
});
209+
210+
// Assert
211+
Assert.NotNull(capturedSubsegment);
212+
Assert.NotNull(capturedSubsegment.Namespace);
213+
}
214+
215+
[Fact]
216+
public void WithSubsegment_WithEntity_HandlesExceptionInAction()
217+
{
218+
// Arrange
219+
var parent = new Segment("parent", TraceId.NewId());
220+
var expectedException = new InvalidOperationException("Test exception");
221+
222+
// Act & Assert
223+
var actualException = Assert.Throws<InvalidOperationException>(() =>
224+
{
225+
Tracing.WithSubsegment("test-namespace", "test-name", parent, subsegment =>
226+
{
227+
throw expectedException;
228+
});
229+
});
230+
231+
Assert.Equal(expectedException, actualException);
232+
// Verify subsegment was still properly cleaned up
233+
Assert.True(parent.IsSubsegmentsAdded);
234+
}
235+
}

0 commit comments

Comments
 (0)