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,32 @@ 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> {
+    /** Enforcing the rule 2.7 of the Reactive Streams spec:
+     * [Subscription.request] and [Subscription.cancel] should be serialized.
+     * The write lock is for [Subscription.cancel] operation, whereas the read lock is for
+     * [Subscription.request].
+     *
+     * The reason for this being so complex is that [Mono] proxies the calls to the actual [Subscription] and
+     * [Subscription.request] that we perform in [Subscriber.onSubscribe] below only modifies the internal state of
+     * the proxy. Only later, after [Subscriber.onSubscribe] is finished, the actual [Subscriber] receives the request.
+     *
+     * However, in general, it is not an error for [Publisher.subscribe] and [Subscriber.onSubscribe] to happen in
+     * different threads, so we need to make sure that `s.request` and `subscribe(subscriber)` are allowed to happen at
+     * the same time. This is why we need a read lock.
+     */
+    val rwLock = ReentrantReadWriteLock()
+    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 {
+                rwLock.writeLock().withLock {
+                    s.cancel()
+                }
+            }
+            rwLock.readLock().withLock {
+                s.request(Long.MAX_VALUE)
+            }
         }
 
         override fun onComplete() {
@@ -62,7 +84,10 @@ public suspend fun <T> Mono<T>.awaitSingleOrNull(): T? = suspendCancellableCorou
         }
 
         override fun onError(error: Throwable) { cont.resumeWithException(error) }
-    })
+    }
+    rwLock.readLock().withLock {
+        injectCoroutineContext(cont.context).subscribe(subscriber)
+    }
 }
 
 /**

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

@@ -0,0 +1,44 @@
+/*
+ * 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
+    fun testAwaitCancellationOrder() = runBlocking {
+        repeat(iterations) {
+            val job = launch(Dispatchers.Default) {
+                TestMono.awaitSingleOrNull()
+            }
+            job.cancelAndJoin()
+        }
+    }
+
+   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()
+                }
+            })
+        }
+    }
+}