clean up DOMs when client state expires, hopefully prevent resource leaks

Author: sanity

src/main/kotlin/kweb/Kweb.kt

@@ -140,6 +140,9 @@ class Kweb private constructor(
     val clientState: Cache<String, RemoteClientState> = CacheBuilder.newBuilder()
         .expireAfterAccess(kwebConfig.clientStateTimeout)
         .apply { if (kwebConfig.clientStateStatsEnabled) recordStats() }
+        .removalListener<String, RemoteClientState> { rl ->
+            rl.value?.triggerCloseListeners()
+        }
         .build()
 
     //: ConcurrentHashMap<String, RemoteClientState> = ConcurrentHashMap()
@@ -339,6 +342,9 @@ class Kweb private constructor(
 
         try {
             val webBrowser = WebBrowser(kwebSessionId, httpRequestInfo, this)
+
+            remoteClientState.addCloseHandler { webBrowser.close() }
+
             webBrowser.htmlDocument.set(htmlDocument)
             if (debug) {
                 warnIfBlocking(maxTimeMs = kwebConfig.buildpageTimeout.toMillis(), onBlock = { thread ->

src/main/kotlin/kweb/WebBrowser.kt

@@ -37,6 +37,8 @@ class WebBrowser(val sessionId: String, val httpRequestInfo: HttpRequestInfo, va
 
     private val idCounter = AtomicInteger(0)
 
+    private val closeListeners : ConcurrentHashMap<Int, () -> Unit> = ConcurrentHashMap()
+
     /**
      * During page render, the initial HTML document will be available for modification as a
      * [JSoup Document](https://jsoup.org/) in this [AtomicReference].
@@ -229,6 +231,24 @@ class WebBrowser(val sessionId: String, val httpRequestInfo: HttpRequestInfo, va
         kweb.addCallback(sessionId, functionCall.callbackId!!, callback)
     }
 
+    fun addCloseListener(listener: () -> Unit) : Int {
+        val id = random.nextInt()
+        this.closeListeners[id] = listener
+        return id
+    }
+
+    fun removeCloseListener(id: Int) {
+        synchronized(this) {
+            this.closeListeners.remove(id)
+        }
+    }
+
+    fun close() {
+        synchronized(this) {
+            this.closeListeners.values.forEach { it() }
+        }
+    }
+
     private fun createCacheFunctionJs(cacheId: Int, functionBody: String, params: String? = null) : String {
         params?.let {
             //language=JavaScript

src/main/kotlin/kweb/client/RemoteClientState.kt

@@ -4,14 +4,38 @@ import kotlinx.serialization.encodeToString
 import kotlinx.serialization.json.Json
 import kotlinx.serialization.json.JsonElement
 import kweb.DebugInfo
+import kweb.util.random
 import java.time.Instant
+import java.util.concurrent.ConcurrentHashMap
 
 data class RemoteClientState(val id: String, @Volatile var clientConnection: ClientConnection,
                              val handlers: MutableMap<Int, (JsonElement) -> Unit> = HashMap(),
+                             val onCloseHandlers : ConcurrentHashMap<Int, () -> Unit> = ConcurrentHashMap(),
                              val debugTokens: MutableMap<String, DebugInfo> = HashMap(), var lastModified: Instant = Instant.now()) {
     fun send(message: Server2ClientMessage) {
         clientConnection.send(Json.encodeToString(message))
     }
 
+    fun addCloseHandler(handler: () -> Unit) : Int {
+        synchronized(this) {
+            val id = random.nextInt()
+            onCloseHandlers[id] = handler
+            return id
+        }
+    }
+
+    fun removeCloseHandler(id: Int) {
+        synchronized(this) {
+            onCloseHandlers.remove(id)
+        }
+    }
+
+    fun triggerCloseListeners() {
+        synchronized(this) {
+            onCloseHandlers.values.forEach { it() }
+            onCloseHandlers.clear()
+        }
+    }
+
     override fun toString() = "Remote2ClientState(id = $id)"
 }

src/main/kotlin/kweb/html/Document.kt

@@ -18,6 +18,7 @@ import kweb.util.random
  * Passed in as `doc` to the `buildPage` lambda of the [Kweb] constructor.
  */
 class Document(val receiver: WebBrowser) : EventGenerator<Document> {
+
     fun getElementById(id: String) = Element(receiver, null, "return document.getElementById(\"$id\")", id = id)
 
     val cookie = CookieReceiver(receiver)
@@ -26,7 +27,11 @@ class Document(val receiver: WebBrowser) : EventGenerator<Document> {
 
     fun body(new: (ElementCreator<BodyElement>.() -> Unit)? = null) : BodyElement {
         if (new != null) {
-            new(ElementCreator(parent = body, insertBefore = null))
+            val ec = ElementCreator(parent = body, insertBefore = null)
+            new(ec)
+            receiver.addCloseListener {
+                ec.cleanup()
+            }
         }
         return body
     }
@@ -35,7 +40,11 @@ class Document(val receiver: WebBrowser) : EventGenerator<Document> {
 
     fun head(new: (ElementCreator<HeadElement>.() -> Unit)? = null) : HeadElement {
         if (new != null) {
-            new(ElementCreator(parent = head, insertBefore = null))
+            val ec = ElementCreator(parent = head, insertBefore = null)
+            new(ec)
+            receiver.addCloseListener {
+                ec.cleanup()
+            }
         }
         return head
     }