Revive TracingMacros and @Traced

Package.swift

@@ -1,14 +1,17 @@
 // swift-tools-version:5.9
+import CompilerPluginSupport
 import PackageDescription
 
 let package = Package(
     name: "swift-distributed-tracing",
     products: [
         .library(name: "Instrumentation", targets: ["Instrumentation"]),
         .library(name: "Tracing", targets: ["Tracing"]),
+        .library(name: "TracingMacros", targets: ["TracingMacros"]),
     ],
     dependencies: [
-        .package(url: "https://github.com/apple/swift-service-context.git", from: "1.1.0")
+        .package(url: "https://github.com/apple/swift-service-context.git", from: "1.1.0"),
+        .package(url: "https://github.com/swiftlang/swift-syntax.git", from: "600.0.0-latest"),
     ],
     targets: [
         // ==== --------------------------------------------------------------------------------------------------------
@@ -43,6 +46,33 @@ let package = Package(
                 .target(name: "Tracing")
             ]
         ),
+
+        // ==== --------------------------------------------------------------------------------------------------------
+        // MARK: TracingMacros
+
+        .target(
+            name: "TracingMacros",
+            dependencies: [
+                .target(name: "Tracing"),
+                .target(name: "TracingMacrosImplementation"),
+            ]
+        ),
+        .macro(
+            name: "TracingMacrosImplementation",
+            dependencies: [
+                .product(name: "SwiftSyntaxMacros", package: "swift-syntax"),
+                .product(name: "SwiftCompilerPlugin", package: "swift-syntax"),
+            ]
+        ),
+        .testTarget(
+            name: "TracingMacrosTests",
+            dependencies: [
+                .target(name: "Tracing"),
+                .target(name: "TracingMacros"),
+                .target(name: "TracingMacrosImplementation"),
+                .product(name: "SwiftSyntaxMacrosTestSupport", package: "swift-syntax"),
+            ]
+        ),
     ]
 )
 

Sources/TracingMacros/Docs.docc/TracedOperationName.md

@@ -0,0 +1,38 @@
+# ``TracingMacros/TracedOperationName``
+
+### Examples
+
+The default behavior is to use the base name of the function, but you can
+explicitly specify this as well. This creates a span named `"preheatOven"`:
+```swift
+@Traced(.baseName)
+func preheatOven(temperature: Int)
+```
+
+You can request the full name of the function as the span name, this
+creates a span named `"preheatOven(temperature:)"`:
+```swift
+@Traced(.fullName)
+func preheatOven(temperature: Int)
+```
+
+And it is also initializable with a string literal for fully custom names,
+this creates a span explicitly named `"preheat oven"`:
+```swift
+@Traced("preheat oven")
+func preheatOven(temperature: Int)
+```
+And if you need to load an existing string value as a name, you can use
+`.string(someString)` to adapt it.
+
+
+## Topics
+
+### Create Operation Names
+- ``baseName``
+- ``fullName``
+- ``string(_:)``
+- ``init(stringLiteral:)``
+
+### Convert an Operation Name to a String
+- ``operationName(baseName:fullName:)``

Sources/TracingMacros/Docs.docc/index.md

@@ -0,0 +1,19 @@
+# ``TracingMacros``
+
+Macro helpers for Tracing.
+
+## Overview
+
+The TracingMacros module provides optional macros to make it easier to write traced code.
+
+The ``Traced(_:context:ofKind:span:)`` macro lets you avoid the extra indentation that comes with
+adopting traced code, and avoids having to keep the throws/try and async/await
+in-sync with the body. You can just attach `@Traced` to a function and get
+started.
+
+## Topics
+
+### Tracing functions
+- ``Traced(_:context:ofKind:span:)``
+- ``TracedOperationName``
+

Sources/TracingMacros/TracedMacro.swift

@@ -0,0 +1,97 @@
+//===----------------------------------------------------------------------===//
+//
+// This source file is part of the Swift Distributed Tracing open source project
+//
+// Copyright (c) 2020-2024 Apple Inc. and the Swift Distributed Tracing project authors
+// Licensed under Apache License v2.0
+//
+// See LICENSE.txt for license information
+// See CONTRIBUTORS.txt for the list of Swift Distributed Tracing project authors
+//
+// SPDX-License-Identifier: Apache-2.0
+//
+//===----------------------------------------------------------------------===//
+@_exported import ServiceContextModule
+import Tracing
+
+/// A span name for a traced operation, either derived from the function name or explicitly specified.
+///
+/// When using the ``Traced(_:context:ofKind:span:)`` macro, you can use this to customize the span name.
+public struct TracedOperationName: ExpressibleByStringLiteral {
+    @usableFromInline
+    let value: Name
+
+    @usableFromInline
+    enum Name {
+        case baseName
+        case fullName
+        case string(String)
+    }
+
+    internal init(value: Name) {
+        self.value = value
+    }
+
+    /// Use a literal string as an operation name.
+    public init(stringLiteral: String) {
+        value = .string(stringLiteral)
+    }
+
+    /// Use the base name of the attached function.
+    ///
+    /// For `func preheatOven(temperature: Int)` this is `"preheatOven"`.
+    public static let baseName = TracedOperationName(value: .baseName)
+
+    /// Use the full name of the attached function.
+    ///
+    /// For `func preheatOven(temperature: Int)` this is `"preheatOven(temperature:)"`.
+    /// This is provided by the `#function` macro.
+    public static let fullName = TracedOperationName(value: .fullName)
+
+    /// Use an explicitly specified operation name.
+    public static func string(_ text: String) -> Self {
+        .init(value: .string(text))
+    }
+
+    /// Helper logic to support the `Traced` macro turning this operation name into a string.
+    /// Provided as an inference guide.
+    ///
+    /// - Parameters:
+    ///   - baseName: The value to use for the ``baseName`` case. Must be
+    ///     specified explicitly because there's no equivalent of `#function`.
+    ///   - fullName: The value to use for the ``fullName`` case.
+    @inlinable
+    @_documentation(visibility: internal)
+    public static func _getOperationName(_ name: Self, baseName: String, fullName: String = #function) -> String {
+        switch name.value {
+        case .baseName: baseName
+        case .fullName: fullName
+        case let .string(text): text
+        }
+    }
+}
+
+#if compiler(>=6.0)
+/// Instrument a function to place the entire body inside a span.
+///
+/// This macro is equivalent to calling ``withSpan`` in the body, but saves an
+/// indentation level and duplication. It introduces a `span` variable into the
+/// body of the function which can be used to add attributes to the span.
+///
+/// Parameters are passed directly to ``withSpan`` where applicable,
+/// and omitting the parameters from the macro omit them from the call, falling
+/// back to the default.
+///
+/// - Parameters:
+///   - operationName: The name of the operation being traced.
+///   - context: The `ServiceContext` providing information on where to start the new ``Span``.
+///   - kind: The ``SpanKind`` of the new ``Span``.
+///   - spanName: The name of the span variable to introduce in the function. Pass `"_"` to omit it.
+@attached(body)
+public macro Traced(
+    _ operationName: TracedOperationName = .baseName,
+    context: ServiceContext? = nil,
+    ofKind kind: SpanKind? = nil,
+    span spanName: String = "span"
+) = #externalMacro(module: "TracingMacrosImplementation", type: "TracedMacro")
+#endif

Sources/TracingMacrosImplementation/TracedMacro.swift

@@ -0,0 +1,159 @@
+//===----------------------------------------------------------------------===//
+//
+// This source file is part of the Swift Distributed Tracing open source project
+//
+// Copyright (c) 2020-2024 Apple Inc. and the Swift Distributed Tracing project authors
+// Licensed under Apache License v2.0
+//
+// See LICENSE.txt for license information
+// See CONTRIBUTORS.txt for the list of Swift Distributed Tracing project authors
+//
+// SPDX-License-Identifier: Apache-2.0
+//
+//===----------------------------------------------------------------------===//
+import SwiftCompilerPlugin
+import SwiftSyntax
+import SwiftSyntaxBuilder
+import SwiftSyntaxMacros
+
+#if compiler(>=6.0)
+public struct TracedMacro: BodyMacro {
+    public static func expansion(
+        of node: AttributeSyntax,
+        providingBodyFor declaration: some DeclSyntaxProtocol & WithOptionalCodeBlockSyntax,
+        in context: some MacroExpansionContext
+    ) throws -> [CodeBlockItemSyntax] {
+        guard let function = declaration.as(FunctionDeclSyntax.self),
+            let body = function.body
+        else {
+            throw MacroExpansionErrorMessage("expected a function with a body")
+        }
+
+        // Construct a withSpan call matching the invocation of the @Traced macro
+        let (operationName, context, kind, spanName) = try extractArguments(from: node)
+        let baseNameExpr = ExprSyntax(StringLiteralExprSyntax(content: function.name.text))
+
+        var withSpanCall = FunctionCallExprSyntax("withSpan()" as ExprSyntax)!
+        let operationNameExpr: ExprSyntax
+        if let operationName {
+            if operationName.is(StringLiteralExprSyntax.self) {
+                operationNameExpr = operationName
+            } else {
+                operationNameExpr = "TracedOperationName._getOperationName(\(operationName), baseName: \(baseNameExpr))"
+            }
+        } else {
+            operationNameExpr = baseNameExpr
+        }
+        withSpanCall.arguments.append(LabeledExprSyntax(expression: operationNameExpr))
+        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.
+        // 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) { \(spanIdentifier) \(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)"]
+    }
+
+    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
+
+@main
+struct TracingMacroPlugin: CompilerPlugin {
+    #if compiler(>=6.0)
+    let providingMacros: [Macro.Type] = [
+        TracedMacro.self
+    ]
+    #else
+    let providingMacros: [Macro.Type] = []
+    #endif
+}