Allow device selection

Author: wenli-cai

workflow-trace-viewer/src/jvmMain/kotlin/com/squareup/workflow1/traceviewer/App.kt

@@ -15,6 +15,7 @@ import androidx.compose.ui.Modifier
 import androidx.compose.ui.geometry.Offset
 import com.squareup.workflow1.traceviewer.model.NodeUpdate
 import com.squareup.workflow1.traceviewer.ui.ColorLegend
+import com.squareup.workflow1.traceviewer.ui.DisplayDevices
 import com.squareup.workflow1.traceviewer.ui.FrameSelectTab
 import com.squareup.workflow1.traceviewer.ui.RightInfoPanel
 import com.squareup.workflow1.traceviewer.ui.TraceModeToggleSwitch
@@ -63,8 +64,9 @@ internal fun App(
       sandboxState = sandboxState,
     ) {
       // if there is not a file selected and trace mode is live, then don't render anything.
-      val readyForFileTrace = traceMode is TraceMode.File && selectedTraceFile != null
-      val readyForLiveTrace = traceMode is TraceMode.Live
+      val readyForFileTrace = validateFileMode(traceMode)
+      val readyForLiveTrace = validateLiveMode(traceMode)
+
       if (readyForFileTrace || readyForLiveTrace) {
         active = true
         RenderTrace(
@@ -92,13 +94,12 @@ internal fun App(
           frameIndex = 0
           TraceMode.File(null)
         } else {
-          // TODO: TraceRecorder needs to be able to take in multiple clients if this is the case
           /*
           We set the frame to -1 here since we always increment it during Live mode as the list of
           frames get populated, so we avoid off by one when indexing into the frames.
            */
           frameIndex = -1
-          TraceMode.Live
+          TraceMode.Live()
         }
       },
       traceMode = traceMode,
@@ -118,6 +119,16 @@ internal fun App(
     }
 
     if (traceMode is TraceMode.Live) {
+      if ((traceMode as TraceMode.Live).device == null) {
+        DisplayDevices(
+          onDeviceSelected = { selectedDevice ->
+            traceMode = TraceMode.Live(selectedDevice)
+          },
+          devices = listDevices(),
+          modifier = Modifier.align(Alignment.Center)
+        )
+      }
+
       FileDump(
         trace = rawRenderPass,
         modifier = Modifier.align(Alignment.BottomStart)
@@ -149,5 +160,23 @@ internal class SandboxState {
 
 internal sealed interface TraceMode {
   data class File(val file: PlatformFile?) : TraceMode
-  data object Live : TraceMode
+  data class Live(val device: String? = null) : TraceMode
+}
+
+/**
+ * Allows users to select from multiple devices that are currently running.
+ */
+internal fun listDevices(): List<String> {
+  val process = ProcessBuilder("adb", "devices", "-l").start()
+  process.waitFor()
+  // We drop the header "List of devices attached"
+  return process.inputStream.bufferedReader().readLines().drop(1).dropLast(1)
+}
+
+internal fun validateLiveMode(traceMode: TraceMode): Boolean {
+  return traceMode is TraceMode.Live && traceMode.device != null
+}
+
+internal fun validateFileMode(traceMode: TraceMode): Boolean {
+  return traceMode is TraceMode.File && traceMode.file != null
 }

workflow-trace-viewer/src/jvmMain/kotlin/com/squareup/workflow1/traceviewer/ui/DisplayDevices.kt

@@ -0,0 +1,60 @@
+package com.squareup.workflow1.traceviewer.ui
+
+import androidx.compose.foundation.BorderStroke
+import androidx.compose.foundation.layout.Box
+import androidx.compose.foundation.layout.Column
+import androidx.compose.foundation.layout.fillMaxWidth
+import androidx.compose.foundation.layout.padding
+import androidx.compose.foundation.shape.RoundedCornerShape
+import androidx.compose.material.Card
+import androidx.compose.material.ExperimentalMaterialApi
+import androidx.compose.material.Text
+import androidx.compose.runtime.Composable
+import androidx.compose.ui.Alignment
+import androidx.compose.ui.Modifier
+import androidx.compose.ui.graphics.Color
+import androidx.compose.ui.unit.dp
+
+@OptIn(ExperimentalMaterialApi::class)
+@Composable
+internal fun DisplayDevices(
+  onDeviceSelected: (String) -> Unit,
+  devices: List<String>,
+  modifier: Modifier = Modifier,
+) {
+  Box(
+    modifier = modifier
+      .fillMaxWidth(),
+    contentAlignment = Alignment.Center
+  ) {
+    if (devices.isEmpty()) {
+      Text(
+        text = "No device available",
+        modifier = Modifier.align(Alignment.Center)
+      )
+      return@Box
+    }
+
+    val emulatorRegex = Regex("""\bemulator-\d+\b""")
+    Column {
+      devices.forEach { device ->
+        Card(
+          onClick = {
+            emulatorRegex.find(device)?.value?.let { emulator ->
+              onDeviceSelected(emulator)
+            }
+          },
+          shape = RoundedCornerShape(16.dp),
+          border = BorderStroke(1.dp, Color.Gray),
+          modifier = Modifier.padding(4.dp),
+          elevation = 2.dp
+        ) {
+          Text(
+            text = device,
+            modifier = Modifier.padding(horizontal = 16.dp, vertical = 8.dp)
+          )
+        }
+      }
+    }
+  }
+}

workflow-trace-viewer/src/jvmMain/kotlin/com/squareup/workflow1/traceviewer/util/SocketClient.kt

@@ -11,12 +11,12 @@ import kotlinx.coroutines.withContext
 import okio.IOException
 import java.net.Socket
 
-internal suspend fun streamRenderPassesFromDevice(parseOnNewRenderPass: (String) -> Unit) {
+internal suspend fun streamRenderPassesFromDevice(device: String, parseOnNewRenderPass: (String) -> Unit) {
   val renderPassChannel: Channel<String> = Channel(Channel.BUFFERED)
   coroutineScope {
     launch {
       try {
-        pollSocket(onNewRenderPass = renderPassChannel::send)
+        pollSocket(device = device, onNewRenderPass = renderPassChannel::send)
       } finally {
         renderPassChannel.close()
       }
@@ -39,10 +39,10 @@ internal suspend fun streamRenderPassesFromDevice(parseOnNewRenderPass: (String)
  * @param onNewRenderPass is called from an arbitrary thread, so it is important to ensure that the
  *        caller is thread safe
  */
-private suspend fun pollSocket(onNewRenderPass: suspend (String) -> Unit) {
+private suspend fun pollSocket(device: String, onNewRenderPass: suspend (String) -> Unit) {
   withContext(Dispatchers.IO) {
     try {
-      runForwardingPortThroughAdb { port ->
+      runForwardingPortThroughAdb(device) { port ->
         Socket("localhost", port).useWithCancellation { socket ->
           val reader = socket.getInputStream().bufferedReader()
           do {
@@ -60,7 +60,6 @@ private suspend fun pollSocket(onNewRenderPass: suspend (String) -> Unit) {
   }
 }
 
-
 /**
  * Force [pollSocket] to exit with exception if the coroutine is cancelled. See comment below.
  */
@@ -89,9 +88,9 @@ private suspend fun Socket.useWithCancellation(block: suspend (Socket) -> Unit)
  * If block throws or returns on finish, the port forwarding is removed via adb (best effort).
  */
 @Suppress("BlockingMethodInNonBlockingContext")
-private suspend inline fun runForwardingPortThroughAdb(block: (port: Int) -> Unit) {
+private suspend inline fun runForwardingPortThroughAdb(device: String, block: (port: Int) -> Unit) {
   val process = ProcessBuilder(
-    "adb", "forward", "tcp:0", "localabstract:workflow-trace"
+    "adb", "-s", device, "forward", "tcp:0", "localabstract:workflow-trace"
   ).start()
 
   // The adb forward command will output the port number it picks to connect.

workflow-trace-viewer/src/jvmMain/kotlin/com/squareup/workflow1/traceviewer/util/TraceParser.kt

@@ -34,9 +34,9 @@ internal fun RenderTrace(
 ) {
   var isLoading by remember(traceSource) { mutableStateOf(true) }
   var error by remember(traceSource) { mutableStateOf<String?>(null) }
-  val frames = remember (traceSource){ mutableStateListOf<Node>() }
-  val fullTree = remember (traceSource){ mutableStateListOf<Node>() }
-  val affectedNodes = remember (traceSource){ mutableStateListOf<Set<Node>>() }
+  val frames = remember(traceSource) { mutableStateListOf<Node>() }
+  val fullTree = remember(traceSource) { mutableStateListOf<Node>() }
+  val affectedNodes = remember(traceSource) { mutableStateListOf<Set<Node>>() }
 
   // Updates current state with the new data from trace source.
   fun addToStates(frame: List<Node>, tree: List<Node>, affected: List<Set<Node>>) {
@@ -83,8 +83,11 @@ internal fun RenderTrace(
       }
 
       is TraceMode.Live -> {
+        checkNotNull(traceSource.device) {
+          "TraceMode.Live requires a selected device"
+        }
         val adapter: JsonAdapter<List<Node>> = createMoshiAdapter<Node>()
-        streamRenderPassesFromDevice { renderPass ->
+        streamRenderPassesFromDevice(traceSource.device) { renderPass ->
           val currentTree = fullTree.lastOrNull()
           val parseResult = parseLiveTrace(renderPass, adapter, currentTree)
           handleParseResult(parseResult, renderPass, onNewFrame)