RUM-17619: Fix deadlock/ANR in DatadogRumMonitor.handleEvent on JVM crash

dd-sdk-android-core/api/apiSurface

@@ -130,7 +130,7 @@ interface com.datadog.android.api.feature.FeatureScope
   val dataStore: com.datadog.android.api.storage.datastore.DataStoreHandler
   fun withWriteContext(Set<String> = emptySet(), (com.datadog.android.api.context.DatadogContext) -> Unit)
   fun withContext(Set<String> = emptySet(), (com.datadog.android.api.context.DatadogContext) -> Unit)
-  fun getWriteContextSync(Set<String> = emptySet()): Pair<com.datadog.android.api.context.DatadogContext, EventWriteScope>?
+  fun withWriteContextSync(Set<String> = emptySet(), (com.datadog.android.api.context.DatadogContext) -> Unit): Boolean
   fun sendEvent(Any)
   fun <T: Feature> unwrap(): T
 typealias EventWriteScope = ((com.datadog.android.api.storage.EventBatchWriter) -> Unit) -> Unit

dd-sdk-android-core/api/dd-sdk-android-core.api

@@ -414,17 +414,17 @@ public abstract interface class com/datadog/android/api/feature/FeatureEventRece
 
 public abstract interface class com/datadog/android/api/feature/FeatureScope {
 	public abstract fun getDataStore ()Lcom/datadog/android/api/storage/datastore/DataStoreHandler;
-	public abstract fun getWriteContextSync (Ljava/util/Set;)Lkotlin/Pair;
 	public abstract fun sendEvent (Ljava/lang/Object;)V
 	public abstract fun unwrap ()Lcom/datadog/android/api/feature/Feature;
 	public abstract fun withContext (Ljava/util/Set;Lkotlin/jvm/functions/Function1;)V
 	public abstract fun withWriteContext (Ljava/util/Set;Lkotlin/jvm/functions/Function2;)V
+	public abstract fun withWriteContextSync (Ljava/util/Set;Lkotlin/jvm/functions/Function2;)Z
 }
 
 public final class com/datadog/android/api/feature/FeatureScope$DefaultImpls {
-	public static synthetic fun getWriteContextSync$default (Lcom/datadog/android/api/feature/FeatureScope;Ljava/util/Set;ILjava/lang/Object;)Lkotlin/Pair;
 	public static synthetic fun withContext$default (Lcom/datadog/android/api/feature/FeatureScope;Ljava/util/Set;Lkotlin/jvm/functions/Function1;ILjava/lang/Object;)V
 	public static synthetic fun withWriteContext$default (Lcom/datadog/android/api/feature/FeatureScope;Ljava/util/Set;Lkotlin/jvm/functions/Function2;ILjava/lang/Object;)V
+	public static synthetic fun withWriteContextSync$default (Lcom/datadog/android/api/feature/FeatureScope;Ljava/util/Set;Lkotlin/jvm/functions/Function2;ILjava/lang/Object;)Z
 }
 
 public final class com/datadog/android/api/feature/FeatureScopeExtKt {

dd-sdk-android-core/src/main/kotlin/com/datadog/android/api/feature/FeatureScope.kt

@@ -52,19 +52,31 @@ interface FeatureScope {
         callback: (datadogContext: DatadogContext) -> Unit
     )
 
-    // TODO RUM-9852 Implement better passthrough mechanism for the JVM crash scenario
     /**
-     * Same as [withWriteContext] but will be executed in the blocking manner.
+     * Same as [withWriteContext], but blocks the calling thread until [callback] has returned.
+     *
+     * The callback still runs on the context processing worker thread, so the calling thread must not
+     * hold any lock that [callback] may need, otherwise the two will deadlock.
      *
      * @param withFeatureContexts Feature contexts ([DatadogContext.featuresContext] property) to include
      * in the [DatadogContext] provided. The value should be the feature names as declared by [Feature.name].
      * Default is empty, meaning that no feature contexts will be included.
+     * @param callback an operation called with an up-to-date [DatadogContext]
+     * and an [EventWriteScope]. Callback will be executed on a single context processing worker thread. Execution of
+     * [EventWriteScope] will be done on a worker thread from I/O pool.
+     * [DatadogContext] is a snapshot taken when [callback] starts executing, which is once every context
+     * operation scheduled before this call has completed.
+     * @return `true` if [callback] was executed, `false` if it could not be, for example because the
+     * SDK core is not initialized or the operation could not be scheduled.
      *
      * **NOTE**: This API is for the internal use only and is not guaranteed to be stable.
      */
     @AnyThread
     @InternalApi
-    fun getWriteContextSync(withFeatureContexts: Set<String> = emptySet()): Pair<DatadogContext, EventWriteScope>?
+    fun withWriteContextSync(
+        withFeatureContexts: Set<String> = emptySet(),
+        callback: (datadogContext: DatadogContext, write: EventWriteScope) -> Unit
+    ): Boolean
 
     /**
      * Send event to a given feature. It will be sent in a synchronous way.

dd-sdk-android-core/src/main/kotlin/com/datadog/android/core/internal/SdkFeature.kt

@@ -215,22 +215,24 @@ internal class SdkFeature(
         )
     }
 
-    override fun getWriteContextSync(
-        withFeatureContexts: Set<String>
-    ): Pair<DatadogContext, EventWriteScope>? {
-        val operationName = "getWriteContextSync-${wrappedFeature.name}"
+    override fun withWriteContextSync(
+        withFeatureContexts: Set<String>,
+        callback: (DatadogContext, EventWriteScope) -> Unit
+    ): Boolean {
+        val operationName = "withWriteContextSync-${wrappedFeature.name}"
         return coreFeature.contextExecutorService
             .submitSafe(
                 operationName,
                 internalLogger,
                 Callable {
-                    if (!coreFeature.initialized.get()) return@Callable null
+                    if (!coreFeature.initialized.get()) return@Callable false
                     val context = contextProvider.getContext(withFeatureContexts)
                     val eventBatchWriteScope = storage.getEventWriteScope(context)
-                    context to eventBatchWriteScope
+                    callback(context, eventBatchWriteScope)
+                    true
                 }
             )
-            .getSafe(operationName, internalLogger)
+            .getSafe(operationName, internalLogger) ?: false
     }
 
     override fun sendEvent(event: Any) {

dd-sdk-android-core/src/test/kotlin/com/datadog/android/core/internal/SdkFeatureTest.kt

@@ -565,18 +565,19 @@ internal class SdkFeatureTest {
     }
 
     @Test
-    fun `M provide write context W getWriteContextSync()`(
+    fun `M provide write context W withWriteContextSync(callback)`(
         @Forgery fakeContext: DatadogContext,
         @StringForgery fakeWithFeatureContexts: Set<String>,
         @Mock mockEventWriteScope: EventWriteScope
     ) {
         // Given
         testedFeature.storage = mockStorage
+        val callback = mock<(DatadogContext, EventWriteScope) -> Unit>()
         whenever(mockContextProvider.getContext(fakeWithFeatureContexts)) doReturn fakeContext
-        whenever(coreFeature.mockInstance.contextExecutorService.submit(any<Callable<*>>())) doAnswer {
-            val callable = it.getArgument<Callable<Pair<DatadogContext, EventWriteScope>>>(0)
-            mock<Future<*>>().apply {
-                whenever(get()) doAnswer { callable.call() }
+        whenever(coreFeature.mockInstance.contextExecutorService.submit(any<Callable<Boolean>>())) doAnswer {
+            val result = it.getArgument<Callable<Boolean>>(0).call()
+            mock<Future<Boolean>>().apply {
+                whenever(get()) doReturn result
             }
         }
 
@@ -585,77 +586,107 @@ internal class SdkFeatureTest {
         ) doReturn mockEventWriteScope
 
         // When
-        val writeContext = testedFeature.getWriteContextSync(fakeWithFeatureContexts)
+        val result = testedFeature.withWriteContextSync(fakeWithFeatureContexts, callback = callback)
 
         // Then
-        checkNotNull(writeContext)
-        assertThat(writeContext.first).isEqualTo(fakeContext)
-        assertThat(writeContext.second).isEqualTo(mockEventWriteScope)
+        verify(callback).invoke(fakeContext, mockEventWriteScope)
+        assertThat(result).isTrue()
     }
 
     @Test
-    fun `M provide null write context W getWriteContextSync() { task rejected }`(
+    fun `M wait for the callback W withWriteContextSync(callback)`(
         @Forgery fakeContext: DatadogContext,
+        @StringForgery fakeWithFeatureContexts: Set<String>,
         @Mock mockEventWriteScope: EventWriteScope
     ) {
         // Given
         testedFeature.storage = mockStorage
-        whenever(
-            coreFeature.mockInstance.contextExecutorService.submit(any<Callable<*>>())
-        ) doThrow RejectedExecutionException()
+        whenever(mockContextProvider.getContext(fakeWithFeatureContexts)) doReturn fakeContext
+        whenever(mockStorage.getEventWriteScope(fakeContext)) doReturn mockEventWriteScope
+        // the submitted task only runs when the future is awaited, so the callback can only have been
+        // invoked if withWriteContextSync blocked on it
+        whenever(coreFeature.mockInstance.contextExecutorService.submit(any<Callable<Boolean>>())) doAnswer {
+            val submittedTask = it.getArgument<Callable<Boolean>>(0)
+            mock<Future<Boolean>>().apply {
+                whenever(get()) doAnswer { submittedTask.call() }
+            }
+        }
+        var callbackInvoked = false
+
+        // When
+        testedFeature.withWriteContextSync(fakeWithFeatureContexts) { _, _ -> callbackInvoked = true }
 
+        // Then
+        assertThat(callbackInvoked).isTrue()
+    }
+
+    @Test
+    fun `M not provide write context W withWriteContextSync(callback) { task rejected }`(
+        @StringForgery fakeWithFeatureContexts: Set<String>
+    ) {
+        // Given
+        testedFeature.storage = mockStorage
+        val callback = mock<(DatadogContext, EventWriteScope) -> Unit>()
         whenever(
-            mockStorage.getEventWriteScope(fakeContext)
-        ) doReturn mockEventWriteScope
+            coreFeature.mockInstance.contextExecutorService.submit(any<Callable<Boolean>>())
+        ) doThrow RejectedExecutionException()
 
         // When
-        val writeContext = testedFeature.getWriteContextSync()
+        val result = testedFeature.withWriteContextSync(fakeWithFeatureContexts, callback = callback)
 
         // Then
-        assertThat(writeContext).isNull()
+        verifyNoInteractions(callback, mockContextProvider, mockStorage)
+        assertThat(result).isFalse()
     }
 
     @Test
-    fun `M provide null write context W getWriteContextSync() { failed to get task result }`(
-        @Forgery fakeContext: DatadogContext,
-        @Mock mockEventWriteScope: EventWriteScope,
+    fun `M not throw W withWriteContextSync(callback) { failed to get task result }`(
+        @StringForgery fakeWithFeatureContexts: Set<String>,
         forge: Forge
     ) {
         // Given
         testedFeature.storage = mockStorage
+        val callback = mock<(DatadogContext, EventWriteScope) -> Unit>()
         val throwable = forge.anElementFrom(
             CancellationException(),
             ExecutionException(forge.aThrowable()),
             InterruptedException()
         )
-        whenever(coreFeature.mockInstance.contextExecutorService.submit(any<Callable<*>>())) doAnswer {
-            mock<Future<*>>().apply {
+        whenever(coreFeature.mockInstance.contextExecutorService.submit(any<Callable<Boolean>>())) doAnswer {
+            mock<Future<Boolean>>().apply {
                 whenever(get()) doThrow throwable
             }
         }
 
-        whenever(
-            mockStorage.getEventWriteScope(fakeContext)
-        ) doReturn mockEventWriteScope
-
         // When
-        val writeContext = testedFeature.getWriteContextSync()
+        val result = testedFeature.withWriteContextSync(fakeWithFeatureContexts, callback = callback)
 
         // Then
-        assertThat(writeContext).isNull()
+        verifyNoInteractions(callback)
+        assertThat(result).isFalse()
     }
 
     @Test
-    fun `M provide null write context W getWriteContextSync() { CoreFeature is not initialized }`() {
+    fun `M not provide write context W withWriteContextSync(callback) { CoreFeature is not initialized }`(
+        @StringForgery fakeWithFeatureContexts: Set<String>
+    ) {
         // Given
+        testedFeature.storage = mockStorage
+        val callback = mock<(DatadogContext, EventWriteScope) -> Unit>()
         whenever(coreFeature.mockInstance.initialized) doReturn AtomicBoolean(false)
+        whenever(coreFeature.mockInstance.contextExecutorService.submit(any<Callable<Boolean>>())) doAnswer {
+            val taskResult = it.getArgument<Callable<Boolean>>(0).call()
+            mock<Future<Boolean>>().apply {
+                whenever(get()) doReturn taskResult
+            }
+        }
 
         // When
-        val writeContext = testedFeature.getWriteContextSync()
+        val result = testedFeature.withWriteContextSync(fakeWithFeatureContexts, callback = callback)
 
         // Then
-        assertThat(writeContext).isNull()
-        verifyNoInteractions(mockContextProvider, mockStorage)
+        verifyNoInteractions(callback, mockContextProvider, mockStorage)
+        assertThat(result).isFalse()
     }
 
     @Test

features/dd-sdk-android-rum/src/main/kotlin/com/datadog/android/rum/internal/monitor/DatadogRumMonitor.kt

@@ -919,26 +919,26 @@ internal class DatadogRumMonitor(
 
     internal fun handleEvent(event: RumRawEvent) {
         if (event is RumRawEvent.AddError && event.isFatal) {
-            synchronized(rootScope) {
-                // TODO RUM-9852 Implement better passthrough mechanism for the JVM crash scenario
-                val writeContext = sdkCore.getFeature(Feature.RUM_FEATURE_NAME)
-                    ?.getWriteContextSync(withFeatureContexts = setOf(Feature.SESSION_REPLAY_FEATURE_NAME))
-                if (writeContext != null) {
-                    val (datadogContext, eventWriteScope) = writeContext
-                    @Suppress("ThreadSafety") // Crash handling, can't delegate to another thread
-                    rootScope.handleEvent(event, datadogContext, eventWriteScope, writer)
-                    val rumContext = currentRumContext()
-                    sdkCore.updateFeatureContext(Feature.RUM_FEATURE_NAME) {
-                        it.clear()
-                        rumContext?.toMap()?.let(it::putAll)
+            val handled = sdkCore.getFeature(Feature.RUM_FEATURE_NAME)
+                ?.withWriteContextSync(
+                    withFeatureContexts = setOf(Feature.SESSION_REPLAY_FEATURE_NAME)
+                ) { datadogContext, eventWriteScope ->
+                    synchronized(rootScope) {
+                        @Suppress("ThreadSafety") // Crash handling, can't delegate to another thread
+                        rootScope.handleEvent(event, datadogContext, eventWriteScope, writer)
+                        val rumContext = currentRumContext()
+                        sdkCore.updateFeatureContext(Feature.RUM_FEATURE_NAME, useContextThread = false) {
+                            it.clear()
+                            rumContext?.toMap()?.let(it::putAll)
+                        }
                     }
-                } else {
-                    sdkCore.internalLogger.log(
-                        InternalLogger.Level.WARN,
-                        InternalLogger.Target.USER,
-                        { CANNOT_WRITE_CRASH_WRITE_CONTEXT_IS_NOT_AVAILABLE }
-                    )
                 }
+            if (handled != true) {
+                sdkCore.internalLogger.log(
+                    InternalLogger.Level.WARN,
+                    InternalLogger.Target.USER,
+                    { CANNOT_WRITE_CRASH_WRITE_CONTEXT_IS_NOT_AVAILABLE }
+                )
             }
         } else if (event is RumRawEvent.TelemetryEventWrapper) {
             telemetryEventHandler.handleEvent(event, writer)

features/dd-sdk-android-rum/src/test/kotlin/com/datadog/android/rum/internal/domain/scope/RumViewManagerScopeTest.kt

@@ -162,6 +162,7 @@ internal class RumViewManagerScopeTest {
         fakeSampleRate = forge.aFloat(min = 0.0f, max = 100.0f)
 
         whenever(mockSdkCore.time) doReturn fakeTime
+        whenever(mockSdkCore.timeProvider) doReturn mock()
 
         whenever(mockParentScope.getRumContext()) doReturn fakeParentContext
         whenever(mockChildScope.handleEvent(any(), any(), any(), any())) doReturn mockChildScope