Improve tracing support

fabianfett · slashmo · commit baff60ce71d0 · 2025-09-13T16:07:32.000+02:00
Signed-off-by: Fabian Fett &lt;fabianfett@apple.com&gt;
diff --git a/Benchmarks/ValkeyBenchmarks/ValkeyConnectionBenchmark.swift b/Benchmarks/ValkeyBenchmarks/ValkeyConnectionBenchmark.swift
@@ -12,12 +12,18 @@ import Logging
 import NIOCore
 import NIOPosix
 import Synchronization
+#if DistributedTracingSupport
+import Tracing
+#endif
 import Valkey
 
 @available(valkeySwift 1.0, *)
 func connectionBenchmarks() {
     makeConnectionCreateAndDropBenchmark()
     makeConnectionGETBenchmark()
+    #if DistributedTracingSupport
+    makeConnectionGETNoOpTracerBenchmark()
+    #endif
     makeConnectionPipelineBenchmark()
 }
 
@@ -58,9 +64,45 @@ func makeConnectionGETBenchmark() -> Benchmark? {
     return Benchmark("Connection: GET benchmark", configuration: .init(metrics: defaultMetrics, scalingFactor: .kilo)) { benchmark in
         let port = serverMutex.withLock { $0 }!.localAddress!.port!
         let logger = Logger(label: "test")
+        #if DistributedTracingSupport
+        // explicitly set tracer to nil, if trait is enabled
+        var configuration = ValkeyConnectionConfiguration()
+        configuration.tracing.tracer = nil
+        #endif
         try await ValkeyConnection.withConnection(
             address: .hostname("127.0.0.1", port: port),
-            configuration: .init(),
+            configuration: configuration,
+            logger: logger
+        ) { connection in
+            benchmark.startMeasurement()
+            for _ in benchmark.scaledIterations {
+                let foo = try await connection.get("foo")
+                precondition(foo.map { String(buffer: $0) } == "Bar")
+            }
+            benchmark.stopMeasurement()
+        }
+    } setup: {
+        let server = try await makeLocalServer()
+        serverMutex.withLock { $0 = server }
+    } teardown: {
+        try await serverMutex.withLock { $0 }?.close().get()
+    }
+}
+
+#if DistributedTracingSupport
+@available(valkeySwift 1.0, *)
+@discardableResult
+func makeConnectionGETNoOpTracerBenchmark() -> Benchmark? {
+    let serverMutex = Mutex<(any Channel)?>(nil)
+
+    return Benchmark("Connection: GET benchmark – NoOpTracer", configuration: .init(metrics: defaultMetrics, scalingFactor: .kilo)) { benchmark in
+        let port = serverMutex.withLock { $0 }!.localAddress!.port!
+        let logger = Logger(label: "test")
+        var configuration = ValkeyConnectionConfiguration()
+        configuration.tracing.tracer = NoOpTracer()
+        try await ValkeyConnection.withConnection(
+            address: .hostname("127.0.0.1", port: port),
+            configuration: configuration,
             logger: logger
         ) { connection in
             benchmark.startMeasurement()
@@ -77,6 +119,7 @@ func makeConnectionGETBenchmark() -> Benchmark? {
         try await serverMutex.withLock { $0 }?.close().get()
     }
 }
+#endif
 
 @available(valkeySwift 1.0, *)
 @discardableResult
diff --git a/Sources/Valkey/Connection/ValkeyConnection.swift b/Sources/Valkey/Connection/ValkeyConnection.swift
@@ -33,6 +33,10 @@ public final actor ValkeyConnection: ValkeyClientProtocol, Sendable {
     public let id: ID
     /// Logger used by Server
     let logger: Logger
+    #if DistributedTracingSupport
+    @usableFromInline
+    let tracer: (any Tracer)?
+    #endif
     @usableFromInline
     let channel: any Channel
     @usableFromInline
@@ -57,6 +61,9 @@ public final actor ValkeyConnection: ValkeyClientProtocol, Sendable {
         self.configuration = configuration
         self.id = connectionID
         self.logger = logger
+        #if DistributedTracingSupport
+        self.tracer = configuration.tracing.tracer
+        #endif
         switch address?.value {
         case let .hostname(host, port):
             self.address = (host, port)
@@ -169,12 +176,11 @@ public final actor ValkeyConnection: ValkeyClientProtocol, Sendable {
     @inlinable
     func _execute<Command: ValkeyCommand>(command: Command) async throws -> RESPToken {
         #if DistributedTracingSupport
-        let span = startSpan(Command.name, ofKind: .client)
-        defer { span.end() }
+        let span = self.tracer?.startSpan(Command.name, ofKind: .client)
+        defer { span?.end() }
 
-        span.updateAttributes { attributes in
-            attributes["db.operation.name"] = Command.name
-            applyCommonAttributes(to: &attributes)
+        span?.updateAttributes { attributes in
+            self.applyCommonAttributes(to: &attributes, commandName: Command.name)
         }
         #endif
 
@@ -193,21 +199,21 @@ public final actor ValkeyConnection: ValkeyClientProtocol, Sendable {
             }
         } catch let error as ValkeyClientError {
             #if DistributedTracingSupport
-            span.recordError(error)
+            span?.recordError(error)
             if let message = error.message {
                 var prefixEndIndex = message.startIndex
                 while prefixEndIndex < message.endIndex, message[prefixEndIndex] != " " {
                     message.formIndex(after: &prefixEndIndex)
                 }
                 let prefix = message[message.startIndex..<prefixEndIndex]
-                span.attributes["db.response.status_code"] = "\(prefix)"
-                span.setStatus(SpanStatus(code: .error))
+                span?.attributes["db.response.status_code"] = "\(prefix)"
+                span?.setStatus(SpanStatus(code: .error))
             }
             #endif
             throw error
         } catch {
             #if DistributedTracingSupport
-            span.recordError(error)
+            span?.recordError(error)
             #endif
             throw error
         }
@@ -224,41 +230,7 @@ public final actor ValkeyConnection: ValkeyClientProtocol, Sendable {
     public func execute<each Command: ValkeyCommand>(
         _ commands: repeat each Command
     ) async -> sending (repeat Result<(each Command).Response, Error>) {
-        #if DistributedTracingSupport
-        let span = startSpan("MULTI", ofKind: .client)
-        defer { span.end() }
-
-        // We want to suffix the `db.operation.name` if all pipelined commands are of the same type.
-        var commandName: String?
-        var operationNameSuffix: String?
-        var commandCount = 0
-
-        for command in repeat each commands {
-            commandCount += 1
-            if commandName == nil {
-                commandName = Swift.type(of: command).name
-                operationNameSuffix = commandName
-            } else if commandName != Swift.type(of: command).name {
-                // We should only add a suffix if all commands in the transaction are the same.
-                operationNameSuffix = nil
-            }
-        }
-        let operationName = operationNameSuffix.map { "MULTI \($0)" } ?? "MULTI"
-
-        span.updateAttributes { attributes in
-            attributes["db.operation.name"] = operationName
-            attributes["db.operation.batch.size"] = commandCount > 1 ? commandCount : nil
-            applyCommonAttributes(to: &attributes)
-        }
-        #endif
-
         func convert<Response: RESPTokenDecodable>(_ result: Result<RESPToken, Error>, to: Response.Type) -> Result<Response, Error> {
-            #if DistributedTracingSupport
-            if case .failure(let error) = result {
-                span.recordError(error)
-            }
-            #endif
-
             return result.flatMap {
                 do {
                     return try .success(Response(fromRESP: $0))
@@ -295,12 +267,13 @@ public final actor ValkeyConnection: ValkeyClientProtocol, Sendable {
     }
 
     @usableFromInline
-    func applyCommonAttributes(to attributes: inout SpanAttributes) {
-        attributes["db.system.name"] = "valkey"
-        attributes["network.peer.address"] = channel.remoteAddress?.ipAddress
-        attributes["network.peer.port"] = channel.remoteAddress?.port
-        attributes["server.address"] = address?.hostOrSocketPath
-        attributes["server.port"] = address?.port == 6379 ? nil : address?.port
+    func applyCommonAttributes(to attributes: inout SpanAttributes, commandName: String) {
+        attributes[self.configuration.tracing.attributeNames.databaseOperationName] = commandName
+        attributes[self.configuration.tracing.attributeNames.databaseSystemName] = self.configuration.tracing.attributeValue.databaseSystem
+        attributes[self.configuration.tracing.attributeNames.networkPeerAddress] = channel.remoteAddress?.ipAddress
+        attributes[self.configuration.tracing.attributeNames.networkPeerPort] = channel.remoteAddress?.port
+        attributes[self.configuration.tracing.attributeNames.serverAddress] = address?.hostOrSocketPath
+        attributes[self.configuration.tracing.attributeNames.serverPort] = address?.port == 6379 ? nil : address?.port
     }
 
     @usableFromInline
diff --git a/Sources/Valkey/Connection/ValkeyConnectionConfiguration.swift b/Sources/Valkey/Connection/ValkeyConnectionConfiguration.swift
@@ -7,6 +7,9 @@
 //
 
 import NIOSSL
+#if DistributedTracingSupport
+import Tracing
+#endif
 
 /// A configuration object that defines how to connect to a Valkey server.
 ///
@@ -112,6 +115,10 @@ public struct ValkeyConnectionConfiguration: Sendable {
     /// Default value is `nil` (no client name is set).
     public var clientName: String?
 
+    #if DistributedTracingSupport
+    public var tracing: ValkeyTracingConfiguration = .init()
+    #endif
+
     /// Creates a new Valkey connection configuration.
     ///
     /// Use this initializer to create a configuration object that can be used to establish
@@ -137,3 +144,26 @@ public struct ValkeyConnectionConfiguration: Sendable {
         self.clientName = clientName
     }
 }
+
+#if DistributedTracingSupport
+public struct ValkeyTracingConfiguration: Sendable {
+
+    public var tracer: (any Tracer)? = InstrumentationSystem.tracer
+
+    public var attributeNames: AttributeNames = .init()
+    public var attributeValue: AttributeValues = .init()
+
+    public struct AttributeNames: Sendable {
+        public var databaseOperationName: String = "db.operation.name"
+        public var databaseSystemName: String = "db.system.name"
+        public var networkPeerAddress: String = "network.peer.address"
+        public var networkPeerPort: String = "network.peer.port"
+        public var serverAddress: String = "server.address"
+        public var serverPort: String = "server.port"
+    }
+
+    public struct AttributeValues: Sendable {
+        public var databaseSystem: String = "valkey"
+    }
+}
+#endif
diff --git a/Sources/Valkey/ValkeyClientConfiguration.swift b/Sources/Valkey/ValkeyClientConfiguration.swift
@@ -118,6 +118,10 @@ public struct ValkeyClientConfiguration: Sendable {
     /// The TLS to use for the Valkey connection.
     public var tls: TLS
 
+    #if DistributedTracingSupport
+    public var tracing: ValkeyTracingConfiguration = .init()
+    #endif
+
     /// Creates a Valkey client connection configuration.
     ///
     /// - Parameters:
diff --git a/Sources/Valkey/ValkeyConnectionFactory.swift b/Sources/Valkey/ValkeyConnectionFactory.swift
@@ -103,7 +103,7 @@ package final class ValkeyConnectionFactory: Sendable {
                 try await .enable(self.cache!.getSSLContext(), tlsServerName: clientName)
             }
 
-        return ValkeyConnectionConfiguration(
+        let newConfig = ValkeyConnectionConfiguration(
             authentication: self.configuration.authentication.flatMap {
                 .init(username: $0.username, password: $0.password)
             },
@@ -112,5 +112,13 @@ package final class ValkeyConnectionFactory: Sendable {
             tls: tls,
             clientName: nil
         )
+
+        #if DistributedTracingSupport
+        var mConfig = newConfig
+        mConfig.tracing = self.configuration.tracing
+        return mConfig
+        #else
+        return newConfig
+        #endif
     }
 }

Original file line number	Diff line number	Diff line change
`@@ -103,7 +103,7 @@ package final class ValkeyConnectionFactory: Sendable {`
`103`	`103`	`try await .enable(self.cache!.getSSLContext(), tlsServerName: clientName)`
`104`	`104`	`}`
`105`	`105`
`106`		`- return ValkeyConnectionConfiguration(`
	`106`	`+ let newConfig = ValkeyConnectionConfiguration(`
`107`	`107`	`authentication: self.configuration.authentication.flatMap {`
`108`	`108`	`.init(username: $0.username, password: $0.password)`
`109`	`109`	`},`
`@@ -112,5 +112,13 @@ package final class ValkeyConnectionFactory: Sendable {`
`112`	`112`	`tls: tls,`
`113`	`113`	`clientName: nil`
`114`	`114`	`)`
	`115`	`+`
	`116`	`+ #if DistributedTracingSupport`
	`117`	`+ var mConfig = newConfig`
	`118`	`+ mConfig.tracing = self.configuration.tracing`
	`119`	`+ return mConfig`
	`120`	`+ #else`
	`121`	`+ return newConfig`
	`122`	`+ #endif`
`115`	`123`	`}`
`116`	`124`	`}`