From 101df6623449c93316deee089defbe343feadf35 Mon Sep 17 00:00:00 2001 From: Cassie Jones Date: Mon, 18 Nov 2024 22:28:20 -0800 Subject: [PATCH] Add support for parameters passed to the @Traced macro 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. --- .../TracedMacro.swift | 77 ++++++- Tests/TracingMacrosTests/TracedTests.swift | 202 ++++++++++++++++++ 2 files changed, 276 insertions(+), 3 deletions(-) diff --git a/Sources/TracingMacrosImplementation/TracedMacro.swift b/Sources/TracingMacrosImplementation/TracedMacro.swift index fc46a87..b9970fa 100644 --- a/Sources/TracingMacrosImplementation/TracedMacro.swift +++ b/Sources/TracingMacrosImplementation/TracedMacro.swift @@ -30,9 +30,28 @@ public struct TracedMacro: BodyMacro { } // Construct a withSpan call matching the invocation of the @Traced macro + let (operationName, context, kind, spanName) = try extractArguments(from: node) - let operationName = StringLiteralExprSyntax(content: function.name.text) - let withSpanCall: ExprSyntax = "withSpan(\(operationName))" + var withSpanCall = FunctionCallExprSyntax("withSpan()" as ExprSyntax)! + withSpanCall.arguments.append(LabeledExprSyntax( + expression: operationName ?? ExprSyntax(StringLiteralExprSyntax(content: function.name.text)))) + func appendComma() { + withSpanCall.arguments[withSpanCall.arguments.index(before: withSpanCall.arguments.endIndex)].trailingComma = .commaToken() + } + if let context { + appendComma() + withSpanCall.arguments.append(LabeledExprSyntax(label: "context", expression: context)) + } + if let kind { + appendComma() + withSpanCall.arguments.append(LabeledExprSyntax(label: "ofKind", expression: kind)) + } + + // Introduce a span identifier in scope + var spanIdentifier: TokenSyntax = "span" + if let spanName { + spanIdentifier = .identifier(spanName) + } // We want to explicitly specify the closure effect specifiers in order // to avoid warnings about unused try/await expressions. @@ -47,7 +66,7 @@ public struct TracedMacro: BodyMacro { throwsClause?.throwsSpecifier = .keyword(.throws) } var withSpanExpr: ExprSyntax = """ - \(withSpanCall) { span \(asyncClause)\(throwsClause)\(returnClause)in \(body.statements) } + \(withSpanCall) { \(spanIdentifier) \(asyncClause)\(throwsClause)\(returnClause)in \(body.statements) } """ // Apply a try / await as necessary to adapt the withSpan expression @@ -62,6 +81,58 @@ public struct TracedMacro: BodyMacro { return ["\(withSpanExpr)"] } + + static func extractArguments( + from node: AttributeSyntax + ) throws -> ( + operationName: ExprSyntax?, + context: ExprSyntax?, + kind: ExprSyntax?, + spanName: String? + ) { + // If there are no arguments, we don't have to do any of these bindings + guard let arguments = node.arguments?.as(LabeledExprListSyntax.self) else { + return (nil, nil, nil, nil) + } + + func getArgument(label: String) -> ExprSyntax? { + arguments.first(where: { $0.label?.identifier?.name == label })?.expression + } + + // The operation name is the first argument if it's unlabeled + var operationName: ExprSyntax? + if let firstArgument = arguments.first, firstArgument.label == nil { + operationName = firstArgument.expression + } + + let context = getArgument(label: "context") + let kind = getArgument(label: "ofKind") + var spanName: String? + let spanNameExpr = getArgument(label: "span") + if let spanNameExpr { + guard let stringLiteral = spanNameExpr.as(StringLiteralExprSyntax.self), + stringLiteral.segments.count == 1, + let segment = stringLiteral.segments.first, + let segmentText = segment.as(StringSegmentSyntax.self) + else { + throw MacroExpansionErrorMessage("span name must be a simple string literal") + } + let text = segmentText.content.text + let isValidIdentifier = DeclReferenceExprSyntax("\(raw: text)" as ExprSyntax)?.hasError == false + let isValidWildcard = text == "_" + guard isValidIdentifier || isValidWildcard else { + throw MacroExpansionErrorMessage("'\(text)' is not a valid parameter name") + } + spanName = text + } + return ( + operationName: operationName, + context: context, + kind: kind, + spanName: spanName, + ) + } + } #endif diff --git a/Tests/TracingMacrosTests/TracedTests.swift b/Tests/TracingMacrosTests/TracedTests.swift index 1af1c23..3152d5f 100644 --- a/Tests/TracingMacrosTests/TracedTests.swift +++ b/Tests/TracingMacrosTests/TracedTests.swift @@ -214,6 +214,203 @@ final class TracedMacroTests: XCTestCase { macros: ["Traced": TracedMacro.self] ) } + + func test_tracedMacro_specifyOperationName() { + assertMacroExpansion( + """ + @Traced("example but with a custom operationName") + func example(param: Int) { + span.attributes["param"] = param + } + """, + expandedSource: """ + func example(param: Int) { + withSpan("example but with a custom operationName") { span in + span.attributes["param"] = param + } + } + """, + macros: ["Traced": TracedMacro.self] + ) + + assertMacroExpansion( + """ + let globalName = "example" + + @Traced(globalName) + func example(param: Int) { + span.attributes["param"] = param + } + """, + expandedSource: """ + let globalName = "example" + func example(param: Int) { + withSpan(globalName) { span in + span.attributes["param"] = param + } + } + """, + macros: ["Traced": TracedMacro.self] + ) + } + + func test_tracedMacro_specifyContext() { + assertMacroExpansion( + """ + @Traced(context: .topLevel) + func example() { + print("Hello") + } + """, + expandedSource: """ + func example() { + withSpan("example", context: .topLevel) { span in + print("Hello") + } + } + """, + macros: ["Traced": TracedMacro.self] + ) + } + + func test_tracedMacro_specifyKind() { + assertMacroExpansion( + """ + @Traced(ofKind: .client) + func example() { + print("Hello") + } + """, + expandedSource: """ + func example() { + withSpan("example", ofKind: .client) { span in + print("Hello") + } + } + """, + macros: ["Traced": TracedMacro.self] + ) + } + + func test_tracedMacro_specifySpanBindingName() { + assertMacroExpansion( + """ + @Traced(span: "customSpan") + func example(span: String) throws { + customSpan.attributes["span"] = span + } + """, + expandedSource: """ + func example(span: String) throws { + try withSpan("example") { customSpan throws in + customSpan.attributes["span"] = span + } + } + """, + macros: ["Traced": TracedMacro.self] + ) + + assertMacroExpansion( + """ + @Traced(span: "_") + func example(span: String) { + print(span) + } + """, + expandedSource: """ + func example(span: String) { + withSpan("example") { _ in + print(span) + } + } + """, + macros: ["Traced": TracedMacro.self] + ) + } + + func test_tracedMacro_specifySpanBindingName_invalid() { + assertMacroExpansion( + """ + @Traced(span: 1) + func example(span: String) throws { + customSpan.attributes["span"] = span + } + """, + expandedSource: """ + func example(span: String) throws { + customSpan.attributes["span"] = span + } + """, + diagnostics: [ + .init(message: "span name must be a simple string literal", line: 1, column: 1), + ], + macros: ["Traced": TracedMacro.self] + ) + + assertMacroExpansion( + """ + @Traced(span: "invalid name") + func example(span: String) throws { + customSpan.attributes["span"] = span + } + + @Traced(span: "123") + func example2(span: String) throws { + customSpan.attributes["span"] = span + } + """, + expandedSource: """ + func example(span: String) throws { + customSpan.attributes["span"] = span + } + func example2(span: String) throws { + customSpan.attributes["span"] = span + } + """, + diagnostics: [ + .init(message: "'invalid name' is not a valid parameter name", line: 1, column: 1), + .init(message: "'123' is not a valid parameter name", line: 6, column: 1), + ], + macros: ["Traced": TracedMacro.self] + ) + + assertMacroExpansion( + """ + @Traced(span: "Hello \\(1)") + func example(span: String) throws { + customSpan.attributes["span"] = span + } + """, + expandedSource: """ + func example(span: String) throws { + customSpan.attributes["span"] = span + } + """, + diagnostics: [ + .init(message: "span name must be a simple string literal", line: 1, column: 1), + ], + macros: ["Traced": TracedMacro.self] + ) + } + + func test_tracedMacro_multipleMacroParameters() { + assertMacroExpansion( + """ + @Traced("custom span name", context: .topLevel, ofKind: .client, span: "customSpan") + func example(span: Int) { + customSpan.attributes["span"] = span + 1 + } + """, + expandedSource: """ + func example(span: Int) { + withSpan("custom span name", context: .topLevel, ofKind: .client) { customSpan in + customSpan.attributes["span"] = span + 1 + } + } + """, + macros: ["Traced": TracedMacro.self] + ) + } } // MARK: Compile tests @@ -261,4 +458,9 @@ func example(param: Int) { span.attributes["param"] = param } +@Traced("custom span name", context: .topLevel, ofKind: .client, span: "customSpan") +func exampleWithParams(span: Int) { + customSpan.attributes["span"] = span + 1 +} + #endif