Fix PR comments

Author: wenli-cai

workflow-trace-viewer/README.md

@@ -12,7 +12,7 @@ It can be run via Gradle using:
 
 By Default, the app will be in file parsing mode, where you are able to select a previously recorded workflow trace file for it to visualize the data.
 
-By hitting the bottom switch, you are able to toggle to live stream mode, where data is directly pulled from the emulator into the visualizer.
+By hitting the bottom switch, you are able to toggle to live stream mode, where data is directly pulled from the emulator into the visualizer. A connection can only happen once. If there needs to be rerecording of the trace, the emulator must first be restarted, and then the app must be restarted as well. This is due to the fact that any open socket will consume all render pass data, meaning there is nothing to read from the emulator.
 
 It is ***important*** to run the emulator first before toggling to live mode.
 

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

@@ -42,12 +42,6 @@ internal fun App(
   var selectedTraceFile by remember { mutableStateOf<PlatformFile?>(null) }
   val socket = remember { SocketClient() }
 
-  Runtime.getRuntime().addShutdownHook(
-    Thread {
-      socket.close()
-    }
-  )
-
   LaunchedEffect(sandboxState) {
     snapshotFlow { frameIndex }.collect {
       sandboxState.reset()
@@ -57,7 +51,7 @@ internal fun App(
   Box(
     modifier = modifier
   ) {
-    fun resetStates() = run {
+    fun resetStates() {
       socket.close()
       selectedTraceFile = null
       selectedNode = null
@@ -112,7 +106,7 @@ internal fun App(
           frames get populated, so we avoid off by one when indexing into the frames.
            */
           frameIndex = -1
-          socket.start()
+          socket.open()
           TraceMode.Live(socket)
         }
       },
@@ -140,11 +134,10 @@ internal class SandboxState {
 
   fun reset() {
     offset = Offset.Zero
-    // scale = 1f
   }
 }
 
 internal sealed interface TraceMode {
   data class File(val file: PlatformFile?) : TraceMode
-  data class Live(val socket: SocketClient = SocketClient()) : TraceMode
+  data class Live(val socket: SocketClient) : TraceMode
 }

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

@@ -14,7 +14,7 @@ fun main() {
         .start().waitFor()
     }
   )
-  singleWindowApplication(title = "Workflow Trace Viewer") {
+  singleWindowApplication(title = "Workflow Trace Viewer", exitProcessOnExit = false) {
     App(Modifier.fillMaxSize())
   }
 }

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

@@ -32,9 +32,10 @@ internal fun FrameSelectTab(
   modifier: Modifier = Modifier
 ) {
   val lazyListState = rememberLazyListState()
-  LaunchedEffect(currentIndex) {
-    if (currentIndex < 0) return@LaunchedEffect
-    lazyListState.animateScrollToItem(currentIndex)
+  if (currentIndex >= 0) {
+    LaunchedEffect(currentIndex) {
+      lazyListState.animateScrollToItem(currentIndex)
+    }
   }
 
   Surface(

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

@@ -34,7 +34,9 @@ import com.squareup.workflow1.traceviewer.util.createMoshiAdapter
 import com.squareup.workflow1.traceviewer.util.parseFileTrace
 import com.squareup.workflow1.traceviewer.util.parseLiveTrace
 import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.launch
 import kotlinx.coroutines.withContext
+import java.net.SocketException
 
 /**
  * Access point for drawing the main content of the app. It will load the trace for given files and
@@ -71,11 +73,10 @@ internal fun RenderTrace(
   fun handleParseResult(
     parseResult: ParseResult,
     onNewFrame: (() -> Unit)? = null
-  ): Boolean {
-    return when (parseResult) {
+  ) {
+    when (parseResult) {
       is ParseResult.Failure -> {
         error = parseResult.error
-        false
       }
       is ParseResult.Success -> {
         addToStates(
@@ -84,28 +85,39 @@ internal fun RenderTrace(
           affected = parseResult.affectedNodes
         )
         onNewFrame?.invoke()
-        true
       }
     }
   }
 
   LaunchedEffect(traceSource) {
     when (traceSource) {
       is TraceMode.File -> {
-        // We guarantee the file is null since this composable can only be called when a file is selected.
-        val parseResult = parseFileTrace(traceSource.file!!)
+        checkNotNull(traceSource.file){
+          "TraceMode.File should have a non-null file to parse."
+        }
+        val parseResult = parseFileTrace(traceSource.file)
         handleParseResult(parseResult)
       }
 
       is TraceMode.Live -> {
         val socket = traceSource.socket
-        socket.beginListen(this)
+        launch {
+          try {
+            socket.pollSocket()
+          } catch (e: SocketException) {
+            error = SocketException("Socket has already been closed or is not available: ${e.message}")
+            return@launch
+          }
+        }
+        if (error != null) {
+          return@LaunchedEffect
+        }
         val adapter: JsonAdapter<List<Node>> = createMoshiAdapter<Node>()
 
-        withContext(Dispatchers.IO) {
+        withContext(Dispatchers.Default) {
           // Since channel implements ChannelIterator, we can for-loop through on the receiver end.
           for (renderPass in socket.renderPassChannel) {
-            val currentTree = if (fullTree.isEmpty()) null else fullTree.last()
+            val currentTree = fullTree.lastOrNull()
             val parseResult = parseLiveTrace(renderPass, adapter, currentTree)
             handleParseResult(parseResult, onNewFrame)
           }

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

@@ -2,20 +2,28 @@ package com.squareup.workflow1.traceviewer.util
 
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.Job
 import kotlinx.coroutines.channels.Channel
 import kotlinx.coroutines.launch
+import kotlinx.coroutines.withContext
+import kotlinx.coroutines.delay
+import kotlinx.coroutines.yield
+
 import java.net.Socket
 import java.net.SocketException
 
 /**
- * This is a client that connects to the `ActionLogger` Unix Domain Socket and listens for any new
- * render passes. Since this app is on JVM and the server is on Android, we use ADB to forward the
- * port onto the socket.
+ * This is a client that can connect to any server socket that sends render pass data while using
+ * the Workflow framework.
+ *
+ * [start] and [close] are idempotent commands, so this socket can only be started and closed once.
+ *
+ * Since this app is on JVM and the server is on Android, we use ADB to forward the port onto the socket.
  */
 internal class SocketClient {
   private lateinit var socket: Socket
   private var initialized = false
-  val renderPassChannel: Channel<String> = Channel(Channel.UNLIMITED)
+  val renderPassChannel: Channel<String> = Channel(Channel.BUFFERED)
 
   /**
    * We use any available ports on the host machine to connect to the emulator.
@@ -24,18 +32,20 @@ internal class SocketClient {
    * `LocalServerSocket` -- which creates a unix socket on the linux abstract namespace -- we use
    * `localabstract:` to connect to it.
    */
-  fun start() {
+  fun open() {
+    if (initialized){
+      return
+    }
     initialized = true
     val process = ProcessBuilder(
       "adb", "forward", "tcp:0", "localabstract:workflow-trace"
     ).start()
 
     // The adb forward command will output the port number it picks to connect.
-    val port = run {
-      process.waitFor()
-      process.inputStream.bufferedReader().readText()
+    process.waitFor()
+    val port = process.inputStream.bufferedReader().readText()
         .trim().toInt()
-    }
+
     socket = Socket("localhost", port)
   }
 
@@ -47,22 +57,25 @@ internal class SocketClient {
   }
 
   /**
-   * This will always be called within an asynchronous call, so we do not need to block/launch a
-   * new coroutine here.
+   * Polls the socket's input stream and sends the data into [renderPassChannel].
+   * The caller should handle the scope of the coroutine that this function is called in.
    *
    * To better separate the responsibility of reading from the socket, we use a channel for the caller
    * to handle parsing and amalgamating the render passes.
    */
-  fun beginListen(scope: CoroutineScope) {
-    scope.launch(Dispatchers.IO) {
+  suspend fun pollSocket() {
+    withContext(Dispatchers.IO) {
       val reader = socket.getInputStream().bufferedReader()
-      try {
-        while (true) {
-          val input = reader.readLine()
-          renderPassChannel.trySend(input)
+      reader.use {
+        try {
+          while (true) {
+            val input = reader.readLine()
+            renderPassChannel.trySend(input)
+            println(input)
+          }
+        } catch (e: SocketException) {
+          e.printStackTrace()
         }
-      } catch (e: SocketException) {
-        println("Exiting socket listener due to: ${e.message}")
       }
     }
   }