Support async/throws functions in the @Traced macro

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.

Author: porglezomp

Sources/TracingMacrosImplementation/TracedMacro.swift

@@ -29,9 +29,36 @@ public struct TracedMacro: BodyMacro {
             throw MacroExpansionErrorMessage("expected a function with a body")
         }
 
+        // Construct a withSpan call matching the invocation of the @Traced macro
+
         let operationName = StringLiteralExprSyntax(content: function.name.text)
         let withSpanCall: ExprSyntax = "withSpan(\(operationName))"
-        let withSpanExpr: ExprSyntax = "\(withSpanCall) { span in \(body.statements) }"
+
+        // We want to explicitly specify the closure effect specifiers in order
+        // to avoid warnings about unused try/await expressions.
+        // We might as well explicitly specify the closure return type to help type inference.
+
+        let asyncClause = function.signature.effectSpecifiers?.asyncSpecifier
+        let returnClause = function.signature.returnClause
+        var throwsClause = function.signature.effectSpecifiers?.throwsClause
+        // You aren't allowed to apply "rethrows" as a closure effect
+        // specifier, so we have to convert this to a "throws" effect
+        if throwsClause?.throwsSpecifier.tokenKind == .keyword(.rethrows) {
+            throwsClause?.throwsSpecifier = .keyword(.throws)
+        }
+        var withSpanExpr: ExprSyntax = """
+        \(withSpanCall) { span \(asyncClause)\(throwsClause)\(returnClause)in \(body.statements) }
+        """
+
+        // Apply a try / await as necessary to adapt the withSpan expression
+
+        if function.signature.effectSpecifiers?.asyncSpecifier != nil {
+            withSpanExpr = "await \(withSpanExpr)"
+        }
+
+        if function.signature.effectSpecifiers?.throwsClause != nil {
+            withSpanExpr = "try \(withSpanExpr)"
+        }
 
         return ["\(withSpanExpr)"]
     }

Tests/TracingMacrosTests/TracedTests.swift

@@ -56,6 +56,146 @@ final class TracedMacroTests: XCTestCase {
         )
     }
 
+    func test_tracedMacro_sync_throws() {
+        assertMacroExpansion(
+            """
+            @Traced
+            func syncThrowingExample(param: Int) throws {
+                struct ExampleError: Error {
+                }
+                throw ExampleError()
+            }
+            """,
+            expandedSource: """
+            func syncThrowingExample(param: Int) throws {
+                try withSpan("syncThrowingExample") { span throws in
+                    struct ExampleError: Error {
+                    }
+                    throw ExampleError()
+                }
+            }
+            """,
+            macros: ["Traced": TracedMacro.self]
+        )
+    }
+
+    func test_tracedMacro_sync_rethrows() {
+        assertMacroExpansion(
+            """
+            @Traced
+            func syncRethrowingExample(body: () throws -> Int) rethrows -> Int {
+                print("Starting")
+                let result = try body()
+                print("Ending")
+                return result
+            }
+            """,
+            expandedSource: """
+            func syncRethrowingExample(body: () throws -> Int) rethrows -> Int {
+                try withSpan("syncRethrowingExample") { span throws -> Int in
+                    print("Starting")
+                    let result = try body()
+                    print("Ending")
+                    return result
+                }
+            }
+            """,
+            macros: ["Traced": TracedMacro.self]
+        )
+    }
+
+    func test_tracedMacro_async_nothrow() {
+        assertMacroExpansion(
+            """
+            @Traced
+            func asyncNonthrowingExample(param: Int) async {
+                print(param)
+            }
+            """,
+            expandedSource: """
+            func asyncNonthrowingExample(param: Int) async {
+                await withSpan("asyncNonthrowingExample") { span async in
+                    print(param)
+                }
+            }
+            """,
+            macros: ["Traced": TracedMacro.self]
+        )
+    }
+
+    func test_tracedMacro_async_throws() {
+        assertMacroExpansion(
+            """
+            @Traced
+            func asyncThrowingExample(param: Int) async throws {
+                try await Task.sleep(for: .seconds(1))
+                print("Hello")
+            }
+            """,
+            expandedSource: """
+            func asyncThrowingExample(param: Int) async throws {
+                try await withSpan("asyncThrowingExample") { span async throws in
+                    try await Task.sleep(for: .seconds(1))
+                    print("Hello")
+                }
+            }
+            """,
+            macros: ["Traced": TracedMacro.self]
+        )
+    }
+
+    func test_tracedMacro_async_rethrows() {
+        assertMacroExpansion(
+            """
+            @Traced
+            func asyncRethrowingExample(body: () async throws -> Int) async rethrows -> Int {
+                try? await Task.sleep(for: .seconds(1))
+                let result = try await body()
+                span.attributes["result"] = result
+                return result
+            }
+            """,
+            expandedSource: """
+            func asyncRethrowingExample(body: () async throws -> Int) async rethrows -> Int {
+                try await withSpan("asyncRethrowingExample") { span async throws -> Int in
+                    try? await Task.sleep(for: .seconds(1))
+                    let result = try await body()
+                    span.attributes["result"] = result
+                    return result
+                }
+            }
+            """,
+            macros: ["Traced": TracedMacro.self]
+        )
+    }
+
+    // Testing that this expands correctly, but not including this as a
+    // compile-test because withSpan doesn't currently support typed throws.
+    func test_tracedMacro_async_typed_throws() {
+        assertMacroExpansion(
+            """
+            @Traced
+            func asyncTypedThrowingExample<Err>(body: () async throws(Err) -> Int) async throws(Err) -> Int {
+                try? await Task.sleep(for: .seconds(1))
+                let result = try await body()
+                span.attributes["result"] = result
+                return result
+            }
+            """,
+            expandedSource: """
+            func asyncTypedThrowingExample<Err>(body: () async throws(Err) -> Int) async throws(Err) -> Int {
+                try await withSpan("asyncTypedThrowingExample") { span async throws(Err) -> Int in
+                    try? await Task.sleep(for: .seconds(1))
+                    let result = try await body()
+                    span.attributes["result"] = result
+                    return result
+                }
+            }
+            """,
+            macros: ["Traced": TracedMacro.self]
+        )
+    }
+
     func test_tracedMacro_accessSpan() {
         assertMacroExpansion(
             """
@@ -83,6 +223,39 @@ func syncNonthrowingExample(param: Int) {
     print(param)
 }
 
+@Traced
+func syncThrowingExample(param: Int) throws {
+    struct ExampleError: Error {}
+    throw ExampleError()
+}
+
+@Traced
+func syncRethrowingExample(body: () throws -> Int) rethrows -> Int {
+    print("Starting")
+    let result = try body()
+    print("Ending")
+    return result
+}
+
+@Traced
+func asyncNonthrowingExample(param: Int) async {
+    print(param)
+}
+
+@Traced
+func asyncThrowingExample(param: Int) async throws {
+    try await Task.sleep(for: .seconds(1))
+    print("Hello")
+}
+
+@Traced
+func asyncRethrowingExample(body: () async throws -> Int) async rethrows -> Int {
+    try? await Task.sleep(for: .seconds(1))
+    let result = try await body()
+    span.attributes["result"] = result
+    return result
+}
+
 @Traced
 func example(param: Int) {
     span.attributes["param"] = param