Add tracing support (#921)

* Add tracing interceptor

* PR feedback: Remove bitdrift trace id header

* iOS test fixes

* PR feedback (JNI cleanups, default to w3c)

Author: FranAguilera

CHANGELOG.md

@@ -7,7 +7,7 @@
 
 **Added**
 
-- Nothing yet!
+- Added distributed tracing support for network instrumentation with W3C (`traceparent`), B3 Single (`b3`), and B3 Multi (`X-B3-*`) propagation formats.
 
 **Changed**
 

examples/swift/hello_world/LoggerCustomer.swift

@@ -108,7 +108,7 @@ final class LoggerCustomer: NSObject, URLSessionDelegate {
                     requestFieldProvider: CustomNetworkFieldProvider(),
                     responseFieldProvider: CustomNetworkResponseFieldProvider()
                 ), ],
-                disableSwizzling: true
+                disableSwizzling: false
             )
 
         Logger.addField(withKey: "field_container_field_key", value: "field_container_value")

platform/jvm/capture-plugin/src/main/kotlin/io/bitdrift/capture/instrumentation/okhttp/visitor/OkHttpEventListenerMethodVisitor.kt

@@ -53,6 +53,8 @@ class OkHttpEventListenerMethodVisitor(
     ) {
     private val captureOkHttpEventListenerFactory =
         "io/bitdrift/capture/network/okhttp/CaptureOkHttpEventListenerFactory"
+    private val captureOkHttpTracingInterceptor =
+        "io/bitdrift/capture/network/okhttp/CaptureOkHttpTracingInterceptor"
 
     override fun onMethodEnter() {
         super.onMethodEnter()
@@ -64,6 +66,8 @@ class OkHttpEventListenerMethodVisitor(
     }
 
     private fun addOverwritingEventListener() {
+        addTracingInterceptor()
+
         // Add the following call at the beginning of the constructor with the Builder parameter:
         // builder.eventListenerFactory(new CaptureOkHttpEventListenerFactory());
 
@@ -95,9 +99,12 @@ class OkHttpEventListenerMethodVisitor(
             "(Lokhttp3/EventListener\$Factory;)Lokhttp3/OkHttpClient\$Builder;",
             false,
         )
+        visitInsn(Opcodes.POP)
     }
 
     private fun addProxyingEventListener() {
+        addTracingInterceptor()
+
         // Add the following call at the beginning of the constructor with the Builder parameter:
         // builder.eventListenerFactory(new CaptureOkHttpEventListenerFactory(builder.eventListenerFactory));
 
@@ -141,5 +148,35 @@ class OkHttpEventListenerMethodVisitor(
             "(Lokhttp3/EventListener\$Factory;)Lokhttp3/OkHttpClient\$Builder;",
             false,
         )
+        visitInsn(Opcodes.POP)
+    }
+
+    private fun addTracingInterceptor() {
+        // Add the following call at the beginning of the constructor with the Builder parameter:
+        // builder.addInterceptor(new CaptureOkHttpTracingInterceptor());
+
+        // OkHttpClient.Builder is the parameter, retrieved here
+        visitVarInsn(Opcodes.ALOAD, 1)
+
+        // Create CaptureOkHttpTracingInterceptor instance
+        visitTypeInsn(Opcodes.NEW, captureOkHttpTracingInterceptor)
+        visitInsn(Opcodes.DUP)
+        visitMethodInsn(
+            Opcodes.INVOKESPECIAL,
+            captureOkHttpTracingInterceptor,
+            "<init>",
+            "()V",
+            false,
+        )
+
+        // Call addInterceptor(Interceptor)
+        visitMethodInsn(
+            Opcodes.INVOKEVIRTUAL,
+            "okhttp3/OkHttpClient\$Builder",
+            "addInterceptor",
+            "(Lokhttp3/Interceptor;)Lokhttp3/OkHttpClient\$Builder;",
+            false,
+        )
+        visitInsn(Opcodes.POP)
     }
 }

platform/jvm/capture/src/main/kotlin/io/bitdrift/capture/Capture.kt

@@ -226,6 +226,14 @@ object Capture {
         val deviceId: String?
             get() = logger()?.deviceId
 
+        /**
+         * Whether workflow-controlled tracing is currently active for this session.
+         * Returns `null` prior to SDK start.
+         */
+        @JvmStatic
+        val isTracingActive: Boolean?
+            get() = logger()?.isTracingActive
+
         /**
          * Defines the initialization of a new session within the currently running logger
          * If no logger is started, this is a no-op.

platform/jvm/capture/src/main/kotlin/io/bitdrift/capture/CaptureJniLibrary.kt

@@ -111,6 +111,11 @@ internal object CaptureJniLibrary : IBridge, IStreamingReportProcessor {
      */
     external fun getDeviceId(loggerId: Long): String?
 
+    /**
+     * Returns true when workflow-controlled tracing is active for the current session.
+     */
+    external fun isTracingActive(loggerId: Long): Boolean
+
     /**
      * Adds a field that should be attached to all logs emitted by the logger going forward.
      * If a field with a given key has already been registered with the logger, its value is

platform/jvm/capture/src/main/kotlin/io/bitdrift/capture/ILogger.kt

@@ -37,6 +37,11 @@ interface ILogger {
      */
     val deviceId: String
 
+    /**
+     * Whether workflow-controlled tracing is currently active for this session.
+     */
+    val isTracingActive: Boolean
+
     /**
      * Defines the initialization of a new session within the current logger.
      */

platform/jvm/capture/src/main/kotlin/io/bitdrift/capture/IRuntimeProvider.kt

@@ -0,0 +1,33 @@
+// capture-sdk - bitdrift's client SDK
+// Copyright Bitdrift, Inc. All rights reserved.
+//
+// Use of this source code is governed by a source available license that can be found in the
+// LICENSE file or at:
+// https://polyformproject.org/wp-content/uploads/2020/06/PolyForm-Shield-1.0.0.txt
+
+package io.bitdrift.capture
+
+import io.bitdrift.capture.common.RuntimeConfig
+import io.bitdrift.capture.common.RuntimeFeature
+import io.bitdrift.capture.common.RuntimeStringConfig
+
+internal interface IRuntimeProvider {
+    fun isRuntimeFeatureEnabled(feature: RuntimeFeature): Boolean
+
+    fun getRuntimeConfigValue(config: RuntimeConfig): Int
+
+    fun getRuntimeStringConfigValue(config: RuntimeStringConfig): String
+}
+
+internal object CaptureRuntimeProvider : IRuntimeProvider {
+    override fun isRuntimeFeatureEnabled(feature: RuntimeFeature): Boolean =
+        getLoggerProvider()?.isRuntimeFeatureEnabled(feature) ?: feature.defaultValue
+
+    override fun getRuntimeConfigValue(config: RuntimeConfig): Int =
+        getLoggerProvider()?.getRuntimeConfigValue(config) ?: config.defaultValue
+
+    override fun getRuntimeStringConfigValue(config: RuntimeStringConfig): String =
+        getLoggerProvider()?.getRuntimeStringConfigValue(config) ?: config.defaultValue
+
+    private fun getLoggerProvider(): IRuntimeProvider? = Capture.logger() as? IRuntimeProvider
+}

platform/jvm/capture/src/main/kotlin/io/bitdrift/capture/JniRuntime.kt

@@ -10,13 +10,17 @@ package io.bitdrift.capture
 import io.bitdrift.capture.common.Runtime
 import io.bitdrift.capture.common.RuntimeConfig
 import io.bitdrift.capture.common.RuntimeFeature
+import io.bitdrift.capture.common.RuntimeStringConfig
 
 internal class JniRuntime(
     private val logger: LoggerId,
 ) : Runtime {
     override fun isEnabled(feature: RuntimeFeature): Boolean = Jni.isRuntimeEnabled(logger, feature.featureName, feature.defaultValue)
 
     override fun getConfigValue(config: RuntimeConfig): Int = Jni.runtimeValue(logger, config.configName, config.defaultValue)
+
+    override fun getConfigValue(config: RuntimeStringConfig): String =
+        Jni.runtimeStringValue(logger, config.configName, config.defaultValue)
 }
 
 internal object Jni {
@@ -31,4 +35,10 @@ internal object Jni {
         variableName: String,
         defaultValue: Int,
     ): Int
+
+    external fun runtimeStringValue(
+        logger: LoggerId,
+        variableName: String,
+        defaultValue: String,
+    ): String
 }

platform/jvm/capture/src/main/kotlin/io/bitdrift/capture/LoggerImpl.kt

@@ -16,7 +16,9 @@ import androidx.lifecycle.ProcessLifecycleOwner
 import io.bitdrift.capture.attributes.ClientAttributes
 import io.bitdrift.capture.attributes.NetworkAttributes
 import io.bitdrift.capture.common.IWindowManager
+import io.bitdrift.capture.common.RuntimeConfig
 import io.bitdrift.capture.common.RuntimeFeature
+import io.bitdrift.capture.common.RuntimeStringConfig
 import io.bitdrift.capture.common.WindowManager
 import io.bitdrift.capture.error.ErrorReporterService
 import io.bitdrift.capture.error.IErrorReporter
@@ -94,7 +96,8 @@ internal class LoggerImpl(
     private val eventListenerDispatcher: CaptureDispatchers.CommonBackground = CaptureDispatchers.CommonBackground,
     windowManager: IWindowManager = WindowManager(errorHandler),
 ) : IInternalLogger,
-    ICompletedReportsProcessor {
+    ICompletedReportsProcessor,
+    IRuntimeProvider {
     @OptIn(ExperimentalBitdriftApi::class)
     internal val webViewConfiguration: WebViewConfiguration? = configuration.webViewConfiguration
 
@@ -315,6 +318,9 @@ internal class LoggerImpl(
     override val deviceId: String
         get() = CaptureJniLibrary.getDeviceId(this.loggerId) ?: "unknown"
 
+    override val isTracingActive: Boolean
+        get() = CaptureJniLibrary.isTracingActive(this.loggerId)
+
     override fun startNewSession() {
         CaptureJniLibrary.startNewSession(this.loggerId)
     }
@@ -542,6 +548,12 @@ internal class LoggerImpl(
         )
     }
 
+    override fun isRuntimeFeatureEnabled(feature: RuntimeFeature): Boolean = runtime.isEnabled(feature)
+
+    override fun getRuntimeConfigValue(config: RuntimeConfig): Int = runtime.getConfigValue(config)
+
+    override fun getRuntimeStringConfigValue(config: RuntimeStringConfig): String = runtime.getConfigValue(config)
+
     override fun logSessionReplayScreenshot(
         fields: Array<Field>,
         duration: Duration,

platform/jvm/capture/src/main/kotlin/io/bitdrift/capture/network/okhttp/CaptureOkHttpEventListener.kt

@@ -102,7 +102,11 @@ internal class CaptureOkHttpEventListener internal constructor(
         callStartTimeMs = clock.elapsedRealtime()
 
         val request = call.request()
-        val extraFields = requestExtraFieldsProvider.provideExtraFields(request)
+        val extraFields =
+            buildExtraFieldsWithOptionalTraceId(
+                request,
+                requestExtraFieldsProvider.provideExtraFields(request),
+            )
 
         val pathTemplateHeaderValues = request.headers.values("x-capture-path-template")
         val pathTemplate =
@@ -356,6 +360,12 @@ internal class CaptureOkHttpEventListener internal constructor(
         // successful or not to keep iOS and Android implementation in sync.
         val isSuccess = (statusCode in 200..<400)
 
+        val extraFields =
+            buildExtraFieldsWithOptionalTraceId(
+                request,
+                responseExtraFieldsProvider.provideExtraFields(response),
+            )
+
         // Capture response URL attributes in case there was a redirect and attributes such as host,
         // path, and query have different values for the original request and the response.
         // https://square.github.io/okhttp/features/interceptors/#application-interceptors
@@ -380,7 +390,7 @@ internal class CaptureOkHttpEventListener internal constructor(
                 response = httpResponse,
                 durationMs = (clock.elapsedRealtime() - callStartTimeMs),
                 metrics = getMetrics(),
-                extraFields = responseExtraFieldsProvider.provideExtraFields(response),
+                extraFields = extraFields,
             )
 
         logger?.log(httpResponseInfo)
@@ -398,6 +408,8 @@ internal class CaptureOkHttpEventListener internal constructor(
         // them i.e., redirected request.
         val request = lastResponse?.request ?: call.request()
 
+        val extraFields = buildExtraFieldsWithOptionalTraceId(request)
+
         // Capture response URL attributes in case there was a redirect and attributes such as host,
         // path, and query have different values for the original request and the response.
         // https://square.github.io/okhttp/features/interceptors/#application-interceptors
@@ -421,6 +433,7 @@ internal class CaptureOkHttpEventListener internal constructor(
                 response = httpResponse,
                 durationMs = (clock.elapsedRealtime() - callStartTimeMs),
                 metrics = getMetrics(),
+                extraFields = extraFields,
             )
         logger?.log(httpResponseInfo)
     }
@@ -497,6 +510,16 @@ internal class CaptureOkHttpEventListener internal constructor(
         }
     }
 
+    private fun buildExtraFieldsWithOptionalTraceId(
+        request: Request,
+        baseFields: Map<String, String> = emptyMap(),
+    ): Map<String, String> {
+        val traceId = TracePropagation.extractSampledTraceId(request) ?: return baseFields
+        val fields = baseFields.toMutableMap()
+        fields[TracePropagation.TRACE_ID_FIELD_KEY] = traceId
+        return fields
+    }
+
     private companion object {
         const val TRAFFIC_STATS_TAG = 0x0017D21F7
     }