feat: traced() (#96)

* feat: traced()

fixes: #38

* chore: Improve tests

---------

Co-authored-by: danthorpe <danthorpe@users.noreply.github.com>

Author: danthorpe

Sources/Helpers/Data+Crypto.swift

@@ -0,0 +1,14 @@
+import CryptoKit
+import Foundation
+
+extension Data {
+
+  static func secureRandomData(length: UInt) -> Data? {
+    let count = Int(length)
+    var bytes = [Int8](repeating: 0, count: count)
+    guard errSecSuccess == SecRandomCopyBytes(kSecRandomDefault, count, &bytes) else {
+      return nil
+    }
+    return Data(bytes: bytes, count: count)
+  }
+}

Sources/Helpers/UniqueIdentifier.swift

@@ -0,0 +1,72 @@
+import ConcurrencyExtras
+import Foundation
+
+package enum UniqueIdentifier: Hashable {
+  package enum Format: Hashable {
+    case base64, hex
+  }
+  case secureBytes(length: UInt = 10, format: Format)
+}
+
+extension UniqueIdentifier {
+
+  func generate() -> String {
+    switch self {
+    case let .secureBytes(length, _):
+      var data = Data()
+      repeat {
+        data = .secureRandomData(length: length) ?? Data()
+      } while data.isEmpty
+      return format(data: data)
+    }
+  }
+
+  func format(data: Data) -> String {
+    switch self {
+    case .secureBytes(_, .base64):
+      return data.base64EncodedString(options: [])
+    case .secureBytes(_, .hex):
+      return data.map { String(format: "%02hhx", $0) }.joined()
+    }
+  }
+}
+
+// MARK: - Generator
+
+extension UniqueIdentifier {
+  package struct Generator: Sendable {
+    private let generate: @Sendable () -> String
+
+    package init(_ id: UniqueIdentifier) {
+      self.init { id.generate() }
+    }
+
+    package init(generate: @escaping @Sendable () -> String) {
+      self.generate = generate
+    }
+
+    @discardableResult
+    package func callAsFunction() -> String {
+      generate()
+    }
+  }
+}
+
+extension UniqueIdentifier.Generator {
+  package static func constant(_ id: UniqueIdentifier) -> Self {
+    let generation = id.generate()
+    return Self { generation }
+  }
+
+  package static func incrementing(_ id: UniqueIdentifier) -> Self {
+    let sequence = LockIsolated<Int>(0)
+    return Self {
+      let number = sequence.withValue {
+        $0 += 1
+        return $0
+      }
+      let data = withUnsafeBytes(of: number.bigEndian) { Data($0) }
+      return id.format(data: data)
+    }
+  }
+}

Sources/Networking/Components/Traced.swift

@@ -0,0 +1,113 @@
+import Dependencies
+import DependenciesMacros
+import Foundation
+import HTTPTypes
+import Helpers
+
+extension NetworkingComponent {
+
+  /// Generates a HTTP Trace Parent header for each request.
+  ///
+  /// - See-Also: [Trace-Context](https://www.w3.org/TR/trace-context/)
+  public func traced() -> some NetworkingComponent {
+    modified(Traced())
+  }
+}
+
+private struct Traced: NetworkingModifier {
+  @Dependency(\.traceParentGenerator) var generate
+  func resolve(upstream: some NetworkingComponent, request: HTTPRequestData) -> HTTPRequestData {
+    guard nil == request.traceParent else {
+      return request
+    }
+    var copy = request
+    copy.traceParent = generate()
+    return copy
+  }
+}
+
+extension HTTPField.Name {
+  public static let traceparent = HTTPField.Name("traceparent")!
+}
+
+extension HTTPRequestData {
+  package fileprivate(set) var traceParent: TraceParent? {
+    get { self[option: TraceParent.self] }
+    set {
+      self[option: TraceParent.self] = newValue
+      self.headerFields[.traceparent] = newValue?.description
+    }
+  }
+
+  public var traceId: String? {
+    traceParent?.traceId
+  }
+
+  public var parentId: String? {
+    traceParent?.parentId
+  }
+}
+
+public struct TraceParent: Sendable, Hashable, HTTPRequestDataOption {
+  public static var defaultOption: Self?
+
+  // Current version of the spec only supports 01 flag
+  // Future versions of the spec will require support for bit-field mask
+  public let traceId: String
+  public let parentId: String
+
+  public var description: String {
+    "00-\(traceId)-\(parentId)-01"
+  }
+
+  public init(traceId: String, parentId: String) {
+    self.traceId = traceId
+    self.parentId = parentId
+  }
+}
+
+// MARK: - Generator
+
+@DependencyClient
+public struct TraceParentGenerator: Sendable {
+  public var generate: @Sendable () -> TraceParent = {
+    TraceParent(traceId: "dummy-trace-id", parentId: "dummy-parent-id")
+  }
+
+  package func callAsFunction() -> TraceParent {
+    generate()
+  }
+}
+
+extension TraceParentGenerator: DependencyKey {
+  public static let liveValue = {
+    let traceId = UniqueIdentifier.Generator(.secureBytes(length: 16, format: .hex))
+    let parentId = UniqueIdentifier.Generator(.secureBytes(length: 8, format: .hex))
+    return TraceParentGenerator {
+      TraceParent(
+        traceId: traceId(),
+        parentId: parentId()
+      )
+    }
+  }()
+}
+
+extension DependencyValues {
+  public var traceParentGenerator: TraceParentGenerator {
+    get { self[TraceParentGenerator.self] }
+    set { self[TraceParentGenerator.self] = newValue }
+  }
+}
+
+extension TraceParentGenerator {
+  public static var incrementing: TraceParentGenerator {
+    let traceId = UniqueIdentifier.Generator.incrementing(.secureBytes(length: 16, format: .hex))
+    let parentId = UniqueIdentifier.Generator.incrementing(.secureBytes(length: 8, format: .hex))
+    return TraceParentGenerator {
+      TraceParent(
+        traceId: traceId(),
+        parentId: parentId()
+      )
+    }
+  }
+}

Sources/TestSupport/Mocked.swift

@@ -2,6 +2,13 @@ import Networking
 
 extension NetworkingComponent {
 
+  /// Mock all requests with a stub
+  public func mocked(
+    all stub: StubbedResponseStream
+  ) -> some NetworkingComponent {
+    mocked(stub) { _ in true }
+  }
+
   /// Mock a given request with a stub
   public func mocked(
     _ request: HTTPRequestData,

Sources/TestSupport/NetworkingTestCase.swift

@@ -23,6 +23,7 @@ open class NetworkingTestCase: XCTestCase {
   ) {
     withDependencies {
       $0.shortID = shortIdGenerator ?? .incrementing
+      $0.traceParentGenerator = .incrementing
       $0.continuousClock = continuousClock ?? TestClock()
       updateValuesForOperation(&$0)
     } operation: {

Tests/NetworkingTests/Components/TracedTests.swift

@@ -0,0 +1,101 @@
+import Foundation
+import TestSupport
+import XCTest
+
+@testable import Networking
+
+final class TracedTests: NetworkingTestCase {
+  override func invokeTest() {
+    withTestDependencies {
+      super.invokeTest()
+    }
+  }
+
+  func test__request_includes_trace() async throws {
+    let reporter = TestReporter()
+
+    let network = TerminalNetworkingComponent()
+      .mocked(all: .ok())
+      .reported(by: reporter)
+      .traced()
+
+    try await withThrowingTaskGroup(of: HTTPResponseData.self) { group in
+      for _ in 0 ..< 10 {
+        group.addTask {
+          try await network.data(HTTPRequestData())
+        }
+      }
+
+      var responses: [HTTPResponseData] = []
+      for try await response in group {
+        responses.append(response)
+      }
+    }
+
+    let sentRequests = await reporter.requests
+
+    XCTAssertEqual(
+      sentRequests.map(\.headerFields[.traceparent]),
+      [
+        "00-0000000000000001-0000000000000001-01",
+        "00-0000000000000002-0000000000000002-01",
+        "00-0000000000000003-0000000000000003-01",
+        "00-0000000000000004-0000000000000004-01",
+        "00-0000000000000005-0000000000000005-01",
+        "00-0000000000000006-0000000000000006-01",
+        "00-0000000000000007-0000000000000007-01",
+        "00-0000000000000008-0000000000000008-01",
+        "00-0000000000000009-0000000000000009-01",
+        "00-000000000000000a-000000000000000a-01",
+      ]
+    )
+
+    XCTAssertEqual(sentRequests.last?.traceId, "000000000000000a")
+    XCTAssertEqual(sentRequests.first?.parentId, "0000000000000001")
+  }
+
+  func test__traced_requests_do_not_get_another_trace() async throws {
+    let reporter = TestReporter()
+
+    let network = TerminalNetworkingComponent()
+      .mocked(all: .ok())
+      .reported(by: reporter)
+      .traced()
+
+    // Make an initial request
+    let response = try await network.data(HTTPRequestData())
+
+    // Resend the request
+    try await network.data(response.request)
+
+    let sentRequests = await reporter.requests
+
+    XCTAssertEqual(
+      sentRequests.map(\.headerFields[.traceparent]),
+      [
+        "00-0000000000000001-0000000000000001-01",
+        "00-0000000000000001-0000000000000001-01",
+      ]
+    )
+  }
+
+  func test__live_trace_generator() async {
+    let generate = TraceParentGenerator.liveValue
+
+    let traces = await withTaskGroup(of: TraceParent.self) { group in
+      for _ in 0 ..< 10 {
+        group.addTask {
+          generate()
+        }
+      }
+
+      var traces: [TraceParent] = []
+      for await trace in group {
+        traces.append(trace)
+      }
+      return traces
+    }
+
+    XCTAssertEqual(Set(traces).count, 10)
+  }
+}