From cd8aaf04b324b0dfeac8e38df4b6ddd88b5b2933 Mon Sep 17 00:00:00 2001
From: firemaples <firemaples@gmail.com>
Date: Sun, 20 Aug 2023 19:32:14 +0900
Subject: [PATCH] Make screen capture cancellable

---
 .../onscreenocr/screenshot/ScreenExtractor.kt          | 10 ++++++++--
 1 file changed, 8 insertions(+), 2 deletions(-)

diff --git a/main/src/main/java/tw/firemaples/onscreenocr/screenshot/ScreenExtractor.kt b/main/src/main/java/tw/firemaples/onscreenocr/screenshot/ScreenExtractor.kt
index c39721c5..1de7c750 100644
--- a/main/src/main/java/tw/firemaples/onscreenocr/screenshot/ScreenExtractor.kt
+++ b/main/src/main/java/tw/firemaples/onscreenocr/screenshot/ScreenExtractor.kt
@@ -18,6 +18,7 @@ import android.os.Handler
 import android.os.HandlerThread
 import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.TimeoutCancellationException
+import kotlinx.coroutines.suspendCancellableCoroutine
 import kotlinx.coroutines.withContext
 import kotlinx.coroutines.withTimeout
 import tw.firemaples.onscreenocr.R
@@ -231,12 +232,13 @@ object ScreenExtractor {
     }
 
     private suspend fun ImageReader.awaitForBitmap(screenSize: Point): Bitmap? =
-        suspendCoroutine {
+        suspendCancellableCoroutine {
+            logger.debug("suspendCancellableCoroutine: Started")
             var counter = 0
             val resumed = AtomicBoolean(false)
             setOnImageAvailableListener({ reader ->
                 logger.info("onImageAvailable()")
-                if (resumed.get()){
+                if (resumed.get()) {
                     reader.setOnImageAvailableListener(null, null)
                     return@setOnImageAvailableListener
                 }
@@ -270,6 +272,10 @@ object ScreenExtractor {
                     image?.close()
                 }
             }, handler)
+            it.invokeOnCancellation {
+                resumed.set(true)
+                setOnImageAvailableListener(null, null)
+            }
         }
 
     @Throws(IllegalArgumentException::class)