allow for tracingInstrument and tracingInstrument(of: Type) APIs (#101)

allow for tracingInstrument and tracingInstrument(of: Type) APIs

Author: ktoso

Sources/Instrumentation/Instrument.swift

@@ -67,14 +67,3 @@ public protocol Instrument {
         Injector.Carrier == Carrier
 }
 
-public struct NoOpInstrument: Instrument {
-    public func inject<Carrier, Injector>(_ baggage: BaggageContext, into carrier: inout Carrier, using injector: Injector)
-        where
-        Injector: InjectorProtocol,
-        Carrier == Injector.Carrier {}
-
-    public func extract<Carrier, Extractor>(_ carrier: Carrier, into baggage: inout BaggageContext, using extractor: Extractor)
-        where
-        Extractor: ExtractorProtocol,
-        Carrier == Extractor.Carrier {}
-}

Sources/Instrumentation/InstrumentationSystem.swift

@@ -61,24 +61,36 @@ public enum InstrumentationSystem {
         }
     }
 
-    /// Returns the globally configured `Instrument`. Defaults to a no-op `Instrument` if `boostrap` wasn't called before.
+    /// Returns the globally configured `Instrument`.
+    ///
+    /// Defaults to a no-op `Instrument` if `boostrap` wasn't called before.
     public static var instrument: Instrument {
         self.lock.withReaderLock { self._instrument }
     }
-}
 
-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? {
+    public static func instrument<I>(of instrumentType: I.Type) -> I? where I: Instrument {
+        self._findInstrument(where: { $0 is I }) as? I
+    }
+
+}
+
+extension InstrumentationSystem {
+    /// :nodoc: INTERNAL API: Do Not Use
+    public static func _findInstrument(where predicate: (Instrument) -> Bool) -> Instrument? {
         self.lock.withReaderLock {
-            if let multiplexInstrument = self._instrument as? MultiplexInstrument {
-                return multiplexInstrument.firstInstance(of: I.self)
+            if let multiplex = self._instrument as? MultiplexInstrument {
+                return multiplex.firstInstrument(where: predicate)
+            } else if predicate(self._instrument) {
+                return self._instrument
+            } else {
+                return nil
             }
-            return self._instrument as? I
         }
     }
 }

Sources/Instrumentation/MultiplexInstrument.swift

@@ -28,8 +28,8 @@ public struct MultiplexInstrument {
 }
 
 extension MultiplexInstrument {
-    func firstInstance<I>(of instrument: I.Type) -> I? {
-        self.instruments.first(where: { $0 is I }) as? I
+    func firstInstrument(where predicate: (Instrument) -> Bool) -> Instrument? {
+        self.instruments.first(where: predicate)
     }
 }
 

Sources/Instrumentation/NoOpInstrument.swift

@@ -0,0 +1,30 @@
+//===----------------------------------------------------------------------===//
+//
+// This source file is part of the Swift Distributed Tracing open source project
+//
+// Copyright (c) 2020 Moritz Lang and the Swift Tracing project authors
+// Licensed under Apache License v2.0
+//
+// See LICENSE.txt for license information
+//
+// SPDX-License-Identifier: Apache-2.0
+//
+//===----------------------------------------------------------------------===//
+
+import Baggage
+
+/// A "no op" implementation of an `Instrument`.
+public struct NoOpInstrument: Instrument {
+
+    public init() {}
+
+    public func inject<Carrier, Injector>(_ baggage: BaggageContext, into carrier: inout Carrier, using injector: Injector)
+        where
+        Injector: InjectorProtocol,
+        Carrier == Injector.Carrier {}
+
+    public func extract<Carrier, Extractor>(_ carrier: Carrier, into baggage: inout BaggageContext, using extractor: Extractor)
+        where
+        Extractor: ExtractorProtocol,
+        Carrier == Extractor.Carrier {}
+}

Sources/TracingInstrumentation/InstrumentationSystem+Tracing.swift

@@ -0,0 +1,42 @@
+//===----------------------------------------------------------------------===//
+//
+// This source file is part of the Swift Distributed Tracing open source project
+//
+// Copyright (c) 2020 Moritz Lang and the Swift Tracing project authors
+// Licensed under Apache License v2.0
+//
+// See LICENSE.txt for license information
+//
+// SPDX-License-Identifier: Apache-2.0
+//
+//===----------------------------------------------------------------------===//
+
+import Baggage
+import Instrumentation
+
+
+extension InstrumentationSystem {
+
+    /// Get a `TracingInstrument` instance of the given type.
+    ///
+    /// When using `MultiplexInstrument`, this returns the first instance of the given type stored in the `MultiplexInstrument`.
+    ///
+    /// Usually tracing libraries will provide their own convenience getter, e.g. CoolTracing could provide `InstrumentationSystem.coolTracer`;
+    /// if available, prefer using those APIs rather than relying on this general function.
+    ///
+    /// - 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 tracingInstrument<T>(of instrumentType: T.Type) -> T? where T: TracingInstrument {
+        self._findInstrument(where: { $0 is T }) as? T
+    }
+
+    /// Returns the `TracingInstrument` bootstrapped as part of the `InstrumentationSystem`.
+    ///
+    /// If the system was bootstrapped with a `MultiplexInstrument` this function attempts to locate the _first_
+    /// tracing instrument as passed to the multiplex instrument. If none is found, a `NoOpTracingInstrument` is returned.
+    ///
+    /// - Returns: A `TracingInstrument` if the system was bootstrapped with one, and `NoOpTracingInstrument` otherwise.
+    public static var tracingInstrument: TracingInstrument {
+        (self._findInstrument(where: { $0 is TracingInstrument }) as? TracingInstrument) ?? NoOpTracingInstrument()
+    }
+}

Sources/TracingInstrumentation/NoOpTracing.swift renamed to Sources/TracingInstrumentation/NoOpTracingInstrument.swift

@@ -47,13 +47,3 @@ 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/TracingInstrumentation/TracingInstrument.swift

@@ -13,7 +13,7 @@
 
 import Baggage
 import BaggageLogging
-import Instrumentation
+@testable import Instrumentation
 import TracingInstrumentation
 import XCTest
 
@@ -27,6 +27,30 @@ final class TracingInstrumentTests: XCTestCase {
 
         httpServer.receive(FakeHTTPRequest(path: "/", headers: []))
     }
+
+    func testItProvidesAccessToATracingInstrument() {
+        let tracer = JaegerTracer()
+
+        XCTAssertNil(InstrumentationSystem.tracingInstrument(of: JaegerTracer.self))
+
+        InstrumentationSystem.bootstrapInternal(tracer)
+        XCTAssertFalse(InstrumentationSystem.instrument is MultiplexInstrument)
+        XCTAssert(InstrumentationSystem.instrument(of: JaegerTracer.self) === tracer)
+        XCTAssertNil(InstrumentationSystem.instrument(of: NoOpInstrument.self))
+
+        XCTAssert(InstrumentationSystem.tracingInstrument(of: JaegerTracer.self) === tracer)
+        XCTAssert(InstrumentationSystem.tracingInstrument is JaegerTracer)
+
+        let multiplexInstrument = MultiplexInstrument([tracer])
+        InstrumentationSystem.bootstrapInternal(multiplexInstrument)
+        XCTAssert(InstrumentationSystem.instrument is MultiplexInstrument)
+        XCTAssert(InstrumentationSystem.instrument(of: JaegerTracer.self) === tracer)
+
+        XCTAssert(InstrumentationSystem.tracingInstrument(of: JaegerTracer.self) === tracer)
+        XCTAssert(InstrumentationSystem.tracingInstrument is JaegerTracer)
+
+
+    }
 }
 
 // MARK: - JaegerTracer