[chain-pr 6/10] DatadogTrace: typed-bus migration

Author: maxep

DatadogCore/Tests/Datadog/FeaturesIntegration/TracingWithLoggingIntegrationTests.swift

@@ -17,7 +17,8 @@ class TracingWithLoggingIntegrationTests: XCTestCase {
 
     override func setUp() {
         super.setUp()
-        core = PassthroughCoreMock(messageReceiver: LogMessageReceiver.mockAny())
+        core = PassthroughCoreMock()
+        core.subscribe(receiver: LogMessageReceiver.mockAny())
     }
 
     override func tearDown() {

DatadogCore/Tests/Datadog/TracerTests.swift

@@ -618,10 +618,9 @@ class TracerTests: XCTestCase {
     // MARK: - Integration With Logging Feature
 
     func testSendingSpanLogs() throws {
-        let logging: LogsFeature = .mockWith(
-            messageReceiver: LogMessageReceiver.mockAny()
-        )
+        let logging: LogsFeature = .mockWith()
         try core.register(feature: logging)
+        core.messageBus.subscribe(receiver: logging.logMessageReceiver)
 
         Trace.enable(with: config, in: core)
         let tracer = Tracer.shared(in: core)
@@ -651,10 +650,9 @@ class TracerTests: XCTestCase {
     }
 
     func testSendingSpanLogsWithErrorFromArguments() throws {
-        let logging: LogsFeature = .mockWith(
-            messageReceiver: LogMessageReceiver.mockAny()
-        )
+        let logging: LogsFeature = .mockWith()
         try core.register(feature: logging)
+        core.messageBus.subscribe(receiver: logging.logMessageReceiver)
 
         Trace.enable(with: config, in: core)
         let tracer = Tracer.shared(in: core)
@@ -676,10 +674,9 @@ class TracerTests: XCTestCase {
     }
 
     func testSendingSpanLogsWithErrorFromNSError() throws {
-        let logging: LogsFeature = .mockWith(
-            messageReceiver: LogMessageReceiver.mockAny()
-        )
+        let logging: LogsFeature = .mockWith()
         try core.register(feature: logging)
+        core.messageBus.subscribe(receiver: logging.logMessageReceiver)
 
         Trace.enable(with: config, in: core)
         let tracer = Tracer.shared(in: core)
@@ -707,10 +704,9 @@ class TracerTests: XCTestCase {
     }
 
     func testSendingSpanLogsWithErrorFromSwiftError() throws {
-        let logging: LogsFeature = .mockWith(
-            messageReceiver: LogMessageReceiver.mockAny()
-        )
+        let logging: LogsFeature = .mockWith()
         try core.register(feature: logging)
+        core.messageBus.subscribe(receiver: logging.logMessageReceiver)
 
         Trace.enable(with: config, in: core)
         let tracer = Tracer.shared(in: core)

DatadogCore/Tests/Datadog/Tracing/TracingURLSessionHandlerTests.swift

@@ -21,10 +21,9 @@ class TracingURLSessionHandlerTests: XCTestCase {
     override func setUp() {
         super.setUp()
         let receiver = ContextMessageReceiver(samplerProvider: SamplerProvider(sampleRate: .mockAny()))
-        core = PassthroughCoreMock(messageReceiver: CombinedFeatureMessageReceiver([
-            LogMessageReceiver.mockAny(),
-            receiver
-        ]))
+        core = PassthroughCoreMock()
+        core.subscribe(receiver: receiver)
+        core.subscribe(receiver: LogMessageReceiver.mockAny())
 
         tracer = .mockWith(
             core: core,
@@ -225,8 +224,7 @@ class TracingURLSessionHandlerTests: XCTestCase {
                 )
             ]
         )
-        let message = FeatureMessage.context(fakeContext)
-        _ = handler.contextReceiver.receive(message: message, from: core)
+        handler.contextReceiver.receive(message: fakeContext, from: core)
         let (modifiedRequest, _, _) = handler.modify(
             request: request,
             headerTypes: [.datadog, .tracecontext, .b3, .b3multi],

DatadogTrace/Sources/Feature/MessageReceivers.swift

@@ -21,7 +21,7 @@ internal struct CoreContext {
     var accountInfo: AccountInfo?
 }
 
-internal final class ContextMessageReceiver: FeatureMessageReceiver {
+internal final class ContextMessageReceiver: BusMessageReceiver {
     /// Creates a new `ContextMessageReceiver`.
     ///
     /// - parameters:
@@ -41,24 +41,7 @@ internal final class ContextMessageReceiver: FeatureMessageReceiver {
     /// The tracer sampler that should be updated with the RUM deterministic sampler.
     let samplerProvider: SamplerProvider
 
-    /// Process messages receives from the bus.
-    ///
-    /// - Parameters:
-    ///   - message: The Feature message
-    ///   - core: The core from which the message is transmitted.
-    func receive(message: FeatureMessage, from core: DatadogCoreProtocol) -> Bool {
-        switch message {
-        case .context(let context):
-            return update(context: context, from: core)
-        default:
-            return false
-        }
-    }
-
-    /// Updates context of the `DatadogTracer` if available.
-    ///
-    /// - Parameter context: The updated core context.
-    private func update(context datadogContext: DatadogContext, from core: DatadogCoreProtocol) -> Bool {
+    func receive(message datadogContext: DatadogContext, from core: DatadogCoreProtocol) {
         let rumContext = datadogContext.additionalContext(ofType: RUMCoreContext.self)
 
         _context.mutate {
@@ -69,7 +52,5 @@ internal final class ContextMessageReceiver: FeatureMessageReceiver {
         }
 
         samplerProvider.updateWith(deterministicSampler: rumContext?.sessionSampler)
-
-        return true
     }
 }

DatadogTrace/Sources/Feature/TraceFeature.swift

@@ -11,7 +11,7 @@ internal final class TraceFeature: DatadogRemoteFeature {
     static let name = "tracing"
 
     let requestBuilder: FeatureRequestBuilder
-    var messageReceiver: FeatureMessageReceiver { contextReceiver }
+    let messageReceiver: FeatureMessageReceiver = NOPFeatureMessageReceiver()
 
     let tracer: DatadogTracer
     let contextReceiver: ContextMessageReceiver

DatadogTrace/Sources/Integrations/TracingWithLoggingIntegration.swift

@@ -61,23 +61,21 @@ internal struct TracingWithLoggingIntegration {
             )
         : nil
 
-        core.send(
-            message: .payload(
-                LogMessage(
-                    logger: "trace",
-                    service: service,
-                    date: date,
-                    message: message,
-                    error: extractedError,
-                    level: level,
-                    thread: Thread.current.dd.name,
-                    networkInfoEnabled: networkInfoEnabled,
-                    userAttributes: userAttributes,
-                    internalAttributes: [
-                        Constants.traceIDKey: String(spanContext.traceID, representation: .hexadecimal),
-                        Constants.spanIDKey: String(spanContext.spanID, representation: .hexadecimal)
-                    ]
-                )
+        core.messageBus.send(
+            message: LogMessage(
+                logger: "trace",
+                service: service,
+                date: date,
+                message: message,
+                error: extractedError,
+                level: level,
+                thread: Thread.current.dd.name,
+                networkInfoEnabled: networkInfoEnabled,
+                userAttributes: userAttributes,
+                internalAttributes: [
+                    Constants.traceIDKey: String(spanContext.traceID, representation: .hexadecimal),
+                    Constants.spanIDKey: String(spanContext.spanID, representation: .hexadecimal)
+                ]
             ),
             else: fallback
         )

DatadogTrace/Sources/Trace.swift

@@ -42,6 +42,10 @@ public enum Trace {
 
         // Register Trace feature:
         let trace = TraceFeature(in: core, configuration: configuration)
+
+        // Subscribe typed-bus receivers before registration so initial context push is received:
+        core.messageBus.subscribe(receiver: trace.contextReceiver)
+
         try core.register(feature: trace)
 
         // If `URLSession` tracking is configured, register `URLSessionHandler` to enable distributed tracing:

DatadogTrace/Tests/ContextMessageReceiverTests.swift

@@ -14,10 +14,8 @@ class ContextMessageReceiverTests: XCTestCase {
     func testItReceivesApplicationStateHistory() throws {
         // Given
         let receiver = ContextMessageReceiver(samplerProvider: SamplerProvider(sampleRate: .mockAny()))
-        let core = PassthroughCoreMock(
-            context: .mockWith(applicationStateHistory: .mockAppInBackground()),
-            messageReceiver: receiver
-        )
+        let core = PassthroughCoreMock(context: .mockWith(applicationStateHistory: .mockAppInBackground()))
+        core.subscribe(receiver: receiver)
 
         XCTAssertEqual(receiver.context.applicationStateHistory?.currentState, .background)
 

DatadogTrace/Tests/TracingURLSessionHandlerTests.swift

@@ -20,7 +20,8 @@ class TracingURLSessionHandlerTests: XCTestCase {
     override func setUp() {
         super.setUp()
         let receiver = ContextMessageReceiver(samplerProvider: SamplerProvider(sampleRate: .mockAny()))
-        core = PassthroughCoreMock(messageReceiver: receiver)
+        core = PassthroughCoreMock()
+        core.subscribe(receiver: receiver)
 
         tracer = .mockWith(
             core: core,