Skip to content

Commit

Permalink
Add support for parameters passed to the @Traced macro
Browse files Browse the repository at this point in the history 
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.
  • Loading branch information
@porglezomp
porglezomp committed Nov 19, 2024
1 parent 50446c1 commit 101df66
Show file tree
Hide file tree
Showing 2 changed files with 276 additions and 3 deletions.
77 changes: 74 additions & 3 deletions Sources/TracingMacrosImplementation/TracedMacro.swift
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand All @@ -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
Expand All @@ -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

Expand Down
202 changes: 202 additions & 0 deletions Tests/TracingMacrosTests/TracedTests.swift
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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

0 comments on commit 101df66

Please sign in to comment.