Fix CompositeMutex concurrent locking

app/src/main/java/org/koitharu/kotatsu/utils/CompositeMutex.kt

@@ -1,6 +1,8 @@
 package org.koitharu.kotatsu.utils
 
 import kotlinx.coroutines.CancellableContinuation
+import kotlinx.coroutines.currentCoroutineContext
+import kotlinx.coroutines.isActive
 import kotlinx.coroutines.suspendCancellableCoroutine
 import kotlinx.coroutines.sync.Mutex
 import kotlinx.coroutines.sync.withLock
@@ -32,11 +34,13 @@ class CompositeMutex<T : Any> : Set<T> {
 	}
 
 	suspend fun lock(element: T) {
-		waitForRemoval(element)
-		mutex.withLock {
-			val lastValue = data.put(element, LinkedList<CancellableContinuation<Unit>>())
-			check(lastValue == null) {
-				"CompositeMutex is double-locked for $element"
+		while (currentCoroutineContext().isActive) {
+			waitForRemoval(element)
+			mutex.withLock {
+				if (data[element] == null) {
+					data[element] = LinkedList<CancellableContinuation<Unit>>()
+					return
+				}
 			}
 		}
 	}

app/src/test/java/org/koitharu/kotatsu/utils/CompositeMutexTest.kt

@@ -0,0 +1,39 @@
+package org.koitharu.kotatsu.utils
+
+import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.launch
+import kotlinx.coroutines.test.runTest
+import kotlinx.coroutines.withTimeoutOrNull
+import kotlinx.coroutines.yield
+import org.junit.Assert.assertNull
+import org.junit.Test
+
+class CompositeMutexTest {
+
+	@Test
+	fun testSingleLock() = runTest {
+		val mutex = CompositeMutex<Int>()
+		mutex.lock(1)
+		mutex.lock(2)
+		mutex.unlock(1)
+		assert(mutex.size == 1)
+		mutex.unlock(2)
+		assert(mutex.isEmpty())
+	}
+
+	@Test
+	fun testDoubleLock() = runTest {
+		val mutex = CompositeMutex<Int>()
+		repeat(2) {
+			launch(Dispatchers.Default) {
+				mutex.lock(1)
+			}
+		}
+		yield()
+		mutex.unlock(1)
+		val tryLock = withTimeoutOrNull(1000) {
+			mutex.lock(1)
+		}
+		assertNull(tryLock)
+	}
+}