Skip to content
This repository was archived by the owner on Jul 11, 2025. It is now read-only.

Commit 148582e

Browse files
porglezompktoso
authored andcommitted
Add support for parameters passed to the @Traced macro
**Motivation:** The `@Traced` macro isn't sufficient for all cases of `withSpan` because it doesn't allow the same customization of the span. We need to expose those parameters to the macro. **Modifications:** First, allow overriding the regular parameters of the withSpan call: setting the operationName, the context, and the kind. If it's not specified, it's not passed to the function, which means the function remains the source of truth for default arguments. Also allow overriding the span name, which controls the variable binding in the closure body. This is primarily useful for avoiding shadowing an outer "span" variable. **Result:** Users of the macro can use the operationName, context, and kind parameters to customize the span like `withSpan`, and can use the span parameter to customize the variable introduced by the macro's expansion.
1 parent daac1d4 commit 148582e

File tree

4 files changed

+303
-10
lines changed

4 files changed

+303
-10
lines changed

Sources/TracingMacros/Docs.docc/index.md

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -6,12 +6,12 @@ Macro helpers for Tracing.
66

77
The TracingMacros module provides optional macros to make it easier to write traced code.
88

9-
The ``Traced()`` macro lets you avoid the extra indentation that comes with
10-
adopting traced code, and avoids having to keep the throws/try and async/await
11-
in-sync with the body. You can just attach `@Traced` to a function and get
12-
started.
9+
The ``Traced(_:context:ofKind:span:)`` macro lets you avoid the extra
10+
indentation that comes with adopting traced code, and avoids having to keep the
11+
throws/try and async/await in-sync with the body. You can just attach `@Traced`
12+
to a function and get started.
1313

1414
## Topics
1515

1616
### Tracing functions
17-
- ``Traced()``
17+
- ``Traced(_:context:ofKind:span:)``

Sources/TracingMacros/TracedMacro.swift

Lines changed: 18 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,23 @@ import Tracing
1818
/// Instrument a function to place the entire body inside a span.
1919
///
2020
/// This macro is equivalent to calling ``withSpan`` in the body, but saves an
21-
/// indentation level and duplication.
21+
/// indentation level and duplication. It introduces a `span` variable into the
22+
/// body of the function which can be used to add attributes to the span.
23+
///
24+
/// Parameters are passed directly to ``withSpan`` where applicable, and
25+
/// omitting the parameters from the macro omit them from the call, falling
26+
/// back to the default.
27+
///
28+
/// - Parameters:
29+
/// - operationName: The name of the operation being traced. Defaults to the name of the function.
30+
/// - context: The `ServiceContext` providing information on where to start the new ``Span``.
31+
/// - kind: The ``SpanKind`` of the new ``Span``.
32+
/// - spanName: The name of the span variable to introduce in the function. Pass `"_"` to omit it.
2233
@attached(body)
23-
public macro Traced() = #externalMacro(module: "TracingMacrosImplementation", type: "TracedMacro")
34+
public macro Traced(
35+
_ operationName: String? = nil,
36+
context: ServiceContext? = nil,
37+
ofKind kind: SpanKind? = nil,
38+
span spanName: String = "span"
39+
) = #externalMacro(module: "TracingMacrosImplementation", type: "TracedMacro")
2440
#endif

Sources/TracingMacrosImplementation/TracedMacro.swift

Lines changed: 78 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -30,9 +30,32 @@ public struct TracedMacro: BodyMacro {
3030
}
3131

3232
// Construct a withSpan call matching the invocation of the @Traced macro
33+
let (operationName, context, kind, spanName) = try extractArguments(from: node)
3334

34-
let operationName = StringLiteralExprSyntax(content: function.name.text)
35-
let withSpanCall: ExprSyntax = "withSpan(\(operationName))"
35+
var withSpanCall = FunctionCallExprSyntax("withSpan()" as ExprSyntax)!
36+
withSpanCall.arguments.append(
37+
LabeledExprSyntax(
38+
expression: operationName ?? ExprSyntax(StringLiteralExprSyntax(content: function.name.text))
39+
)
40+
)
41+
func appendComma() {
42+
withSpanCall.arguments[withSpanCall.arguments.index(before: withSpanCall.arguments.endIndex)]
43+
.trailingComma = .commaToken()
44+
}
45+
if let context {
46+
appendComma()
47+
withSpanCall.arguments.append(LabeledExprSyntax(label: "context", expression: context))
48+
}
49+
if let kind {
50+
appendComma()
51+
withSpanCall.arguments.append(LabeledExprSyntax(label: "ofKind", expression: kind))
52+
}
53+
54+
// Introduce a span identifier in scope
55+
var spanIdentifier: TokenSyntax = "span"
56+
if let spanName {
57+
spanIdentifier = .identifier(spanName)
58+
}
3659

3760
// We want to explicitly specify the closure effect specifiers in order
3861
// to avoid warnings about unused try/await expressions.
@@ -47,7 +70,7 @@ public struct TracedMacro: BodyMacro {
4770
throwsClause?.throwsSpecifier = .keyword(.throws)
4871
}
4972
var withSpanExpr: ExprSyntax = """
50-
\(withSpanCall) { span \(asyncClause)\(throwsClause)\(returnClause)in \(body.statements) }
73+
\(withSpanCall) { \(spanIdentifier) \(asyncClause)\(throwsClause)\(returnClause)in \(body.statements) }
5174
"""
5275

5376
// Apply a try / await as necessary to adapt the withSpan expression
@@ -62,6 +85,58 @@ public struct TracedMacro: BodyMacro {
6285

6386
return ["\(withSpanExpr)"]
6487
}
88+
89+
static func extractArguments(
90+
from node: AttributeSyntax
91+
) throws -> (
92+
operationName: ExprSyntax?,
93+
context: ExprSyntax?,
94+
kind: ExprSyntax?,
95+
spanName: String?
96+
) {
97+
// If there are no arguments, we don't have to do any of these bindings
98+
guard let arguments = node.arguments?.as(LabeledExprListSyntax.self) else {
99+
return (nil, nil, nil, nil)
100+
}
101+
102+
func getArgument(label: String) -> ExprSyntax? {
103+
arguments.first(where: { $0.label?.identifier?.name == label })?.expression
104+
}
105+
106+
// The operation name is the first argument if it's unlabeled
107+
var operationName: ExprSyntax?
108+
if let firstArgument = arguments.first, firstArgument.label == nil {
109+
operationName = firstArgument.expression
110+
}
111+
112+
let context = getArgument(label: "context")
113+
let kind = getArgument(label: "ofKind")
114+
var spanName: String?
115+
let spanNameExpr = getArgument(label: "span")
116+
if let spanNameExpr {
117+
guard let stringLiteral = spanNameExpr.as(StringLiteralExprSyntax.self),
118+
stringLiteral.segments.count == 1,
119+
let segment = stringLiteral.segments.first,
120+
let segmentText = segment.as(StringSegmentSyntax.self)
121+
else {
122+
throw MacroExpansionErrorMessage("span name must be a simple string literal")
123+
}
124+
let text = segmentText.content.text
125+
let isValidIdentifier = DeclReferenceExprSyntax("\(raw: text)" as ExprSyntax)?.hasError == false
126+
let isValidWildcard = text == "_"
127+
guard isValidIdentifier || isValidWildcard else {
128+
throw MacroExpansionErrorMessage("'\(text)' is not a valid parameter name")
129+
}
130+
spanName = text
131+
}
132+
return (
133+
operationName: operationName,
134+
context: context,
135+
kind: kind,
136+
spanName: spanName
137+
)
138+
}
139+
65140
}
66141
#endif
67142

Tests/TracingMacrosTests/TracedTests.swift

Lines changed: 202 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -213,6 +213,203 @@ final class TracedMacroTests: XCTestCase {
213213
macros: ["Traced": TracedMacro.self]
214214
)
215215
}
216+
217+
func test_tracedMacro_specifyOperationName() {
218+
assertMacroExpansion(
219+
"""
220+
@Traced("example but with a custom operationName")
221+
func example(param: Int) {
222+
span.attributes["param"] = param
223+
}
224+
""",
225+
expandedSource: """
226+
func example(param: Int) {
227+
withSpan("example but with a custom operationName") { span in
228+
span.attributes["param"] = param
229+
}
230+
}
231+
""",
232+
macros: ["Traced": TracedMacro.self]
233+
)
234+
235+
assertMacroExpansion(
236+
"""
237+
let globalName = "example"
238+
239+
@Traced(globalName)
240+
func example(param: Int) {
241+
span.attributes["param"] = param
242+
}
243+
""",
244+
expandedSource: """
245+
let globalName = "example"
246+
func example(param: Int) {
247+
withSpan(globalName) { span in
248+
span.attributes["param"] = param
249+
}
250+
}
251+
""",
252+
macros: ["Traced": TracedMacro.self]
253+
)
254+
}
255+
256+
func test_tracedMacro_specifyContext() {
257+
assertMacroExpansion(
258+
"""
259+
@Traced(context: .topLevel)
260+
func example() {
261+
print("Hello")
262+
}
263+
""",
264+
expandedSource: """
265+
func example() {
266+
withSpan("example", context: .topLevel) { span in
267+
print("Hello")
268+
}
269+
}
270+
""",
271+
macros: ["Traced": TracedMacro.self]
272+
)
273+
}
274+
275+
func test_tracedMacro_specifyKind() {
276+
assertMacroExpansion(
277+
"""
278+
@Traced(ofKind: .client)
279+
func example() {
280+
print("Hello")
281+
}
282+
""",
283+
expandedSource: """
284+
func example() {
285+
withSpan("example", ofKind: .client) { span in
286+
print("Hello")
287+
}
288+
}
289+
""",
290+
macros: ["Traced": TracedMacro.self]
291+
)
292+
}
293+
294+
func test_tracedMacro_specifySpanBindingName() {
295+
assertMacroExpansion(
296+
"""
297+
@Traced(span: "customSpan")
298+
func example(span: String) throws {
299+
customSpan.attributes["span"] = span
300+
}
301+
""",
302+
expandedSource: """
303+
func example(span: String) throws {
304+
try withSpan("example") { customSpan throws in
305+
customSpan.attributes["span"] = span
306+
}
307+
}
308+
""",
309+
macros: ["Traced": TracedMacro.self]
310+
)
311+
312+
assertMacroExpansion(
313+
"""
314+
@Traced(span: "_")
315+
func example(span: String) {
316+
print(span)
317+
}
318+
""",
319+
expandedSource: """
320+
func example(span: String) {
321+
withSpan("example") { _ in
322+
print(span)
323+
}
324+
}
325+
""",
326+
macros: ["Traced": TracedMacro.self]
327+
)
328+
}
329+
330+
func test_tracedMacro_specifySpanBindingName_invalid() {
331+
assertMacroExpansion(
332+
"""
333+
@Traced(span: 1)
334+
func example(span: String) throws {
335+
customSpan.attributes["span"] = span
336+
}
337+
""",
338+
expandedSource: """
339+
func example(span: String) throws {
340+
customSpan.attributes["span"] = span
341+
}
342+
""",
343+
diagnostics: [
344+
.init(message: "span name must be a simple string literal", line: 1, column: 1)
345+
],
346+
macros: ["Traced": TracedMacro.self]
347+
)
348+
349+
assertMacroExpansion(
350+
"""
351+
@Traced(span: "invalid name")
352+
func example(span: String) throws {
353+
customSpan.attributes["span"] = span
354+
}
355+
356+
@Traced(span: "123")
357+
func example2(span: String) throws {
358+
customSpan.attributes["span"] = span
359+
}
360+
""",
361+
expandedSource: """
362+
func example(span: String) throws {
363+
customSpan.attributes["span"] = span
364+
}
365+
func example2(span: String) throws {
366+
customSpan.attributes["span"] = span
367+
}
368+
""",
369+
diagnostics: [
370+
.init(message: "'invalid name' is not a valid parameter name", line: 1, column: 1),
371+
.init(message: "'123' is not a valid parameter name", line: 6, column: 1),
372+
],
373+
macros: ["Traced": TracedMacro.self]
374+
)
375+
376+
assertMacroExpansion(
377+
"""
378+
@Traced(span: "Hello \\(1)")
379+
func example(span: String) throws {
380+
customSpan.attributes["span"] = span
381+
}
382+
""",
383+
expandedSource: """
384+
func example(span: String) throws {
385+
customSpan.attributes["span"] = span
386+
}
387+
""",
388+
diagnostics: [
389+
.init(message: "span name must be a simple string literal", line: 1, column: 1)
390+
],
391+
macros: ["Traced": TracedMacro.self]
392+
)
393+
}
394+
395+
func test_tracedMacro_multipleMacroParameters() {
396+
assertMacroExpansion(
397+
"""
398+
@Traced("custom span name", context: .topLevel, ofKind: .client, span: "customSpan")
399+
func example(span: Int) {
400+
customSpan.attributes["span"] = span + 1
401+
}
402+
""",
403+
expandedSource: """
404+
func example(span: Int) {
405+
withSpan("custom span name", context: .topLevel, ofKind: .client) { customSpan in
406+
customSpan.attributes["span"] = span + 1
407+
}
408+
}
409+
""",
410+
macros: ["Traced": TracedMacro.self]
411+
)
412+
}
216413
#endif
217414
}
218415

@@ -263,4 +460,9 @@ func example(param: Int) {
263460
span.attributes["param"] = param
264461
}
265462

463+
@Traced("custom span name", context: .topLevel, ofKind: .client, span: "customSpan")
464+
func exampleWithParams(span: Int) {
465+
customSpan.attributes["span"] = span + 1
466+
}
467+
266468
#endif

0 commit comments

Comments
 (0)