Move TracingInstrument & related types to separate library

Unit-test MultiplexInstrument instance casting

Adjust naming of InstrumentationSystem getters

Document InstrumentationSystem getters 📖

Author: slashmo

Package.resolved

@@ -6,6 +6,7 @@ let package = Package(
     products: [
         .library(name: "BaggageLogging", targets: ["BaggageLogging"]),
         .library(name: "Instrumentation", targets: ["Instrumentation"]),
+        .library(name: "TracingInstrumentation", targets: ["TracingInstrumentation"]),
         .library(name: "NIOInstrumentation", targets: ["NIOInstrumentation"])
     ],
     dependencies: [
@@ -52,6 +53,22 @@ let package = Package(
             ]
         ),
 
+        .target(
+            name: "TracingInstrumentation",
+            dependencies: [
+                .product(name: "Baggage", package: "swift-baggage-context"),
+                "Instrumentation"
+            ]
+        ),
+        .testTarget(
+            name: "TracingInstrumentationTests",
+            dependencies: [
+                "Instrumentation",
+                "TracingInstrumentation",
+                "BaggageLogging"
+            ]
+        ),
+
         .target(
             name: "NIOInstrumentation",
             dependencies: [

Package.swift

@@ -14,13 +14,24 @@
 import Baggage
 
 /// `InstrumentationSystem` is a global facility where the default cross-cutting tool can be configured.
-/// It is set up just once in a given program select the desired `Instrument` implementation.
+/// It is set up just once in a given program to select the desired `Instrument` implementation.
 ///
-/// - Note: If you need to use more that one cross-cutting tool you can do so by using `MultiplexInstrument`.
+/// # Bootstrap multiple Instruments
+/// If you need to use more that one cross-cutting tool you can do so by using `MultiplexInstrument`.
+///
+/// # Access the Instrument
+/// There are two ways of getting the bootstrapped instrument.
+/// 1. `InstrumentationSystem.instrument`: Returns whatever you passed to `.bootstrap` as an `Instrument`.
+/// 2. `InstrumentationSystem.instrument(of: MyInstrument.self)`: Returns the bootstrapped `Instrument` if it's
+/// an instance of the given type or the first instance of `MyInstrument` if it's part of a `MultiplexInstrument`.
+///
+/// ## What getter to use
+/// - Default to using `InstrumentationSystem.instrument`
+/// - Use `InstrumentationSystem.instrument(of: MyInstrument.self)` only if you need to use specific `MyInstrument` APIs
+/// - Protocols that "inherit" from `Instrument` may also wrap `.instrument(of: TheirInstrument.self)` in a convenience method
 public enum InstrumentationSystem {
     private static let lock = ReadWriteLock()
     private static var _instrument: Instrument = NoOpInstrument()
-    private static var _tracer: TracingInstrument?
     private static var isInitialized = false
 
     /// Globally select the desired `Instrument` implementation.
@@ -35,21 +46,14 @@ public enum InstrumentationSystem {
                 you need to use multiple instruments.
                 """
             )
-            if let tracer = instrument as? TracingInstrument {
-                self._tracer = tracer
-            }
             self._instrument = instrument
-
             self.isInitialized = true
         }
     }
 
     // for our testing we want to allow multiple bootstrapping
     internal static func bootstrapInternal(_ instrument: Instrument) {
         self.lock.withWriterLock {
-            if let tracer = instrument as? TracingInstrument {
-                self._tracer = tracer
-            }
             self._instrument = instrument
         }
     }
@@ -58,13 +62,20 @@ public enum InstrumentationSystem {
     public static var instrument: Instrument {
         self.lock.withReaderLock { self._instrument }
     }
+}
 
-    // FIXME: smarter impl
-    public static var tracer: TracingInstrument {
+extension InstrumentationSystem {
+    /// Get an `Instrument` instance of the given type.
+    ///
+    /// When using `MultiplexInstrument`, this returns the first instance of the given type stored in the `MultiplexInstrument`.
+    /// - Parameter instrumentType: The type of `Instrument` you want to retrieve an instance for.
+    /// - Returns: An `Instrument` instance of the given type or `nil` if no `Instrument` of that type has been bootstrapped.
+    public static func instrument<I>(of instrumentType: I.Type) -> I? {
         self.lock.withReaderLock {
-            let tracer: TracingInstrument? = self._tracer
-            let res: TracingInstrument = tracer ?? NoOpTracingInstrument()
-            return res
+            if let multiplexInstrument = self._instrument as? MultiplexInstrument {
+                return multiplexInstrument.firstInstance(of: I.self)
+            }
+            return self._instrument as? I
         }
     }
 }

Sources/Instrumentation/InstrumentationSystem.swift

@@ -27,6 +27,12 @@ public struct MultiplexInstrument {
     }
 }
 
+extension MultiplexInstrument {
+    func firstInstance<I>(of instrument: I.Type) -> I? {
+        self.instruments.first(where: { $0 is I }) as? I
+    }
+}
+
 extension MultiplexInstrument: Instrument {
     public func inject<Carrier, Injector>(
         _ baggage: BaggageContext, into carrier: inout Carrier, using injector: Injector

Sources/Instrumentation/MultiplexInstrument.swift

@@ -12,6 +12,7 @@
 //===----------------------------------------------------------------------===//
 
 import Baggage
+import Instrumentation
 
 /// No operation TracingInstrument, used when no tracing is required.
 public struct NoOpTracingInstrument: TracingInstrument {

Sources/Instrumentation/Tracing/NoOpTracing.swift renamed to Sources/TracingInstrumentation/NoOpTracing.swift

@@ -12,6 +12,7 @@
 //===----------------------------------------------------------------------===//
 
 import Baggage
+import Instrumentation
 
 /// An `Instrument` with added functionality for distributed tracing. Is uses the span-based tracing model and is
 /// based on the OpenTracing/OpenTelemetry spec.
@@ -46,3 +47,13 @@ extension TracingInstrument {
         self.startSpan(named: operationName, context: context, ofKind: .internal, at: nil)
     }
 }
+
+extension InstrumentationSystem {
+    /// Returns the `TracingInstrument` bootstrapped as part of the `InstrumentationSystem`.
+    ///
+    /// - Warning: Only call this after you bootstrapped the `InstrumentationSystem`. Calling it before is
+    /// considered a programmer error and leads to a crash.
+    public static var tracingInstrument: TracingInstrument {
+        instrument(of: TracingInstrument.self)!
+    }
+}

Sources/Instrumentation/Tracing/Span.swift renamed to Sources/TracingInstrumentation/Span.swift

@@ -12,26 +12,43 @@
 //===----------------------------------------------------------------------===//
 
 import Baggage
-import Instrumentation
+@testable import Instrumentation
 import XCTest
 
 final class InstrumentationSystemTests: XCTestCase {
     func testItProvidesAccessToASingletonInstrument() {
         let tracer = FakeTracer()
+        let instrument = FakeInstrument()
+        let multiplexInstrument = MultiplexInstrument([tracer, instrument])
 
-        InstrumentationSystem.bootstrap(tracer) // FIXME, must use bootstrapInternal
-        XCTAssertTrue(InstrumentationSystem.instrument as? FakeTracer === tracer)
+        XCTAssertNil(InstrumentationSystem.instrument(of: FakeTracer.self))
+        XCTAssertNil(InstrumentationSystem.instrument(of: FakeInstrument.self))
+
+        InstrumentationSystem.bootstrapInternal(multiplexInstrument)
+        XCTAssert(InstrumentationSystem.instrument is MultiplexInstrument)
+        XCTAssert(InstrumentationSystem.instrument(of: FakeTracer.self) === tracer)
+        XCTAssert(InstrumentationSystem.instrument(of: FakeInstrument.self) === instrument)
+
+        InstrumentationSystem.bootstrapInternal(tracer)
+        XCTAssertFalse(InstrumentationSystem.instrument is MultiplexInstrument)
+        XCTAssert(InstrumentationSystem.instrument(of: FakeTracer.self) === tracer)
+        XCTAssertNil(InstrumentationSystem.instrument(of: FakeInstrument.self))
     }
 }
 
 private final class FakeTracer: Instrument {
-    enum TraceIDKey: BaggageContextKey {
-        typealias Value = String
-    }
+    func inject<Carrier, Injector>(_ baggage: BaggageContext, into carrier: inout Carrier, using injector: Injector)
+        where
+        Injector: InjectorProtocol,
+        Carrier == Injector.Carrier {}
 
-    static let headerName = "fake-trace-id"
-    static let defaultTraceID = UUID().uuidString
+    func extract<Carrier, Extractor>(_ carrier: Carrier, into baggage: inout BaggageContext, using extractor: Extractor)
+        where
+        Extractor: ExtractorProtocol,
+        Carrier == Extractor.Carrier {}
+}
 
+private final class FakeInstrument: Instrument {
     func inject<Carrier, Injector>(_ baggage: BaggageContext, into carrier: inout Carrier, using injector: Injector)
         where
         Injector: InjectorProtocol,

Sources/Instrumentation/Tracing/Timestamp.swift renamed to Sources/TracingInstrumentation/Timestamp.swift

@@ -13,6 +13,7 @@
 
 import Baggage
 import Instrumentation
+import TracingInstrumentation
 import XCTest
 
 final class SpanTests: XCTestCase {