More socket refactoring

Author: wenli-cai

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

@@ -1,7 +1,9 @@
 package com.squareup.workflow1.traceviewer.util
 
+import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.awaitCancellation
+import kotlinx.coroutines.channels.Channel
 import kotlinx.coroutines.coroutineScope
 import kotlinx.coroutines.ensureActive
 import kotlinx.coroutines.launch
@@ -10,6 +12,22 @@ import kotlinx.coroutines.withContext
 import okio.IOException
 import java.net.Socket
 
+internal suspend fun streamRenderPassesFromDevice(parseOnNewRenderPass: (String) -> Unit) {
+  val renderPassChannel: Channel<String> = Channel(Channel.BUFFERED)
+  coroutineScope {
+    try {
+      pollSocket(onNewRenderPass = renderPassChannel::send)
+    } finally {
+      renderPassChannel.close()
+    }
+  }
+
+  // Since channel implements ChannelIterator, we can for-loop through on the receiver end.
+  for (renderPass in renderPassChannel) {
+    parseOnNewRenderPass(renderPass)
+  }
+}
+
 /**
  * Collects data from a server socket and serves them back to the caller via callback.
  *
@@ -20,7 +38,7 @@ import java.net.Socket
  * @param onNewRenderPass is called from an arbitrary thread, so it is important to ensure that the
  *        caller is thread safe
  */
-suspend fun pollSocket(onNewRenderPass: suspend (String) -> Unit) {
+private suspend fun pollSocket(onNewRenderPass: suspend (String) -> Unit) {
   withContext(Dispatchers.IO) {
     try {
       runForwardingPortThroughAdb { port ->

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

@@ -14,9 +14,6 @@ import com.squareup.moshi.JsonAdapter
 import com.squareup.workflow1.traceviewer.TraceMode
 import com.squareup.workflow1.traceviewer.model.Node
 import com.squareup.workflow1.traceviewer.ui.DrawTree
-import kotlinx.coroutines.channels.Channel
-import kotlinx.coroutines.launch
-import java.net.SocketException
 
 /**
  * Handles parsing the trace's after JsonParser has turned all render passes into frames. Also calls
@@ -58,6 +55,7 @@ internal fun RenderTrace(
       is ParseResult.Failure -> {
         error = parseResult.error.toString()
       }
+
       is ParseResult.Success -> {
         addToStates(
           frame = parseResult.trace,
@@ -80,23 +78,13 @@ internal fun RenderTrace(
       }
 
       is TraceMode.Live -> {
-        val renderPassChannel: Channel<String> = Channel(Channel.BUFFERED)
-        launch {
-          try {
-            pollSocket(onNewRenderPass = renderPassChannel::send)
-            error = "Socket has already been closed or is not available."
-          } finally {
-            renderPassChannel.close()
-          }
-        }
         val adapter: JsonAdapter<List<Node>> = createMoshiAdapter<Node>()
-
-        // Since channel implements ChannelIterator, we can for-loop through on the receiver end.
-        for (renderPass in renderPassChannel) {
+        streamRenderPassesFromDevice { renderPass ->
           val currentTree = fullTree.lastOrNull()
           val parseResult = parseLiveTrace(renderPass, adapter, currentTree)
           handleParseResult(parseResult, onNewFrame)
         }
+        error = "Socket has already been closed or is not available."
       }
     }
   }