Merge pull request #284 from ooni/stop-while-uploading

Make uploading missing results stoppable and part of the same background flow

composeApp/src/androidInstrumentedTest/kotlin/org/ooni/probe/uitesting/RunningTest.kt → composeApp/src/androidInstrumentedTest/kotlin/org/ooni/probe/uitesting/RunningTestsTest.kt

@@ -43,7 +43,7 @@ import org.ooni.probe.uitesting.helpers.wait
 import kotlin.time.Duration.Companion.minutes
 
 @RunWith(AndroidJUnit4::class)
-class RunningTest {
+class RunningTestsTest {
     @get:Rule
     val compose = createEmptyComposeRule()
 

composeApp/src/androidMain/kotlin/org/ooni/probe/background/RunWorker.kt

@@ -30,8 +30,8 @@ import org.jetbrains.compose.resources.getString
 import org.ooni.probe.AndroidApplication
 import org.ooni.probe.MainActivity
 import org.ooni.probe.R
+import org.ooni.probe.data.models.RunBackgroundState
 import org.ooni.probe.data.models.RunSpecification
-import org.ooni.probe.data.models.TestRunState
 import org.ooni.probe.di.Dependencies
 import org.ooni.probe.domain.UploadMissingMeasurements
 import org.ooni.probe.ui.primaryLight
@@ -46,6 +46,9 @@ class RunWorker(
     private val json by lazy { dependencies.json }
     private val runBackgroundTask by lazy { dependencies.runBackgroundTask }
     private val cancelCurrentTest by lazy { dependencies.cancelCurrentTest }
+    private val setRunBackgroundState by lazy {
+        dependencies.runBackgroundStateManager::updateState
+    }
 
     private val notificationManager by lazy {
         appContext.getSystemService(NotificationManager::class.java)
@@ -55,7 +58,7 @@ class RunWorker(
 
     override suspend fun getForegroundInfo(): ForegroundInfo {
         buildNotificationChannelIfNeeded()
-        return ForegroundInfo(NOTIFICATION_ID, buildNotification(TestRunState.Running()))
+        return ForegroundInfo(NOTIFICATION_ID, buildNotification(RunBackgroundState.RunningTests()))
     }
 
     override suspend fun doWork(): Result {
@@ -67,6 +70,7 @@ class RunWorker(
             work()
         } catch (e: CancellationException) {
             Logger.i("Run Worker: cancelled")
+            setRunBackgroundState { RunBackgroundState.Idle() }
         } finally {
             notificationManager.cancel(NOTIFICATION_ID)
             unregisterReceiver()
@@ -85,19 +89,17 @@ class RunWorker(
 
         runBackgroundTask(getSpecification())
             .collectLatest { state ->
-                notificationManager.notify(
-                    NOTIFICATION_ID,
-                    when (state) {
-                        is RunBackgroundTask.State.UploadingMissingResults ->
-                            buildNotification(state.state)
-
-                        is RunBackgroundTask.State.RunningTests ->
-                            buildNotification(state.state)
-
-                        is RunBackgroundTask.State.StoppingTests ->
-                            buildStoppingNotification()
-                    },
-                )
+                val notification = when (state) {
+                    is RunBackgroundState.Idle -> null
+                    is RunBackgroundState.UploadingMissingResults -> buildNotification(state.state)
+                    is RunBackgroundState.RunningTests -> buildNotification(state)
+                    is RunBackgroundState.Stopping -> buildStoppingNotification()
+                }
+                if (notification != null) {
+                    notificationManager.notify(NOTIFICATION_ID, notification)
+                } else {
+                    notificationManager.cancel(NOTIFICATION_ID)
+                }
             }
     }
 
@@ -146,23 +148,18 @@ class RunWorker(
                     ),
                 )
                     .setProgress(state.total, progress, false)
+                    .addAction(buildNotificationStopAction())
             } else {
                 setProgress(1, 0, true)
             }
         }
 
-    private suspend fun buildNotification(state: TestRunState.Running) =
+    private suspend fun buildNotification(state: RunBackgroundState.RunningTests) =
         buildNotification {
             setContentText(state.testType?.labelRes?.let { labelRes -> getString(labelRes) })
                 .setColor(state.descriptor?.color?.toArgb() ?: primaryLight.toArgb())
                 .setProgress(1000, (state.progress * 1000).roundToInt(), false)
-                .addAction(
-                    NotificationCompat.Action.Builder(
-                        null,
-                        getString(Res.string.Notification_StopTest),
-                        stopRunIntent,
-                    ).build(),
-                )
+                .addAction(buildNotificationStopAction())
         }
 
     private suspend fun buildStoppingNotification() =
@@ -188,6 +185,13 @@ class RunWorker(
         ).build()
     }
 
+    private suspend fun buildNotificationStopAction() =
+        NotificationCompat.Action.Builder(
+            null,
+            getString(Res.string.Notification_StopTest),
+            stopRunIntent,
+        ).build()
+
     private val openAppIntent
         get() = PendingIntent.getActivity(
             applicationContext,

composeApp/src/commonMain/kotlin/org/ooni/engine/Engine.kt

@@ -2,11 +2,9 @@ package org.ooni.engine
 
 import androidx.annotation.VisibleForTesting
 import co.touchlab.kermit.Logger
-import kotlinx.coroutines.async
 import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.channelFlow
 import kotlinx.coroutines.flow.flowOn
-import kotlinx.coroutines.flow.take
 import kotlinx.coroutines.isActive
 import kotlinx.serialization.encodeToString
 import kotlinx.serialization.json.Json
@@ -36,7 +34,7 @@ class Engine(
     private val isBatteryCharging: suspend () -> Boolean,
     private val platformInfo: PlatformInfo,
     private val getEnginePreferences: suspend () -> EnginePreferences,
-    private val observeCancelTestRun: () -> Flow<Unit>,
+    private val addRunCancelListener: (() -> Unit) -> Unit,
     private val backgroundContext: CoroutineContext,
 ) {
     fun startTask(
@@ -47,30 +45,21 @@ class Engine(
     ): Flow<TaskEvent> =
         channelFlow {
             val preferences = getEnginePreferences()
-            val taskSettings = buildTaskSettings(name, inputs, taskOrigin, preferences, descriptorId)
+            val taskSettings =
+                buildTaskSettings(name, inputs, taskOrigin, preferences, descriptorId)
             val settingsSerialized = json.encodeToString(taskSettings)
 
             var task: OonimkallBridge.Task? = null
             try {
                 task = bridge.startTask(settingsSerialized)
 
-                val cancelJob = async {
-                    observeCancelTestRun()
-                        .take(1)
-                        .collect {
-                            task.interrupt()
-                        }
-                }
+                addRunCancelListener { task.interrupt() }
 
                 while (!task.isDone() && isActive) {
                     val eventJson = task.waitForNextEvent()
                     val taskEventResult = json.decodeFromString<TaskEventResult>(eventJson)
                     taskEventMapper(taskEventResult)?.let { send(it) }
                 }
-
-                if (cancelJob.isActive) {
-                    cancelJob.cancel()
-                }
             } catch (e: Exception) {
                 Logger.d("Error while running task", e)
                 throw MkException(e)

composeApp/src/commonMain/kotlin/org/ooni/probe/background/RunBackgroundTask.kt

@@ -1,20 +1,26 @@
 package org.ooni.probe.background
 
 import co.touchlab.kermit.Logger
+import kotlinx.coroutines.CancellationException
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.async
 import kotlinx.coroutines.coroutineScope
 import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.channelFlow
 import kotlinx.coroutines.flow.collect
 import kotlinx.coroutines.flow.collectLatest
 import kotlinx.coroutines.flow.first
+import kotlinx.coroutines.flow.onCompletion
 import kotlinx.coroutines.flow.onEach
 import kotlinx.coroutines.flow.takeWhile
+import kotlinx.coroutines.launch
+import kotlinx.coroutines.supervisorScope
 import org.ooni.engine.models.NetworkType
 import org.ooni.probe.data.models.ResultModel
+import org.ooni.probe.data.models.RunBackgroundState
 import org.ooni.probe.data.models.RunSpecification
 import org.ooni.probe.data.models.SettingsKey
-import org.ooni.probe.data.models.TestRunState
 import org.ooni.probe.domain.UploadMissingMeasurements
 
 class RunBackgroundTask(
@@ -24,18 +30,56 @@ class RunBackgroundTask(
     private val getNetworkType: () -> NetworkType,
     private val getAutoRunSpecification: suspend () -> RunSpecification,
     private val runDescriptors: suspend (RunSpecification) -> Unit,
-    private val getCurrentTestState: () -> Flow<TestRunState>,
+    private val setRunBackgroundState: ((RunBackgroundState) -> RunBackgroundState) -> Unit,
+    private val getRunBackgroundState: () -> Flow<RunBackgroundState>,
+    private val addRunCancelListener: (() -> Unit) -> Unit,
+    private val clearRunCancelListeners: () -> Unit,
 ) {
-    operator fun invoke(spec: RunSpecification?): Flow<State> =
+    operator fun invoke(spec: RunSpecification?): Flow<RunBackgroundState> =
         channelFlow {
+            val initialState = getRunBackgroundState().first()
+            if (initialState !is RunBackgroundState.Idle) {
+                Logger.i("Background task is already running, so we won't start another one")
+                return@channelFlow
+            }
+
+            var isCancelled = false
+
             if (spec == null &&
                 getPreferenceValueByKey(SettingsKey.UPLOAD_RESULTS).first() == true
             ) {
-                uploadMissingMeasurements(null).collectLatest {
-                    send(State.UploadingMissingResults(it))
+                supervisorScope {
+                    val uploadJob = async {
+                        uploadMissingMeasurements(null)
+                            .collectLatest { uploadState ->
+                                setRunBackgroundState {
+                                    RunBackgroundState.UploadingMissingResults(uploadState)
+                                }
+                                send(RunBackgroundState.UploadingMissingResults(uploadState))
+                            }
+                    }
+
+                    addRunCancelListener {
+                        isCancelled = true
+                        if (uploadJob.isActive) uploadJob.cancel()
+                        setRunBackgroundState { RunBackgroundState.Stopping }
+                        CoroutineScope(Dispatchers.Default).launch { send(RunBackgroundState.Stopping) }
+                    }
+
+                    try {
+                        uploadJob.await()
+                    } catch (e: CancellationException) {
+                        Logger.i("Upload Missing Results: cancelled")
+                    }
                 }
             }
 
+            if (isCancelled) {
+                setRunBackgroundState { RunBackgroundState.Idle() }
+                send(RunBackgroundState.Idle())
+                return@channelFlow
+            }
+
             if (checkSkipAutoRunNotUploadedLimit().first()) {
                 Logger.i("Skipping auto-run tests: too many not-uploaded results")
                 return@channelFlow
@@ -52,34 +96,22 @@ class RunBackgroundTask(
                 }
 
                 var testStarted = false
-                getCurrentTestState()
+                getRunBackgroundState()
                     .takeWhile { state ->
-                        state is TestRunState.Running ||
-                            state is TestRunState.Stopping ||
-                            (state is TestRunState.Idle && !testStarted)
+                        state is RunBackgroundState.RunningTests ||
+                            state is RunBackgroundState.Stopping ||
+                            (state is RunBackgroundState.Idle && !testStarted)
                     }
                     .onEach { state ->
-                        if (state is TestRunState.Idle) return@onEach
+                        if (state is RunBackgroundState.Idle) return@onEach
                         testStarted = true
-                        send(
-                            if (state is TestRunState.Running) {
-                                State.RunningTests(state)
-                            } else {
-                                State.StoppingTests
-                            },
-                        )
+                        send(state)
                     }
                     .collect()
 
                 runJob.await()
             }
+        }.onCompletion {
+            clearRunCancelListeners()
         }
-
-    sealed interface State {
-        data class UploadingMissingResults(val state: UploadMissingMeasurements.State) : State
-
-        data class RunningTests(val state: TestRunState.Running) : State
-
-        data object StoppingTests : State
-    }
 }

composeApp/src/commonMain/kotlin/org/ooni/probe/data/models/TestRunState.kt → composeApp/src/commonMain/kotlin/org/ooni/probe/data/models/RunBackgroundState.kt

@@ -2,16 +2,21 @@ package org.ooni.probe.data.models
 
 import kotlinx.datetime.LocalDateTime
 import org.ooni.engine.models.TestType
+import org.ooni.probe.domain.UploadMissingMeasurements
 import kotlin.time.Duration
 import kotlin.time.Duration.Companion.seconds
 
-sealed interface TestRunState {
+sealed interface RunBackgroundState {
     data class Idle(
         val lastTestAt: LocalDateTime? = null,
         val justFinishedTest: Boolean = false,
-    ) : TestRunState
+    ) : RunBackgroundState
 
-    data class Running(
+    data class UploadingMissingResults(
+        val state: UploadMissingMeasurements.State,
+    ) : RunBackgroundState
+
+    data class RunningTests(
         val descriptor: Descriptor? = null,
         private val descriptorIndex: Int = 0,
         val testType: TestType? = null,
@@ -20,7 +25,7 @@ sealed interface TestRunState {
         private val testIndex: Int = 0,
         private val testTotal: Int = 1,
         val log: String? = "",
-    ) : TestRunState {
+    ) : RunBackgroundState {
         val estimatedTimeLeft: Duration?
             get() {
                 if (estimatedRuntimeOfDescriptors.isNullOrEmpty()) return null
@@ -45,5 +50,5 @@ sealed interface TestRunState {
             }
     }
 
-    data object Stopping : TestRunState
+    data object Stopping : RunBackgroundState
 }