Add documentation on new APIs

ktoso · ktoso · commit f1ffc1c82626 · 2025-09-11T12:34:46.000+09:00
diff --git a/Sources/InMemoryTracing/InMemorySpanContext.swift b/Sources/InMemoryTracing/InMemorySpanContext.swift
@@ -14,9 +14,16 @@
 
 import ServiceContextModule
 
+/// Encapsulates the `traceID`, `spanID` and `parentSpanID` of an `InMemorySpan`.
+/// Generally used through the `ServiceContext/inMemorySpanContext` task local value.
 public struct InMemorySpanContext: Sendable, Hashable {
+    /// Idenfifier of top-level trace of which this span is a part of.
     public let traceID: String
+
+    /// Identifier of this specific span.
     public let spanID: String
+
+    // Identifier of the parent of this span, if any.
     public let parentSpanID: String?
 
     public init(traceID: String, spanID: String, parentSpanID: String?) {
@@ -27,7 +34,8 @@ public struct InMemorySpanContext: Sendable, Hashable {
 }
 
 extension ServiceContext {
-    var inMemorySpanContext: InMemorySpanContext? {
+    /// Task-local value representing the current tracing ``Span`` as set by the ``InMemoryTracer``.
+    public var inMemorySpanContext: InMemorySpanContext? {
         get {
             self[InMemorySpanContextKey.self]
         }
diff --git a/Sources/InMemoryTracing/InMemoryTracer.swift b/Sources/InMemoryTracing/InMemoryTracer.swift
@@ -15,22 +15,50 @@
 @_spi(Locking) import Instrumentation
 import Tracing
 
+/// An in-memory implementation of the ``Tracer`` protocol which can be used either in testing,
+/// or in manual collecting and interrogating traces within a process, and acting on them programatically.
+/// 
+/// ### Span lifecycle
+/// This tracer does _not_ automatically remove spans once they end.
+/// Finished spans are retained and available for inspection using the `finishedSpans` property.
+/// Spans which have been started but have not yet been called `Span/end()` on are also available 
+/// for inspection using the ``activeSpans`` property.
+/// 
+/// Spans are retained by the `InMemoryTracer` until they are explicitly removed, e.g. by using 
+/// ``popFinishedSpans()`` or any of the `clear...` methods (e.g. ``clearFinishedSpans()``)
 public struct InMemoryTracer: Tracer {
 
-    public static let traceIDKey = "test-trace-id"
-    public static let spanIDKey = "test-span-id"
-
     public let idGenerator: IDGenerator
 
+    public let recordInjections: Bool
+    public let recordExtractions: Bool
+
     private let _activeSpans = LockedValueBox<[InMemorySpanContext: InMemorySpan]>([:])
     private let _finishedSpans = LockedValueBox<[FinishedInMemorySpan]>([])
     private let _numberOfForceFlushes = LockedValueBox<Int>(0)
 
-    public init(idGenerator: IDGenerator = .incrementing) {
+    private let _injections = LockedValueBox<[Injection]>([])
+    private let _extractions = LockedValueBox<[Extraction]>([])
+
+    /// Create a new ``InMemoryTracer``.
+    /// 
+    /// - Parameters: 
+    ///   - Parameter idGenerator: strategy for generating trace and span identifiers
+    ///   - Parameter idGenerator: strategy for generating trace and span identifiers
+    public init(
+        idGenerator: IDGenerator = .incrementing,
+        recordInjections: Bool = true,
+        recordExtractions: Bool = true
+        ) {
         self.idGenerator = idGenerator
+        self.recordInjections = recordInjections
+        self.recordExtractions = recordExtractions
     }
+}
 
-    // MARK: - Tracer
+// MARK: - Tracer
+
+extension InMemoryTracer {
 
     public func startSpan<Instant>(
         _ operationName: String,
@@ -77,24 +105,66 @@ public struct InMemoryTracer: Tracer {
         return span
     }
 
+    public func forceFlush() {
+        _numberOfForceFlushes.withValue { $0 += 1 }
+    }
+}
+
+// MARK: - InMemoryTracer querying
+
+extension InMemoryTracer {
+
+    /// Array of active spans, i.e. spans which have been started by have not yet finished (by calling `Span/end()`).
+    public var activeSpans: [InMemorySpan] {
+        _activeSpans.withValue { active in Array(active.values) }
+    }
+
+    /// Retrives a specific _active_ span, identified by the specific span, trace, and parent ID's
+    /// stored in the `inMemorySpanContext`
     public func activeSpan(identifiedBy context: ServiceContext) -> InMemorySpan? {
         guard let spanContext = context.inMemorySpanContext else { return nil }
         return _activeSpans.withValue { $0[spanContext] }
     }
 
-    public func forceFlush() {
-        _numberOfForceFlushes.withValue { $0 += 1 }
-    }
-
+    /// Count of the number of times ``Tracer/forceFlush()`` was called on this tracer.
     public var numberOfForceFlushes: Int {
         _numberOfForceFlushes.withValue { $0 }
     }
 
+    /// Gets, without removing, all the finished spans recorded by this tracer.
+    /// 
+    /// - SeeAlso: `popFinishedSpans()`
     public var finishedSpans: [FinishedInMemorySpan] {
         _finishedSpans.withValue { $0 }
     }
+    
+    /// Returns, and removes, all finished spans recorded by this tracer.
+    public func popFinishedSpans() -> [FinishedInMemorySpan] {
+        _finishedSpans.withValue { spans in 
+            defer { spans = [] }
+            return spans
+        }
+    }
 
-    // MARK: - Instrument
+    /// Atomically clears any stored finished spans in this tracer.
+    public func clearFinishedSpans() {
+        _finishedSpans.withValue { $0 = [] }
+    }
+
+    /// Clears all registered finished spans, as well as injections/extractions performed by this tracer.
+    public func clearAll() {
+        _finishedSpans.withValue { $0 = [] }
+        _injections.withValue { $0 = [] }
+        _extractions.withValue { $0 = [] }
+    }
+}
+
+// MARK: - Instrument
+
+extension InMemoryTracer {
+
+    public static let traceIDKey = "in-memory-trace-id"
+    public static let spanIDKey = "in-memory-span-id"
 
     public func inject<Carrier, Inject: Injector>(
         _ context: ServiceContext,
@@ -110,22 +180,44 @@ public struct InMemoryTracer: Tracer {
             values[Self.spanIDKey] = spanContext.spanID
         }
 
-        let injection = Injection(context: context, values: values)
-        _injections.withValue { $0.append(injection) }
+        if recordInjections {
+            let injection = Injection(context: context, values: values)
+            _injections.withValue { $0.append(injection) }
+        }
     }
 
-    public var injections: [Injection] {
+    /// Lists all recorded calls to this tracer's ``Instrument/inject(_:into:using:)`` method.
+    /// This may be used to inspect what span identifiers are being propagated by this tracer.
+    public var performedContextInjections: [Injection] {
         _injections.withValue { $0 }
     }
 
+    /// Clear the list of recorded context injections (calls to ``Instrument/inject(_:into:using:)``).
+    public func clearPerformedContextInjections() {
+        _injections.withValue { $0 = [] }
+    }
+
+    /// Represents a recorded call to the InMemoryTracer's ``Instrument/inject(_:into:using:)`` method.
+    public struct Injection: Sendable {
+        /// The context from which values were being injected.
+        public let context: ServiceContext
+        /// The injected values, these will be specifically the trace and span identifiers of the propagated span.
+        public let values: [String: String]
+    }
+}
+
+extension InMemoryTracer {
+
     public func extract<Carrier, Extract: Extractor>(
         _ carrier: Carrier,
         into context: inout ServiceContext,
         using extractor: Extract
     ) where Carrier == Extract.Carrier {
         defer {
-            let extraction = Extraction(carrier: carrier, context: context)
-            _extractions.withValue { $0.append(extraction) }
+            if self.recordExtractions {
+                let extraction = Extraction(carrier: carrier, context: context)
+                _extractions.withValue { $0.append(extraction) }
+            }
         }
 
         guard let traceID = extractor.extract(key: Self.traceIDKey, from: carrier),
@@ -137,27 +229,30 @@ public struct InMemoryTracer: Tracer {
         context.inMemorySpanContext = InMemorySpanContext(traceID: traceID, spanID: spanID, parentSpanID: nil)
     }
 
-    public var extractions: [Extraction] {
+    /// Lists all recorded calls to this tracer's ``Instrument/extract(_:into:using:)`` method.
+    /// This may be used to inspect what span identifiers were extracted from an incoming carrier object into ``ServiceContext``.
+    public var performedContextExtractions: [Extraction] {
         _extractions.withValue { $0 }
     }
 
-    public struct Injection: Sendable {
-        public let context: ServiceContext
-        public let values: [String: String]
-    }
-
+    
+    /// Represents a recorded call to the InMemoryTracer's ``Instrument/extract(_:into:using:)`` method.
     public struct Extraction: Sendable {
+        /// The carrier object from which the context values were extracted from,
+        /// e.g. this frequently is an HTTP request or similar.
         public let carrier: any Sendable
+        /// The constructed service context, containing the extracted ``ServiceContext/inMemoryTraceContext``.
         public let context: ServiceContext
     }
-
-    private let _injections = LockedValueBox<[Injection]>([])
-    private let _extractions = LockedValueBox<[Extraction]>([])
 }
 
 // MARK: - ID Generator
 
 extension InMemoryTracer {
+
+    /// Can be used to customize how trace and span IDs are generated by the ``InMemoryTracer``.
+    /// 
+    /// Defaults to a simple sequential numeric scheme (`span-1`, `span-2`, `trace-1`, `trace-2` etc).
     public struct IDGenerator: Sendable {
         public let nextTraceID: @Sendable () -> String
         public let nextSpanID: @Sendable () -> String
diff --git a/Tests/InMemoryTracingTests/InMemoryTracerTests.swift b/Tests/InMemoryTracingTests/InMemoryTracerTests.swift
@@ -106,15 +106,15 @@ struct InMemoryTracerTests {
                 spanID: "stub",
                 parentSpanID: "stub"
             )
-            context.testSpanContext = spanContext
+            context.inMemorySpanContext = spanContext
 
             var values = [String: String]()
             tracer.inject(context, into: &values, using: DictionaryInjector())
 
             #expect(values == [InMemoryTracer.traceIDKey: "stub", InMemoryTracer.spanIDKey: "stub"])
 
-            let injection = try #require(tracer.injections.first)
-            #expect(injection.context.testSpanContext == spanContext)
+            let injection = try #require(tracer.performedContextInjections.first)
+            #expect(injection.context.inMemorySpanContext == spanContext)
             #expect(injection.values == values)
         }
 
@@ -128,8 +128,8 @@ struct InMemoryTracerTests {
 
             #expect(values.isEmpty)
 
-            let injection = try #require(tracer.injections.first)
-            #expect(injection.context.testSpanContext == nil)
+            let injection = try #require(tracer.performedContextInjections.first)
+            #expect(injection.context.inMemorySpanContext == nil)
             #expect(injection.values.isEmpty)
         }
 
@@ -141,13 +141,13 @@ struct InMemoryTracerTests {
             let values = [InMemoryTracer.traceIDKey: "stub", InMemoryTracer.spanIDKey: "stub"]
             tracer.extract(values, into: &context, using: DictionaryExtractor())
 
-            let spanContext = try #require(context.testSpanContext)
+            let spanContext = try #require(context.inMemorySpanContext)
 
             #expect(spanContext == InMemorySpanContext(traceID: "stub", spanID: "stub", parentSpanID: nil))
 
-            let extraction = try #require(tracer.extractions.first)
+            let extraction = try #require(tracer.performedContextExtractions.first)
             #expect(extraction.carrier as? [String: String] == values)
-            #expect(extraction.context.testSpanContext == spanContext)
+            #expect(extraction.context.inMemorySpanContext == spanContext)
         }
 
         @Test("Does not extract span context without values but records extraction")
@@ -158,11 +158,11 @@ struct InMemoryTracerTests {
             let values = ["foo": "bar"]
             tracer.extract(values, into: &context, using: DictionaryExtractor())
 
-            #expect(context.testSpanContext == nil)
+            #expect(context.inMemorySpanContext == nil)
 
-            let extraction = try #require(tracer.extractions.first)
+            let extraction = try #require(tracer.performedContextExtractions.first)
             #expect(extraction.carrier as? [String: String] == values)
-            #expect(extraction.context.testSpanContext == nil)
+            #expect(extraction.context.inMemorySpanContext == nil)
         }
     }
 
@@ -212,18 +212,18 @@ struct InMemoryTracerTests {
 
             let spanContext1 = InMemorySpanContext(traceID: "1", spanID: "1", parentSpanID: nil)
             var context1 = ServiceContext.topLevel
-            context1.testSpanContext = spanContext1
+            context1.inMemorySpanContext = spanContext1
             span.addLink(SpanLink(context: context1, attributes: ["foo": "1"]))
             let link1 = try #require(span.links.first)
-            #expect(link1.context.testSpanContext == spanContext1)
+            #expect(link1.context.inMemorySpanContext == spanContext1)
             #expect(link1.attributes == ["foo": "1"])
 
             let spanContext2 = InMemorySpanContext(traceID: "2", spanID: "2", parentSpanID: nil)
             var context2 = ServiceContext.topLevel
-            context2.testSpanContext = spanContext2
+            context2.inMemorySpanContext = spanContext2
             span.addLink(SpanLink(context: context2, attributes: ["foo": "2"]))
             let link2 = try #require(span.links.last)
-            #expect(link2.context.testSpanContext == spanContext2)
+            #expect(link2.context.inMemorySpanContext == spanContext2)
             #expect(link2.attributes == ["foo": "2"])
         }
 
@@ -282,7 +282,7 @@ struct InMemoryTracerTests {
             span.addEvent("foo")
             let otherSpanContext = InMemorySpanContext(traceID: "other", spanID: "other", parentSpanID: nil)
             var otherContext = ServiceContext.topLevel
-            otherContext.testSpanContext = otherSpanContext
+            otherContext.inMemorySpanContext = otherSpanContext
             span.addLink(SpanLink(context: otherContext, attributes: [:]))
             struct TestError: Error {}
             span.recordError(TestError())
