Promote more lifecycle to ZiplineCodeSession (#1655)

* Promote more lifecycle to ZiplineCodeSession

Previously TreehouseApp was tracking exceptions on behalf
of a class that could do it better.

This introduces listeners at the CodeSession level in
addition to the existing listeners at the CodeHost level.

* Scope the FrameClock to a CodeSession

redwood-treehouse-host/src/androidInstrumentedTest/kotlin/app/cash/redwood/treehouse/AndroidChoreographerFrameClockTest.kt

@@ -20,7 +20,7 @@ import kotlinx.coroutines.runBlocking
 
 class AndroidChoreographerFrameClockTest : AbstractFrameClockTest() {
   // Tests run on a background thread but Choreographer can only be grabbed from the main thread.
-  override val frameClock = runBlocking(Dispatchers.Main) {
-    AndroidChoreographerFrameClock()
+  override val frameClockFactory = runBlocking(Dispatchers.Main) {
+    AndroidChoreographerFrameClock.Factory()
   }
 }

redwood-treehouse-host/src/androidMain/kotlin/app/cash/redwood/treehouse/AndroidChoreographerFrameClock.kt

@@ -22,18 +22,11 @@ import kotlinx.coroutines.launch
 /**
  * A [FrameClock] that sends frames using [Choreographer].
  */
-internal class AndroidChoreographerFrameClock : FrameClock {
-  private val choreographer = Choreographer.getInstance()
-  private lateinit var scope: CoroutineScope
-  private lateinit var dispatchers: TreehouseDispatchers
-
-  override fun start(
-    scope: CoroutineScope,
-    dispatchers: TreehouseDispatchers,
-  ) {
-    this.scope = scope
-    this.dispatchers = dispatchers
-  }
+internal class AndroidChoreographerFrameClock private constructor(
+  private val choreographer: Choreographer,
+  private val scope: CoroutineScope,
+  private val dispatchers: TreehouseDispatchers,
+) : FrameClock {
 
   override fun requestFrame(appLifecycle: AppLifecycle) {
     choreographer.postFrameCallback { frameTimeNanos ->
@@ -42,4 +35,12 @@ internal class AndroidChoreographerFrameClock : FrameClock {
       }
     }
   }
+
+  class Factory : FrameClock.Factory {
+    private val choreographer = Choreographer.getInstance()
+    override fun create(
+      scope: CoroutineScope,
+      dispatchers: TreehouseDispatchers,
+    ) = AndroidChoreographerFrameClock(choreographer, scope, dispatchers)
+  }
 }

redwood-treehouse-host/src/androidMain/kotlin/app/cash/redwood/treehouse/treehouseAppFactoryAndroid.kt

@@ -38,7 +38,7 @@ public fun TreehouseAppFactory(
   dispatchers = AndroidTreehouseDispatchers(),
   eventListener = eventListener,
   httpClient = httpClient.asZiplineHttpClient(),
-  frameClock = AndroidChoreographerFrameClock(),
+  frameClockFactory = AndroidChoreographerFrameClock.Factory(),
   manifestVerifier = manifestVerifier,
   embeddedDir = embeddedDir,
   embeddedFileSystem = embeddedFileSystem,

redwood-treehouse-host/src/commonMain/kotlin/app/cash/redwood/treehouse/CodeHost.kt

@@ -15,49 +15,18 @@
  */
 package app.cash.redwood.treehouse
 
-import kotlin.coroutines.CoroutineContext
-import kotlinx.coroutines.CoroutineExceptionHandler
-
 /** Manages loading and hot-reloading a series of code sessions. */
 internal interface CodeHost<A : AppService> {
   val stateStore: StateStore
 
   /** Only accessed on [TreehouseDispatchers.ui]. */
   val session: CodeSession<A>?
 
-  fun newServiceScope(): ServiceScope<A>
-
-  /** Cancels the current code and propagates [exception] to all listeners. */
-  fun handleUncaughtException(exception: Throwable)
-
   fun addListener(listener: Listener<A>)
 
   fun removeListener(listener: Listener<A>)
 
   interface Listener<A : AppService> {
     fun codeSessionChanged(next: CodeSession<A>)
-    fun uncaughtException(exception: Throwable)
-  }
-
-  /**
-   * Tracks all of the services created to produce a UI, and offers a single mechanism to close
-   * them all. Note that closing this does not close the app services it was applied to.
-   */
-  interface ServiceScope<A : AppService> {
-    /**
-     * Returns a new instance that forwards calls to [appService] and keeps track of returned
-     * instances so they may be closed.
-     */
-    fun apply(appService: A): A
-    fun close()
-  }
-}
-
-internal fun CodeHost<*>.asExceptionHandler() = object : CoroutineExceptionHandler {
-  override val key: CoroutineContext.Key<*>
-    get() = CoroutineExceptionHandler.Key
-
-  override fun handleException(context: CoroutineContext, exception: Throwable) {
-    handleUncaughtException(exception)
   }
 }

redwood-treehouse-host/src/commonMain/kotlin/app/cash/redwood/treehouse/CodeSession.kt

@@ -15,6 +15,9 @@
  */
 package app.cash.redwood.treehouse
 
+import kotlin.coroutines.CoroutineContext
+import kotlinx.coroutines.CoroutineExceptionHandler
+import kotlinx.coroutines.CoroutineScope
 import kotlinx.serialization.json.Json
 
 /** The host state for a single code load. We get a new session each time we get new code. */
@@ -23,7 +26,44 @@ internal interface CodeSession<A : AppService> {
 
   val json: Json
 
-  fun start()
+  fun start(sessionScope: CoroutineScope, frameClock: FrameClock)
+
+  fun addListener(listener: Listener<A>)
+
+  fun removeListener(listener: Listener<A>)
+
+  fun newServiceScope(): ServiceScope<A>
+
+  /** Propagates [exception] to all listeners and cancels this session. */
+  fun handleUncaughtException(exception: Throwable)
 
   fun cancel()
+
+  /**
+   * Tracks all of the services created to produce a UI, and offers a single mechanism to close
+   * them all. Note that closing this does not close the app services it was applied to.
+   */
+  interface ServiceScope<A : AppService> {
+    /**
+     * Returns a new instance that forwards calls to [appService] and keeps track of returned
+     * instances so they may be closed.
+     */
+    fun apply(appService: A): A
+    fun close()
+  }
+
+  interface Listener<A : AppService> {
+    fun onUncaughtException(codeSession: CodeSession<A>, exception: Throwable)
+    fun onCancel(codeSession: CodeSession<A>)
+  }
 }
+
+internal val CodeSession<*>.coroutineExceptionHandler: CoroutineExceptionHandler
+  get() = object : CoroutineExceptionHandler {
+    override val key: CoroutineContext.Key<*>
+      get() = CoroutineExceptionHandler.Key
+
+    override fun handleException(context: CoroutineContext, exception: Throwable) {
+      handleUncaughtException(exception)
+    }
+  }

redwood-treehouse-host/src/commonMain/kotlin/app/cash/redwood/treehouse/FrameClock.kt

@@ -18,17 +18,16 @@ package app.cash.redwood.treehouse
 import kotlinx.coroutines.CoroutineScope
 
 internal interface FrameClock {
-  /** Run this clock until [scope] is canceled. */
-  fun start(
-    scope: CoroutineScope,
-    dispatchers: TreehouseDispatchers,
-  )
-
   /**
    * Request a call to [AppLifecycle.sendFrame]. It is an error to call [requestFrame] again before
    * that call is made.
    *
    * It is an error to call this before [start].
    */
   fun requestFrame(appLifecycle: AppLifecycle)
+
+  interface Factory {
+    /** Creates a new FrameClock that sends frames on [dispatchers] in [scope]. */
+    fun create(scope: CoroutineScope, dispatchers: TreehouseDispatchers): FrameClock
+  }
 }

redwood-treehouse-host/src/commonMain/kotlin/app/cash/redwood/treehouse/TreehouseApp.kt

@@ -16,16 +16,12 @@
 package app.cash.redwood.treehouse
 
 import app.cash.zipline.Zipline
-import app.cash.zipline.ZiplineScope
 import app.cash.zipline.loader.LoadResult
 import app.cash.zipline.loader.ManifestVerifier
 import app.cash.zipline.loader.ZiplineCache
 import app.cash.zipline.loader.ZiplineHttpClient
 import app.cash.zipline.loader.ZiplineLoader
-import app.cash.zipline.withScope
 import kotlin.native.ObjCName
-import kotlinx.coroutines.CoroutineDispatcher
-import kotlinx.coroutines.CoroutineExceptionHandler
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.SupervisorJob
 import kotlinx.coroutines.flow.Flow
@@ -50,8 +46,7 @@ public class TreehouseApp<A : AppService> private constructor(
 ) {
   private val codeHost = ZiplineCodeHost<A>()
 
-  public val dispatchers: TreehouseDispatchers =
-    TreehouseDispatchersWithExceptionHandler(factory.dispatchers, codeHost.asExceptionHandler())
+  public val dispatchers: TreehouseDispatchers = factory.dispatchers
 
   private val eventPublisher = RealEventPublisher(factory.eventListener, this)
 
@@ -123,14 +118,10 @@ public class TreehouseApp<A : AppService> private constructor(
    * Continuously polls for updated code, and emits a new [LoadResult] instance when new code is
    * found.
    */
-  @OptIn(ExperimentalStdlibApi::class)
   private fun ziplineFlow(): Flow<LoadResult> {
-    val dispatcher = dispatchers.zipline[CoroutineDispatcher.Key]
-      ?: error("expected TreehouseDispatchers.zipline to include a CoroutineDispatcher")
-
     // Loads applications from the network only. The cache is neither read nor written.
     var loader = ZiplineLoader(
-      dispatcher = dispatcher,
+      dispatcher = dispatchers.zipline,
       manifestVerifier = factory.manifestVerifier,
       httpClient = factory.httpClient,
       eventListener = eventPublisher.ziplineEventListener,
@@ -179,20 +170,14 @@ public class TreehouseApp<A : AppService> private constructor(
     closed = true
     appScope.launch(dispatchers.ui) {
       val session = codeHost.session ?: return@launch
+      session.removeListener(codeHost)
       session.cancel()
       codeHost.session = null
     }
     eventPublisher.appCanceled()
   }
 
-  private class TreehouseDispatchersWithExceptionHandler(
-    val delegate: TreehouseDispatchers,
-    exceptionHandler: CoroutineExceptionHandler,
-  ) : TreehouseDispatchers by delegate {
-    override val zipline = delegate.zipline + exceptionHandler
-  }
-
-  private inner class ZiplineCodeHost<A : AppService> : CodeHost<A> {
+  private inner class ZiplineCodeHost<A : AppService> : CodeHost<A>, CodeSession.Listener<A> {
     /**
      * Contents that this app is currently responsible for.
      *
@@ -204,20 +189,6 @@ public class TreehouseApp<A : AppService> private constructor(
 
     override var session: ZiplineCodeSession<A>? = null
 
-    override fun newServiceScope(): CodeHost.ServiceScope<A> {
-      val ziplineScope = ZiplineScope()
-
-      return object : CodeHost.ServiceScope<A> {
-        override fun apply(appService: A): A {
-          return appService.withScope(ziplineScope)
-        }
-
-        override fun close() {
-          ziplineScope.close()
-        }
-      }
-    }
-
     override fun addListener(listener: CodeHost.Listener<A>) {
       dispatchers.checkUi()
       listeners += listener
@@ -228,47 +199,42 @@ public class TreehouseApp<A : AppService> private constructor(
       listeners -= listener
     }
 
-    override fun handleUncaughtException(exception: Throwable) {
-      appScope.launch(dispatchers.ui) {
-        for (listener in listeners) {
-          listener.uncaughtException(exception)
-        }
-        codeHost.session?.cancel()
-        codeHost.session = null
-      }
+    override fun onUncaughtException(codeSession: CodeSession<A>, exception: Throwable) {
+    }
 
-      eventPublisher.onUncaughtException(exception)
+    override fun onCancel(codeSession: CodeSession<A>) {
+      check(codeSession == this.session)
+      this.session = null
     }
 
     fun onCodeChanged(zipline: Zipline, appService: A) {
-      val sessionScope = CoroutineScope(SupervisorJob(appScope.coroutineContext.job))
+      val next = ZiplineCodeSession(
+        dispatchers = dispatchers,
+        eventPublisher = eventPublisher,
+        appScope = appScope,
+        appService = appService,
+        zipline = zipline,
+      )
+
+      val sessionScope = CoroutineScope(
+        SupervisorJob(appScope.coroutineContext.job) + next.coroutineExceptionHandler,
+      )
+
       sessionScope.launch(dispatchers.ui) {
         val previous = session
+        previous?.removeListener(this@ZiplineCodeHost)
+        previous?.cancel()
 
-        val next = ZiplineCodeSession(
-          codeHost = this@ZiplineCodeHost,
-          dispatchers = dispatchers,
-          eventPublisher = eventPublisher,
-          appScope = appScope,
-          frameClock = factory.frameClock,
+        session = next
+        next.addListener(this@ZiplineCodeHost)
+        next.start(
           sessionScope = sessionScope,
-          appService = appService,
-          zipline = zipline,
+          frameClock = factory.frameClockFactory.create(sessionScope, dispatchers),
         )
 
-        next.start()
-
         for (listener in listeners) {
           listener.codeSessionChanged(next)
         }
-
-        if (previous != null) {
-          sessionScope.launch(dispatchers.zipline) {
-            previous.cancel()
-          }
-        }
-
-        session = next
       }
     }
   }
@@ -285,7 +251,7 @@ public class TreehouseApp<A : AppService> private constructor(
     public val dispatchers: TreehouseDispatchers,
     internal val eventListener: EventListener,
     internal val httpClient: ZiplineHttpClient,
-    internal val frameClock: FrameClock,
+    internal val frameClockFactory: FrameClock.Factory,
     internal val manifestVerifier: ManifestVerifier,
     internal val embeddedDir: Path?,
     internal val embeddedFileSystem: FileSystem?,