Updated atomicCoroutineScope to run operations under supplied scope,

updated tests for the same

Author: sacOO7

chat-android/src/main/java/com/ably/chat/AtomicCoroutineScope.kt

@@ -2,23 +2,22 @@ package com.ably.chat
 
 import io.ably.annotation.Experimental
 import java.util.concurrent.PriorityBlockingQueue
-import java.util.concurrent.atomic.AtomicInteger
 import kotlinx.coroutines.CompletableDeferred
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.Dispatchers
-import kotlinx.coroutines.coroutineScope
 import kotlinx.coroutines.launch
 
 /**
- * AtomicCoroutineScope makes sure all operations are atomic and run with given priority.
- * Accepts scope as an optional parameter to run atomic operations under the given scope.
+ * AtomicCoroutineScope is a thread safe wrapper to run multiple operations mutually exclusive.
+ * All operations are atomic and run with given priority.
+ * Accepts scope as a constructor parameter to run operations under the given scope.
  * See [Kotlin Dispatchers](https://kt.academy/article/cc-dispatchers) for more information.
  */
 class AtomicCoroutineScope(private val scope: CoroutineScope = CoroutineScope(Dispatchers.Default)) {
 
     private val sequentialScope: CoroutineScope = CoroutineScope(Dispatchers.Default.limitedParallelism(1))
 
-    private class Job private constructor(
+    private class Job(
         private val priority: Int,
         val coroutineBlock: suspend CoroutineScope.() -> Any,
         val deferredResult: CompletableDeferred<Any>,
@@ -31,31 +30,28 @@ class AtomicCoroutineScope(private val scope: CoroutineScope = CoroutineScope(Di
             }
             return this.priority.compareTo(other.priority)
         }
-
-        companion object {
-            var counter: AtomicInteger = AtomicInteger()
-            fun create(priority: Int, coroutineBlock: suspend CoroutineScope.() -> Any, deferredResult: CompletableDeferred<Any>): Job {
-                return Job(priority, coroutineBlock, deferredResult, counter.getAndIncrement())
-            }
-        }
     }
 
-    private val jobs: PriorityBlockingQueue<Job> = PriorityBlockingQueue()
-    private var isRunning = false
+    private val jobs: PriorityBlockingQueue<Job> = PriorityBlockingQueue() // Accessed from both sequentialScope and async method
+    private var isRunning = false // Only accessed from sequentialScope
+    private var queueCounter = 0 // Only accessed from synchronized async method
 
     /**
      * @param priority Defines priority for the operation execution.
      * @param coroutineBlock Suspended function that needs to be executed mutually exclusive under given scope.
      */
+    @Synchronized
     fun <T : Any>async(priority: Int = 0, coroutineBlock: suspend CoroutineScope.() -> T): CompletableDeferred<T> {
         val deferredResult = CompletableDeferred<Any>()
-        jobs.add(Job.create(priority, coroutineBlock, deferredResult))
+        jobs.add(Job(priority, coroutineBlock, deferredResult, queueCounter++))
         sequentialScope.launch {
             if (!isRunning) {
                 isRunning = true
                 while (jobs.isNotEmpty()) {
                     val job = jobs.poll()
-                    safeExecute(job)
+                    job?.let {
+                        safeExecute(it)
+                    }
                 }
                 isRunning = false
             }
@@ -65,32 +61,22 @@ class AtomicCoroutineScope(private val scope: CoroutineScope = CoroutineScope(Di
         return deferredResult as CompletableDeferred<T>
     }
 
-    private suspend fun safeExecute(job: Job?) {
-        job?.let {
-            runCatching {
-                scope.launch {
-                    try {
-                        val result = it.coroutineBlock(this)
-                        it.deferredResult.complete(result)
-                    } catch (t: Throwable) {
-                        it.deferredResult.completeExceptionally(t)
-                    }
-                }.join()
-            }
+    private suspend fun safeExecute(job: Job) {
+        runCatching {
+            scope.launch {
+                try {
+                    val result = job.coroutineBlock(this)
+                    job.deferredResult.complete(result)
+                } catch (t: Throwable) {
+                    job.deferredResult.completeExceptionally(t)
+                }
+            }.join()
+        }.onFailure {
+            job.deferredResult.completeExceptionally(it)
         }
     }
 
     @Experimental
     val finishedProcessing: Boolean
         get() = jobs.isEmpty() && !isRunning
-
-    /**
-     * Clears queuedJobs
-     */
-    fun cancel(message: String? = "Atomic coroutine scope cancelled", cause: Throwable? = null) {
-        jobs.clear()
-        Job.counter.set(0)
-        // Once sequentialScope.cancel called, AtomicCoroutineScope can't be reused
-        // So, once all jobs are executed, it should be garbage collected.
-    }
 }

chat-android/src/test/java/com/ably/chat/AtomicCoroutineScopeTest.kt

@@ -2,15 +2,20 @@ package com.ably.chat
 
 import io.ably.lib.types.AblyException
 import io.ably.lib.types.ErrorInfo
+import java.util.concurrent.LinkedBlockingQueue
 import kotlin.time.DurationUnit
 import kotlin.time.toDuration
+import kotlinx.coroutines.CompletableDeferred
+import kotlinx.coroutines.CoroutineName
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.Deferred
 import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.awaitAll
 import kotlinx.coroutines.delay
+import kotlinx.coroutines.launch
 import kotlinx.coroutines.runBlocking
 import kotlinx.coroutines.test.runTest
+import kotlinx.coroutines.withContext
 import org.hamcrest.CoreMatchers.containsString
 import org.junit.Assert
 import org.junit.Test
@@ -69,8 +74,8 @@ class AtomicCoroutineScopeTest {
                 }
                 operationInProgress = true
                 delay((200..800).random().toDuration(DurationUnit.MILLISECONDS))
-                operationInProgress = false
                 val returnValue = counter++
+                operationInProgress = false
                 return@async returnValue
             }
             deferredResults.add(result)
@@ -86,24 +91,63 @@ class AtomicCoroutineScopeTest {
     }
 
     @Test
-    fun `should perform mutually exclusive operations with custom scope`() = runTest {
-        val sequentialScope = CoroutineScope(Dispatchers.Default.limitedParallelism(1))
+    fun `Concurrently perform mutually exclusive operations`() = runTest {
+        val atomicCoroutineScope = AtomicCoroutineScope()
+        val deferredResults = LinkedBlockingQueue<CompletableDeferred<Unit>>()
+
+        var operationInProgress = false
+        var counter = 0
+        val countedValues = mutableListOf<Int>()
+
+        // Concurrently schedule 100000 jobs from multiple threads
+        withContext(Dispatchers.IO) {
+            repeat(100000) {
+                launch {
+                    val result = atomicCoroutineScope.async {
+                        if (operationInProgress) {
+                            error("Can't perform operation when other operation is going on")
+                        }
+                        operationInProgress = true
+                        countedValues.add(counter++)
+                        operationInProgress = false
+                    }
+                    deferredResults.add(result)
+                }
+            }
+        }
+
+        Assert.assertFalse(atomicCoroutineScope.finishedProcessing)
+        assertWaiter { deferredResults.size == 100000 }
+
+        deferredResults.awaitAll()
+        assertWaiter { atomicCoroutineScope.finishedProcessing }
+        Assert.assertEquals((0..99999).toList(), countedValues)
+    }
+
+    @Test
+    fun `should perform mutually exclusive operations with custom named scope`() = runTest {
+        val sequentialScope = CoroutineScope(Dispatchers.Default.limitedParallelism(1) + CoroutineName("roomId"))
         val atomicCoroutineScope = AtomicCoroutineScope(sequentialScope)
         val deferredResults = mutableListOf<Deferred<Int>>()
+
         val contexts = mutableListOf<String>()
+        val contextNames = mutableListOf<String>()
+
         var operationInProgress = false
         var counter = 0
 
         repeat(10) {
             val result = atomicCoroutineScope.async {
-                contexts.add(this.coroutineContext.toString())
                 if (operationInProgress) {
                     error("Can't perform operation when other operation is going on")
                 }
                 operationInProgress = true
+                contexts.add(coroutineContext.toString())
+                contextNames.add(coroutineContext[CoroutineName]!!.name)
+
                 delay((200..800).random().toDuration(DurationUnit.MILLISECONDS))
-                operationInProgress = false
                 val returnValue = counter++
+                operationInProgress = false
                 return@async returnValue
             }
             deferredResults.add(result)
@@ -113,6 +157,7 @@ class AtomicCoroutineScopeTest {
         val results = deferredResults.awaitAll()
         repeat(10) {
             Assert.assertEquals(it, results[it])
+            Assert.assertEquals("roomId", contextNames[it])
             Assert.assertThat(contexts[it], containsString("Dispatchers.Default.limitedParallelism(1)"))
         }
         Assert.assertTrue(atomicCoroutineScope.finishedProcessing)
@@ -138,14 +183,14 @@ class AtomicCoroutineScopeTest {
         // Add more jobs, will be processed based on priority
         repeat(10) {
             val result = atomicCoroutineScope.async(10 - it) {
-                contexts.add(this.coroutineContext.toString())
                 if (operationInProgress) {
                     error("Can't perform operation when other operation is going on")
                 }
                 operationInProgress = true
+                contexts.add(this.coroutineContext.toString())
                 delay((200..800).random().toDuration(DurationUnit.MILLISECONDS))
-                operationInProgress = false
                 val returnValue = counter++
+                operationInProgress = false
                 return@async returnValue
             }
             deferredResults.add(result)
@@ -161,6 +206,46 @@ class AtomicCoroutineScopeTest {
     }
 
     @Test
-    fun `reuse AtomicCoroutineScope once cancelled`() = runTest {
+    fun `Concurrently execute mutually exclusive operations with given priority`() = runTest {
+        val atomicCoroutineScope = AtomicCoroutineScope()
+        val deferredResults = LinkedBlockingQueue<Deferred<Unit>>()
+
+        var operationInProgress = false
+        val processedValues = mutableListOf<Int>()
+
+//         This will start first internal operation
+        deferredResults.add(
+            atomicCoroutineScope.async {
+                delay(1000)
+                processedValues.add(1000)
+                Unit
+            },
+        )
+
+        // Add more jobs, will be processed based on priority
+        // Concurrently schedule 1000 jobs with incremental priority from multiple threads
+        withContext(Dispatchers.IO) {
+            repeat(1000) {
+                launch {
+                    val result = atomicCoroutineScope.async(1000 - it) {
+                        if (operationInProgress) {
+                            error("Can't perform operation when other operation is going on")
+                        }
+                        operationInProgress = true
+                        processedValues.add(it)
+                        operationInProgress = false
+                    }
+                    deferredResults.add(result)
+                }
+            }
+        }
+
+        Assert.assertFalse(atomicCoroutineScope.finishedProcessing)
+        deferredResults.awaitAll()
+        val expectedResults = (1000 downTo 0).toList()
+        repeat(1001) {
+            Assert.assertEquals(expectedResults[it], processedValues[it])
+        }
+        Assert.assertTrue(atomicCoroutineScope.finishedProcessing)
     }
 }