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

Commit daac1d4

Browse files
porglezompktoso
authored andcommitted
Support async/throws functions in the @Traced macro
**Motivation:** We need to support `@Traced` on async and throws functions to make it compatible with everywhere people want to use them. **Modifications:** Based on the effects signature of the attached function we apply try/await as appropriate. But, if the function is async/throws but those effects aren't actually used in the function body, this causes a new warning because the closure isn't inferred to be async and/or throws. To avoid those warnings, we also apply matching effects specifiers to the withSpan closure. **Result:** The `@Traced` macro is now usable on all functions.
1 parent 8632d7a commit daac1d4

File tree

3 files changed

+203
-2
lines changed

3 files changed

+203
-2
lines changed

Sources/TracingMacros/Docs.docc/index.md

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,8 @@ Macro helpers for Tracing.
77
The TracingMacros module provides optional macros to make it easier to write traced code.
88

99
The ``Traced()`` macro lets you avoid the extra indentation that comes with
10-
adopting traced code. You can just attach `@Traced` to a function and get
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
1112
started.
1213

1314
## Topics

Sources/TracingMacrosImplementation/TracedMacro.swift

Lines changed: 28 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -29,9 +29,36 @@ public struct TracedMacro: BodyMacro {
2929
throw MacroExpansionErrorMessage("expected a function with a body")
3030
}
3131

32+
// Construct a withSpan call matching the invocation of the @Traced macro
33+
3234
let operationName = StringLiteralExprSyntax(content: function.name.text)
3335
let withSpanCall: ExprSyntax = "withSpan(\(operationName))"
34-
let withSpanExpr: ExprSyntax = "\(withSpanCall) { span in \(body.statements) }"
36+
37+
// We want to explicitly specify the closure effect specifiers in order
38+
// to avoid warnings about unused try/await expressions.
39+
// We might as well explicitly specify the closure return type to help type inference.
40+
41+
let asyncClause = function.signature.effectSpecifiers?.asyncSpecifier
42+
let returnClause = function.signature.returnClause
43+
var throwsClause = function.signature.effectSpecifiers?.throwsClause
44+
// You aren't allowed to apply "rethrows" as a closure effect
45+
// specifier, so we have to convert this to a "throws" effect
46+
if throwsClause?.throwsSpecifier.tokenKind == .keyword(.rethrows) {
47+
throwsClause?.throwsSpecifier = .keyword(.throws)
48+
}
49+
var withSpanExpr: ExprSyntax = """
50+
\(withSpanCall) { span \(asyncClause)\(throwsClause)\(returnClause)in \(body.statements) }
51+
"""
52+
53+
// Apply a try / await as necessary to adapt the withSpan expression
54+
55+
if function.signature.effectSpecifiers?.asyncSpecifier != nil {
56+
withSpanExpr = "await \(withSpanExpr)"
57+
}
58+
59+
if function.signature.effectSpecifiers?.throwsClause != nil {
60+
withSpanExpr = "try \(withSpanExpr)"
61+
}
3562

3663
return ["\(withSpanExpr)"]
3764
}

Tests/TracingMacrosTests/TracedTests.swift

Lines changed: 173 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -55,6 +55,146 @@ final class TracedMacroTests: XCTestCase {
5555
)
5656
}
5757

58+
func test_tracedMacro_sync_throws() {
59+
assertMacroExpansion(
60+
"""
61+
@Traced
62+
func syncThrowingExample(param: Int) throws {
63+
struct ExampleError: Error {
64+
}
65+
throw ExampleError()
66+
}
67+
""",
68+
expandedSource: """
69+
func syncThrowingExample(param: Int) throws {
70+
try withSpan("syncThrowingExample") { span throws in
71+
struct ExampleError: Error {
72+
}
73+
throw ExampleError()
74+
}
75+
}
76+
""",
77+
macros: ["Traced": TracedMacro.self]
78+
)
79+
}
80+
81+
func test_tracedMacro_sync_rethrows() {
82+
assertMacroExpansion(
83+
"""
84+
@Traced
85+
func syncRethrowingExample(body: () throws -> Int) rethrows -> Int {
86+
print("Starting")
87+
let result = try body()
88+
print("Ending")
89+
return result
90+
}
91+
""",
92+
expandedSource: """
93+
func syncRethrowingExample(body: () throws -> Int) rethrows -> Int {
94+
try withSpan("syncRethrowingExample") { span throws -> Int in
95+
print("Starting")
96+
let result = try body()
97+
print("Ending")
98+
return result
99+
}
100+
}
101+
""",
102+
macros: ["Traced": TracedMacro.self]
103+
)
104+
}
105+
106+
func test_tracedMacro_async_nothrow() {
107+
assertMacroExpansion(
108+
"""
109+
@Traced
110+
func asyncNonthrowingExample(param: Int) async {
111+
print(param)
112+
}
113+
""",
114+
expandedSource: """
115+
func asyncNonthrowingExample(param: Int) async {
116+
await withSpan("asyncNonthrowingExample") { span async in
117+
print(param)
118+
}
119+
}
120+
""",
121+
macros: ["Traced": TracedMacro.self]
122+
)
123+
}
124+
125+
func test_tracedMacro_async_throws() {
126+
assertMacroExpansion(
127+
"""
128+
@Traced
129+
func asyncThrowingExample(param: Int) async throws {
130+
try await Task.sleep(for: .seconds(1))
131+
print("Hello")
132+
}
133+
""",
134+
expandedSource: """
135+
func asyncThrowingExample(param: Int) async throws {
136+
try await withSpan("asyncThrowingExample") { span async throws in
137+
try await Task.sleep(for: .seconds(1))
138+
print("Hello")
139+
}
140+
}
141+
""",
142+
macros: ["Traced": TracedMacro.self]
143+
)
144+
}
145+
146+
func test_tracedMacro_async_rethrows() {
147+
assertMacroExpansion(
148+
"""
149+
@Traced
150+
func asyncRethrowingExample(body: () async throws -> Int) async rethrows -> Int {
151+
try? await Task.sleep(for: .seconds(1))
152+
let result = try await body()
153+
span.attributes["result"] = result
154+
return result
155+
}
156+
""",
157+
expandedSource: """
158+
func asyncRethrowingExample(body: () async throws -> Int) async rethrows -> Int {
159+
try await withSpan("asyncRethrowingExample") { span async throws -> Int in
160+
try? await Task.sleep(for: .seconds(1))
161+
let result = try await body()
162+
span.attributes["result"] = result
163+
return result
164+
}
165+
}
166+
""",
167+
macros: ["Traced": TracedMacro.self]
168+
)
169+
}
170+
171+
// Testing that this expands correctly, but not including this as a
172+
// compile-test because withSpan doesn't currently support typed throws.
173+
func test_tracedMacro_async_typed_throws() {
174+
assertMacroExpansion(
175+
"""
176+
@Traced
177+
func asyncTypedThrowingExample<Err>(body: () async throws(Err) -> Int) async throws(Err) -> Int {
178+
try? await Task.sleep(for: .seconds(1))
179+
let result = try await body()
180+
span.attributes["result"] = result
181+
return result
182+
}
183+
""",
184+
expandedSource: """
185+
func asyncTypedThrowingExample<Err>(body: () async throws(Err) -> Int) async throws(Err) -> Int {
186+
try await withSpan("asyncTypedThrowingExample") { span async throws(Err) -> Int in
187+
try? await Task.sleep(for: .seconds(1))
188+
let result = try await body()
189+
span.attributes["result"] = result
190+
return result
191+
}
192+
}
193+
""",
194+
macros: ["Traced": TracedMacro.self]
195+
)
196+
}
197+
58198
func test_tracedMacro_accessSpan() {
59199
assertMacroExpansion(
60200
"""
@@ -85,6 +225,39 @@ func syncNonthrowingExample(param: Int) {
85225
print(param)
86226
}
87227

228+
@Traced
229+
func syncThrowingExample(param: Int) throws {
230+
struct ExampleError: Error {}
231+
throw ExampleError()
232+
}
233+
234+
@Traced
235+
func syncRethrowingExample(body: () throws -> Int) rethrows -> Int {
236+
print("Starting")
237+
let result = try body()
238+
print("Ending")
239+
return result
240+
}
241+
242+
@Traced
243+
func asyncNonthrowingExample(param: Int) async {
244+
print(param)
245+
}
246+
247+
@Traced
248+
func asyncThrowingExample(param: Int) async throws {
249+
try await Task.sleep(for: .seconds(1))
250+
print("Hello")
251+
}
252+
253+
@Traced
254+
func asyncRethrowingExample(body: () async throws -> Int) async rethrows -> Int {
255+
try? await Task.sleep(for: .seconds(1))
256+
let result = try await body()
257+
span.attributes["result"] = result
258+
return result
259+
}
260+
88261
@Traced
89262
func example(param: Int) {
90263
span.attributes["param"] = param

0 commit comments

Comments
 (0)