Speculative fix for runBlockingWithTracing NPE

There may have been a memory race in runBlockingWithTracing's
handling of the return variable, with the updated return value
not always visible upon return.

This patch refactors to use standard kotlin constructs which we
preume don't suffer from these defects.

Bug: 227758814
Test: This is speculative, was unable to reproduce the bug.
Change-Id: I070d93e16040654c4e117bc34848734cff0f6300

diff --git a/wear/watchface/watchface-client-guava/src/androidTest/java/androidx/wear/watchface/client/guava/ListenableWatchFaceMetadataClientTest.kt b/wear/watchface/watchface-client-guava/src/androidTest/java/androidx/wear/watchface/client/guava/ListenableWatchFaceMetadataClientTest.kt
index 2c071f6..98ce3be 100644
--- a/wear/watchface/watchface-client-guava/src/androidTest/java/androidx/wear/watchface/client/guava/ListenableWatchFaceMetadataClientTest.kt
+++ b/wear/watchface/watchface-client-guava/src/androidTest/java/androidx/wear/watchface/client/guava/ListenableWatchFaceMetadataClientTest.kt

@@ -22,9 +22,7 @@
 import android.content.Context
 import android.content.Intent
 import android.content.res.XmlResourceParser
-import android.os.Handler
 import android.os.IBinder
-import android.os.Looper
 import androidx.test.core.app.ApplicationProvider
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.MediumTest
@@ -36,6 +34,7 @@
 import org.junit.Test
 import org.junit.runner.RunWith
 import java.util.concurrent.TimeUnit
+import kotlinx.coroutines.MainScope
 
 private const val TIMEOUT_MS = 500L
 
@@ -47,7 +46,7 @@
     private val realService = object : WatchFaceControlService() {
         @SuppressLint("NewApi")
         override fun createServiceStub(): IWatchFaceInstanceServiceStub =
-            IWatchFaceInstanceServiceStub(this, Handler(Looper.getMainLooper()))
+            IWatchFaceInstanceServiceStub(this, MainScope())
 
         init {
             setContext(ApplicationProvider.getApplicationContext<Context>())

diff --git a/wear/watchface/watchface-client/src/androidTest/java/androidx/wear/watchface/client/test/WatchFaceControlClientTest.kt b/wear/watchface/watchface-client/src/androidTest/java/androidx/wear/watchface/client/test/WatchFaceControlClientTest.kt
index 1220dfe..0dd334c 100644
--- a/wear/watchface/watchface-client/src/androidTest/java/androidx/wear/watchface/client/test/WatchFaceControlClientTest.kt
+++ b/wear/watchface/watchface-client/src/androidTest/java/androidx/wear/watchface/client/test/WatchFaceControlClientTest.kt

@@ -751,14 +751,9 @@
             ]!!
         }
 
-        var isFirstCall = true
         handlerCoroutineScope.launch {
             leftComplicationSlot.complicationData.collect {
-                if (!isFirstCall) {
-                    updateCountDownLatch.countDown()
-                } else {
-                    isFirstCall = false
-                }
+                updateCountDownLatch.countDown()
             }
         }
 

diff --git a/wear/watchface/watchface-client/src/androidTest/java/androidx/wear/watchface/client/test/WatchFaceControlTestService.kt b/wear/watchface/watchface-client/src/androidTest/java/androidx/wear/watchface/client/test/WatchFaceControlTestService.kt
index 98b337e..621ae29 100644
--- a/wear/watchface/watchface-client/src/androidTest/java/androidx/wear/watchface/client/test/WatchFaceControlTestService.kt
+++ b/wear/watchface/watchface-client/src/androidTest/java/androidx/wear/watchface/client/test/WatchFaceControlTestService.kt

@@ -20,13 +20,12 @@
 import android.content.Context
 import android.content.Intent
 import android.os.Build
-import android.os.Handler
 import android.os.IBinder
-import android.os.Looper
 import androidx.annotation.RequiresApi
 import androidx.test.core.app.ApplicationProvider
 import androidx.wear.watchface.control.IWatchFaceInstanceServiceStub
 import androidx.wear.watchface.control.WatchFaceControlService
+import kotlinx.coroutines.MainScope
 
 /**
  * Test shim to allow us to connect to WatchFaceControlService from
@@ -42,9 +41,7 @@
 
     private val realService = object : WatchFaceControlService() {
         override fun createServiceStub(): IWatchFaceInstanceServiceStub =
-            object : IWatchFaceInstanceServiceStub(
-                this@WatchFaceControlTestService, Handler(Looper.getMainLooper())
-            ) {
+            object : IWatchFaceInstanceServiceStub(this@WatchFaceControlTestService, MainScope()) {
                 @RequiresApi(Build.VERSION_CODES.O_MR1)
                 override fun getApiVersion(): Int = apiVersionOverride ?: super.getApiVersion()
             }

diff --git a/wear/watchface/watchface-guava/src/androidTest/java/androidx/wear/watchface/WatchFaceControlTestService.kt b/wear/watchface/watchface-guava/src/androidTest/java/androidx/wear/watchface/WatchFaceControlTestService.kt
index 65312b5..9cdf97c 100644
--- a/wear/watchface/watchface-guava/src/androidTest/java/androidx/wear/watchface/WatchFaceControlTestService.kt
+++ b/wear/watchface/watchface-guava/src/androidTest/java/androidx/wear/watchface/WatchFaceControlTestService.kt

@@ -44,6 +44,7 @@
 import java.util.concurrent.CountDownLatch
 import java.util.concurrent.TimeUnit
 import java.util.concurrent.TimeoutException
+import kotlinx.coroutines.MainScope
 
 internal const val TIMEOUT_MILLIS = 1000L
 
@@ -64,7 +65,7 @@
         override fun createServiceStub(): IWatchFaceInstanceServiceStub =
             object : IWatchFaceInstanceServiceStub(
                 ApplicationProvider.getApplicationContext<Context>(),
-                Handler(Looper.getMainLooper())
+                MainScope()
             ) {
                 @RequiresApi(Build.VERSION_CODES.O_MR1)
                 override fun getApiVersion(): Int = apiVersionOverride ?: super.getApiVersion()

diff --git a/wear/watchface/watchface/src/main/java/androidx/wear/watchface/WatchFaceService.kt b/wear/watchface/watchface/src/main/java/androidx/wear/watchface/WatchFaceService.kt
index 8a4926a..cfc5e8b 100644
--- a/wear/watchface/watchface/src/main/java/androidx/wear/watchface/WatchFaceService.kt
+++ b/wear/watchface/watchface/src/main/java/androidx/wear/watchface/WatchFaceService.kt

@@ -84,7 +84,6 @@
 import kotlinx.coroutines.Runnable
 import kotlinx.coroutines.android.asCoroutineDispatcher
 import kotlinx.coroutines.cancel
-import kotlinx.coroutines.flow.collect
 import kotlinx.coroutines.launch
 import kotlinx.coroutines.runBlocking
 import java.io.ByteArrayInputStream
@@ -96,7 +95,6 @@
 import java.io.ObjectOutputStream
 import java.io.PrintWriter
 import java.time.Instant
-import java.util.concurrent.CountDownLatch
 import kotlinx.coroutines.CancellationException
 import kotlinx.coroutines.withContext
 
@@ -2399,50 +2397,6 @@
 }
 
 /**
- * Runs the supplied task on the handler thread. If we're not on the handler thread a task is posted
- * and we block until it's been processed.
- *
- * AIDL calls are dispatched from a thread pool, but for simplicity WatchFaceImpl code is largely
- * single threaded so we need to post tasks to the UI thread and wait for them to execute.
- *
- * @param traceEventName The name of the trace event to emit.
- * @param task The task to post on the handler.
- * @return [R] the return value of [task].
- */
-internal fun <R> Handler.runBlockingOnHandlerWithTracing(
-    traceEventName: String,
-    task: () -> R
-): R = TraceEvent(traceEventName).use {
-    if (looper == Looper.myLooper()) {
-        task.invoke()
-    } else {
-        val latch = CountDownLatch(1)
-        var returnVal: R? = null
-        var exception: Exception? = null
-        if (post {
-            try {
-                returnVal = TraceEvent("$traceEventName invokeTask").use {
-                    task.invoke()
-                }
-            } catch (e: Exception) {
-                // Will rethrow on the calling thread.
-                exception = e
-            }
-            latch.countDown()
-        }
-        ) {
-            latch.await()
-            if (exception != null) {
-                throw exception as Exception
-            }
-        }
-        // R might be nullable so we can't assert nullability here.
-        @Suppress("UNCHECKED_CAST")
-        returnVal as R
-    }
-}
-
-/**
  * Runs the supplied task on the handler thread. If we're not on the handler thread a task is
  * posted.
  *
@@ -2472,19 +2426,16 @@
     traceEventName: String,
     task: suspend () -> R
 ): R = TraceEvent(traceEventName).use {
-    val latch = CountDownLatch(1)
-    var r: R? = null
-    launch {
-        try {
-            r = task()
-        } catch (e: Exception) {
-            Log.e("CoroutineScope", "Exception in traceEventName", e)
-            throw e
+    try {
+        return runBlocking {
+            withContext(coroutineContext) {
+                task()
+            }
         }
-        latch.countDown()
+    } catch (e: Exception) {
+        Log.e("CoroutineScope", "Exception in traceEventName", e)
+        throw e
     }
-    latch.await()
-    return r!!
 }
 
 /**

diff --git a/wear/watchface/watchface/src/main/java/androidx/wear/watchface/control/WatchFaceControlService.kt b/wear/watchface/watchface/src/main/java/androidx/wear/watchface/control/WatchFaceControlService.kt
index ed79006..544256e 100644
--- a/wear/watchface/watchface/src/main/java/androidx/wear/watchface/control/WatchFaceControlService.kt
+++ b/wear/watchface/watchface/src/main/java/androidx/wear/watchface/control/WatchFaceControlService.kt

@@ -20,9 +20,7 @@
 import android.content.ComponentName
 import android.content.Context
 import android.content.Intent
-import android.os.Handler
 import android.os.IBinder
-import android.os.Looper
 import android.util.Log
 import androidx.annotation.RequiresApi
 import androidx.annotation.RestrictTo
@@ -42,11 +40,13 @@
 import androidx.wear.watchface.control.data.WallpaperInteractiveWatchFaceInstanceParams
 import androidx.wear.watchface.data.ComplicationSlotMetadataWireFormat
 import androidx.wear.watchface.editor.EditorService
-import androidx.wear.watchface.runBlockingOnHandlerWithTracing
+import androidx.wear.watchface.runBlockingWithTracing
 import androidx.wear.watchface.style.data.UserStyleFlavorsWireFormat
 import androidx.wear.watchface.style.data.UserStyleSchemaWireFormat
 import java.io.FileDescriptor
 import java.io.PrintWriter
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.MainScope
 
 /**
  * A service for creating and controlling watch face instances.
@@ -77,7 +77,7 @@
     @VisibleForTesting
     public open fun createServiceStub(): IWatchFaceInstanceServiceStub =
         TraceEvent("WatchFaceControlService.createServiceStub").use {
-            IWatchFaceInstanceServiceStub(this, Handler(Looper.getMainLooper()))
+            IWatchFaceInstanceServiceStub(this, MainScope())
         }
 
     @VisibleForTesting
@@ -95,29 +95,12 @@
     }
 }
 
-/**
- * Factory for use by on watch face editors to create [IWatchFaceControlService].
- *
- * @hide
- */
-@RestrictTo(RestrictTo.Scope.LIBRARY_GROUP)
-@RequiresApi(27)
-public class WatchFaceControlServiceFactory {
-    public companion object {
-        @JvmStatic
-        public fun createWatchFaceControlService(
-            context: Context,
-            uiThreadHandler: Handler
-        ): IWatchFaceControlService = IWatchFaceInstanceServiceStub(context, uiThreadHandler)
-    }
-}
-
 /** @hide */
 @RequiresApi(27)
 @RestrictTo(RestrictTo.Scope.LIBRARY_GROUP)
 public open class IWatchFaceInstanceServiceStub(
     private val context: Context,
-    private val uiThreadHandler: Handler
+    private val uiThreadCoroutineScope: CoroutineScope
 ) : IWatchFaceControlService.Stub() {
     override fun getApiVersion(): Int = IWatchFaceControlService.API_VERSION
 
@@ -139,7 +122,7 @@
         val engine = createHeadlessEngine(params.watchFaceName, context)
         engine?.let {
             // This is serviced on a background thread so it should be fine to block.
-            uiThreadHandler.runBlockingOnHandlerWithTracing("createHeadlessInstance") {
+            uiThreadCoroutineScope.runBlockingWithTracing("createHeadlessInstance") {
                 // However the WatchFaceService.createWatchFace method needs to be run on the UI
                 // thread.
                 it.createHeadlessInstance(params)