Enable mutating sampled flag (#9)

* Represent trace flags as OptionSet

* Document TraceFlags 📖

Author: slashmo

Sources/W3CTraceContext/TraceContext.swift

@@ -30,6 +30,22 @@ public struct TraceContext: Equatable {
         self.state = state
     }
 
+    /// Whether the caller may have recorded trace data.
+    ///
+    /// - SeeAlso: [W3C TraceContext: Sampled flag](https://www.w3.org/TR/2020/REC-trace-context-1-20200206/#sampled-flag)
+    public var sampled: Bool {
+        get {
+            self.parent.traceFlags.contains(.sampled)
+        }
+        set {
+            if newValue {
+                self.parent.traceFlags.insert(.sampled)
+            } else {
+                self.parent.traceFlags.remove(.sampled)
+            }
+        }
+    }
+
     /// Create a `TraceContext` by parsing the given header values.
     ///
     /// - Parameters:

Sources/W3CTraceContext/TraceFlags.swift

@@ -0,0 +1,32 @@
+//===----------------------------------------------------------------------===//
+//
+// This source file is part of the Swift W3C Trace Context open source project
+//
+// Copyright (c) 2020 Moritz Lang and the Swift W3C Trace Context project authors
+// Licensed under Apache License v2.0
+//
+// See LICENSE.txt for license information
+//
+// SPDX-License-Identifier: Apache-2.0
+//
+//===----------------------------------------------------------------------===//
+
+/// Represents an 8-bit field that controls tracing flags such as sampling.
+public struct TraceFlags: OptionSet {
+    public let rawValue: UInt8
+
+    public init(rawValue: UInt8) {
+        self.rawValue = rawValue
+    }
+
+    /// [W3C TraceContext: Sampled flag](https://www.w3.org/TR/2020/REC-trace-context-1-20200206/#sampled-flag)
+    public static let sampled = TraceFlags(rawValue: 1 << 0)
+}
+
+extension TraceFlags: CustomStringConvertible {
+    public var description: String {
+        let radix = 2
+        let unpadded = String(self.rawValue, radix: radix, uppercase: false)
+        return String(repeating: "0", count: radix - unpadded.count) + unpadded
+    }
+}

Sources/W3CTraceContext/TraceParent.swift

@@ -31,22 +31,14 @@ public struct TraceParent {
     /// An 8-bit field that controls tracing flags such as sampling, trace level, etc.
     ///
     /// - SeeAlso: [W3C TraceContext: trace-flags](https://www.w3.org/TR/2020/REC-trace-context-1-20200206/#trace-flags)
-    public let traceFlags: String
+    public internal(set) var traceFlags: TraceFlags
 
-    init(traceID: String, parentID: String, traceFlags: String) {
+    init(traceID: String, parentID: String, traceFlags: TraceFlags) {
         self.traceID = traceID
         self.parentID = parentID
         self.traceFlags = traceFlags
     }
 
-    /// When `true`, the least significant bit (right-most), denotes that the caller may have recorded trace data.
-    /// When `false`, the caller did not record trace data out-of-band.
-    ///
-    /// - SeeAlso: [W3C TraceContext: Sampled flag](https://www.w3.org/TR/2020/REC-trace-context-1-20200206/#sampled-flag)
-    public var sampled: Bool {
-        self.traceFlags == "01"
-    }
-
     /// The HTTP header name for `TraceParent`.
     public static let headerName = "traceparent"
 
@@ -90,9 +82,8 @@ extension TraceParent: RawRepresentable {
         self.parentID = String(parentIDComponent)
 
         // trace-flags
-        let traceFlagsComponent = components[3]
-        guard traceFlagsComponent.count == 2 else { return nil }
-        self.traceFlags = String(traceFlagsComponent)
+        guard let traceFlags = UInt8(components[3], radix: 2).map(TraceFlags.init) else { return nil }
+        self.traceFlags = traceFlags.rawValue <= 1 ? traceFlags : []
     }
 
     /// A `String` representation of this trace parent, suitable for injecting into HTTP headers.
@@ -115,7 +106,7 @@ extension TraceParent {
     public static func random<G: RandomNumberGenerator>(using generator: inout G) -> TraceParent {
         let traceID = Self.randomTraceID(using: &generator)
         let parentID = Self.randomParentID(using: &generator)
-        return .init(traceID: traceID, parentID: parentID, traceFlags: UInt64(0).paddedHexString(radix: 2))
+        return .init(traceID: traceID, parentID: parentID, traceFlags: [])
     }
 
     /// Returns a random `TraceParent` using the system random number generator.

Tests/W3CTraceContextTests/TraceContextTests.swift

@@ -30,7 +30,7 @@ final class TraceContextTests: XCTestCase {
                 parent: TraceParent(
                     traceID: "0af7651916cd43dd8448eb211c80319c",
                     parentID: "b7ad6b7169203331",
-                    traceFlags: "01"
+                    traceFlags: .sampled
                 ),
                 state: TraceState([("rojo", "00f067aa0ba902b7")])
             )
@@ -55,7 +55,7 @@ final class TraceContextTests: XCTestCase {
                 parent: TraceParent(
                     traceID: "0af7651916cd43dd8448eb211c80319c",
                     parentID: "b7ad6b7169203331",
-                    traceFlags: "01"
+                    traceFlags: .sampled
                 ),
                 state: TraceState([])
             )
@@ -77,4 +77,17 @@ final class TraceContextTests: XCTestCase {
 
         XCTAssertNotEqual(newTraceContext.parent.parentID, traceContext.parent.parentID)
     }
+
+    func test_update_sampled_flag_sets_parent_trace_flags() {
+        var traceContext = TraceContext(parent: .random(), state: .none)
+        XCTAssertFalse(traceContext.sampled)
+
+        traceContext.sampled = true
+        XCTAssert(traceContext.sampled)
+        XCTAssert(traceContext.parent.traceFlags.contains(.sampled))
+
+        traceContext.sampled = false
+        XCTAssertFalse(traceContext.sampled)
+        XCTAssertFalse(traceContext.parent.traceFlags.contains(.sampled))
+    }
 }

Tests/W3CTraceContextTests/TraceParentTests.swift

@@ -21,7 +21,7 @@ final class TraceParentRawRepresentableTests: XCTestCase {
         let traceParent = TraceParent(
             traceID: "0af7651916cd43dd8448eb211c80319c",
             parentID: "b7ad6b7169203331",
-            traceFlags: "01"
+            traceFlags: .sampled
         )
 
         XCTAssertEqual(traceParent.rawValue, "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01")
@@ -38,10 +38,21 @@ final class TraceParentRawRepresentableTests: XCTestCase {
 
         XCTAssertEqual(
             traceParent,
-            TraceParent(traceID: "0af7651916cd43dd8448eb211c80319c", parentID: "b7ad6b7169203331", traceFlags: "01")
+            TraceParent(traceID: "0af7651916cd43dd8448eb211c80319c", parentID: "b7ad6b7169203331", traceFlags: .sampled)
         )
+    }
 
-        XCTAssert(traceParent.sampled)
+    func testDecodeValidTraceParentStringWithUnsupportedTraceFlags() {
+        let rawValue = "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-11"
+        guard let traceParent = TraceParent(rawValue: rawValue) else {
+            XCTFail("Could not decode valid trace parent")
+            return
+        }
+
+        XCTAssertEqual(
+            traceParent,
+            TraceParent(traceID: "0af7651916cd43dd8448eb211c80319c", parentID: "b7ad6b7169203331", traceFlags: [])
+        )
     }
 
     func testDecodeValidTraceParentStringWithoutSampledFlag() {
@@ -53,10 +64,10 @@ final class TraceParentRawRepresentableTests: XCTestCase {
 
         XCTAssertEqual(
             traceParent,
-            TraceParent(traceID: "0af7651916cd43dd8448eb211c80319c", parentID: "b7ad6b7169203331", traceFlags: "00")
+            TraceParent(traceID: "0af7651916cd43dd8448eb211c80319c", parentID: "b7ad6b7169203331", traceFlags: [])
         )
 
-        XCTAssertFalse(traceParent.sampled)
+        XCTAssertFalse(traceParent.traceFlags.contains(.sampled))
     }
 
     func testDecodeFailsWithTooLongRawValue() {