Fix Mono.awaitSingleOrNull non-compliance with 2.7

reactive/kotlinx-coroutines-reactive/test/AwaitCancellationStressTest.kt

@@ -1,5 +1,5 @@
 /*
- * Copyright 2016-2020 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license.
+ * Copyright 2016-2022 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license.
  */
 
 package kotlinx.coroutines.reactive

reactive/kotlinx-coroutines-reactor/src/Mono.kt

@@ -13,6 +13,8 @@ import reactor.core.*
 import reactor.core.publisher.*
 import kotlin.coroutines.*
 import kotlinx.coroutines.internal.*
+import java.util.concurrent.locks.*
+import kotlin.concurrent.*
 
 /**
  * Creates a cold [mono][Mono] that runs a given [block] in a coroutine and emits its result.
@@ -44,12 +46,18 @@ public fun <T> mono(
  * function immediately cancels its [Subscription] and resumes with [CancellationException].
  */
 public suspend fun <T> Mono<T>.awaitSingleOrNull(): T? = suspendCancellableCoroutine { cont ->
-    injectCoroutineContext(cont.context).subscribe(object : Subscriber<T> {
+    val subscriber = object : Subscriber<T> {
         private var seenValue = false
 
         override fun onSubscribe(s: Subscription) {
-            cont.invokeOnCancellation { s.cancel() }
-            s.request(Long.MAX_VALUE)
+            cont.invokeOnCancellation {
+                withSubscriptionLock {
+                    s.cancel()
+                }
+            }
+            withSubscriptionLock {
+                s.request(Long.MAX_VALUE)
+            }
         }
 
         override fun onComplete() {
@@ -62,7 +70,16 @@ public suspend fun <T> Mono<T>.awaitSingleOrNull(): T? = suspendCancellableCorou
         }
 
         override fun onError(error: Throwable) { cont.resumeWithException(error) }
-    })
+
+        /**
+         * Enforce rule 2.7: [Subscription.request] and [Subscription.cancel] must be executed serially
+         */
+        @Synchronized
+        private fun withSubscriptionLock(block: () -> Unit) {
+            block()
+        }
+    }
+    injectCoroutineContext(cont.context).subscribe(subscriber)
 }
 
 /**

reactive/kotlinx-coroutines-reactor/test/MonoAwaitCancellationStressTest.kt

@@ -0,0 +1,45 @@
+/*
+ * Copyright 2016-2022 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license.
+ */
+
+package kotlinx.coroutines.reactor
+
+import kotlinx.coroutines.*
+import org.junit.*
+import org.reactivestreams.*
+import reactor.core.*
+import reactor.core.publisher.*
+import java.util.concurrent.locks.*
+
+class MonoAwaitCancellationStressTest {
+    private val iterations = 10_000 * stressTestMultiplier
+
+    @Test
+    @Ignore // Will be able to test this after https://github.com/reactor/reactor-core/issues/3117 is fixed
+    fun testAwaitCancellationOrder() = runBlocking {
+        repeat(iterations) {
+            val job = launch(Dispatchers.Default) {
+                TestMono.awaitSingleOrNull()
+            }
+            job.cancelAndJoin()
+        }
+    }
+
+    private object TestMono: Mono<Int>() {
+        override fun subscribe(s: CoreSubscriber<in Int>) {
+            val lock = ReentrantLock()
+            s.onSubscribe(object : Subscription {
+                override fun request(n: Long) {
+                    check(lock.tryLock())
+                    s.onNext(42)
+                    lock.unlock()
+                }
+
+                override fun cancel() {
+                    check(lock.tryLock())
+                    lock.unlock()
+                }
+            })
+        }
+    }
+}