Introduce EXACTLY_ONCE contracts to coroutineScope, supervisorScope, …

…withContext, runBlocking, withTimeout and select (Kotlin#2030)

integration/kotlinx-coroutines-jdk8/src/time/Time.kt

@@ -1,13 +1,16 @@
 /*
  * Copyright 2016-2020 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license.
  */
+@file:OptIn(ExperimentalContracts::class)
+
 package kotlinx.coroutines.time
 
 import kotlinx.coroutines.*
 import kotlinx.coroutines.flow.*
 import kotlinx.coroutines.selects.*
 import java.time.*
 import java.time.temporal.*
+import kotlin.contracts.*
 
 /**
  * "java.time" adapter method for [kotlinx.coroutines.delay].
@@ -35,8 +38,12 @@ public fun <R> SelectBuilder<R>.onTimeout(duration: Duration, block: suspend ()
 /**
  * "java.time" adapter method for [kotlinx.coroutines.withTimeout].
  */
-public suspend fun <T> withTimeout(duration: Duration, block: suspend CoroutineScope.() -> T): T =
- kotlinx.coroutines.withTimeout(duration.coerceToMillis(), block)
+public suspend fun <T> withTimeout(duration: Duration, block: suspend CoroutineScope.() -> T): T {
+ contract {
+ callsInPlace(block, InvocationKind.EXACTLY_ONCE)
+ }
+ return kotlinx.coroutines.withTimeout(duration.coerceToMillis(), block)
+}
 
 /**
  * "java.time" adapter method for [kotlinx.coroutines.withTimeoutOrNull].

kotlinx-coroutines-core/common/src/Builders.common.kt

@@ -4,13 +4,15 @@
 
 @file:JvmMultifileClass
 @file:JvmName("BuildersKt")
+@file:OptIn(ExperimentalContracts::class)
 
 package kotlinx.coroutines
 
 import kotlinx.atomicfu.*
 import kotlinx.coroutines.internal.*
 import kotlinx.coroutines.intrinsics.*
 import kotlinx.coroutines.selects.*
+import kotlin.contracts.*
 import kotlin.coroutines.*
 import kotlin.coroutines.intrinsics.*
 import kotlin.jvm.*
@@ -134,31 +136,36 @@ private class LazyDeferredCoroutine<T>(
 public suspend fun <T> withContext(
  context: CoroutineContext,
  block: suspend CoroutineScope.() -> T
-): T = suspendCoroutineUninterceptedOrReturn sc@ { uCont ->
- // compute new context
- val oldContext = uCont.context
- val newContext = oldContext + context
- // always check for cancellation of new context
- newContext.checkCompletion()
- // FAST PATH #1 -- new context is the same as the old one
- if (newContext === oldContext) {
- val coroutine = ScopeCoroutine(newContext, uCont)
- return@sc coroutine.startUndispatchedOrReturn(coroutine, block)
+): T {
+ contract {
+ callsInPlace(block, InvocationKind.EXACTLY_ONCE)
  }
- // FAST PATH #2 -- the new dispatcher is the same as the old one (something else changed)
- // `equals` is used by design (see equals implementation is wrapper context like ExecutorCoroutineDispatcher)
- if (newContext[ContinuationInterceptor] == oldContext[ContinuationInterceptor]) {
- val coroutine = UndispatchedCoroutine(newContext, uCont)
- // There are changes in the context, so this thread needs to be updated
- withCoroutineContext(newContext, null) {
+ return suspendCoroutineUninterceptedOrReturn sc@ { uCont ->
+ // compute new context
+ val oldContext = uCont.context
+ val newContext = oldContext + context
+ // always check for cancellation of new context
+ newContext.checkCompletion()
+ // FAST PATH #1 -- new context is the same as the old one
+ if (newContext === oldContext) {
+ val coroutine = ScopeCoroutine(newContext, uCont)
  return@sc coroutine.startUndispatchedOrReturn(coroutine, block)
  }
+ // FAST PATH #2 -- the new dispatcher is the same as the old one (something else changed)
+ // `equals` is used by design (see equals implementation is wrapper context like ExecutorCoroutineDispatcher)
+ if (newContext[ContinuationInterceptor] == oldContext[ContinuationInterceptor]) {
+ val coroutine = UndispatchedCoroutine(newContext, uCont)
+ // There are changes in the context, so this thread needs to be updated
+ withCoroutineContext(newContext, null) {
+ return@sc coroutine.startUndispatchedOrReturn(coroutine, block)
+ }
+ }
+ // SLOW PATH -- use new dispatcher
+ val coroutine = DispatchedCoroutine(newContext, uCont)
+ coroutine.initParentJob()
+ block.startCoroutineCancellable(coroutine, coroutine)
+ coroutine.getResult()
  }
- // SLOW PATH -- use new dispatcher
- val coroutine = DispatchedCoroutine(newContext, uCont)
- coroutine.initParentJob()
- block.startCoroutineCancellable(coroutine, coroutine)
- coroutine.getResult()
 }
 
 /**

kotlinx-coroutines-core/common/src/CoroutineScope.kt

@@ -1,11 +1,13 @@
 /*
  * Copyright 2016-2020 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license.
  */
+@file:OptIn(ExperimentalContracts::class)
 
 package kotlinx.coroutines
 
 import kotlinx.coroutines.internal.*
 import kotlinx.coroutines.intrinsics.*
+import kotlin.contracts.*
 import kotlin.coroutines.*
 import kotlin.coroutines.intrinsics.*
 
@@ -183,11 +185,15 @@ public object GlobalScope : CoroutineScope {
  * or may throw a corresponding unhandled [Throwable] if there is any unhandled exception in this scope
  * (for example, from a crashed coroutine that was started with [launch][CoroutineScope.launch] in this scope).
  */
-public suspend fun <R> coroutineScope(block: suspend CoroutineScope.() -> R): R =
- suspendCoroutineUninterceptedOrReturn { uCont ->
+public suspend fun <R> coroutineScope(block: suspend CoroutineScope.() -> R): R {
+ contract {
+ callsInPlace(block, InvocationKind.EXACTLY_ONCE)
+ }
+ return suspendCoroutineUninterceptedOrReturn { uCont ->
  val coroutine = ScopeCoroutine(uCont.context, uCont)
  coroutine.startUndispatchedOrReturn(coroutine, block)
  }
+}
 
 /**
  * Creates a [CoroutineScope] that wraps the given coroutine [context].

kotlinx-coroutines-core/common/src/Supervisor.kt

@@ -1,13 +1,14 @@
 /*
  * Copyright 2016-2020 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license.
  */
-
+@file:OptIn(ExperimentalContracts::class)
 @file:Suppress("DEPRECATION_ERROR")
 
 package kotlinx.coroutines
 
 import kotlinx.coroutines.internal.*
 import kotlinx.coroutines.intrinsics.*
+import kotlin.contracts.*
 import kotlin.coroutines.*
 import kotlin.coroutines.intrinsics.*
 import kotlin.jvm.*
@@ -47,11 +48,15 @@ public fun SupervisorJob0(parent: Job? = null) : Job = SupervisorJob(parent)
  * A failure of the scope itself (exception thrown in the [block] or cancellation) fails the scope with all its children,
  * but does not cancel parent job.
  */
-public suspend fun <R> supervisorScope(block: suspend CoroutineScope.() -> R): R =
- suspendCoroutineUninterceptedOrReturn { uCont ->
+public suspend fun <R> supervisorScope(block: suspend CoroutineScope.() -> R): R {
+ contract {
+ callsInPlace(block, InvocationKind.EXACTLY_ONCE)
+ }
+ return suspendCoroutineUninterceptedOrReturn { uCont ->
  val coroutine = SupervisorCoroutine(uCont.context, uCont)
  coroutine.startUndispatchedOrReturn(coroutine, block)
  }
+}
 
 private class SupervisorJobImpl(parent: Job?) : JobImpl(parent) {
  override fun childCancelled(cause: Throwable): Boolean = false

kotlinx-coroutines-core/common/src/Timeout.kt

@@ -1,12 +1,14 @@
 /*
  * Copyright 2016-2020 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license.
  */
+@file:OptIn(ExperimentalContracts::class)
 
 package kotlinx.coroutines
 
 import kotlinx.coroutines.internal.*
 import kotlinx.coroutines.intrinsics.*
 import kotlinx.coroutines.selects.*
+import kotlin.contracts.*
 import kotlin.coroutines.*
 import kotlin.coroutines.intrinsics.*
 import kotlin.jvm.*
@@ -27,6 +29,9 @@ import kotlin.time.*
  * @param timeMillis timeout time in milliseconds.
  */
 public suspend fun <T> withTimeout(timeMillis: Long, block: suspend CoroutineScope.() -> T): T {
+ contract {
+ callsInPlace(block, InvocationKind.EXACTLY_ONCE)
+ }
  if (timeMillis <= 0L) throw TimeoutCancellationException("Timed out immediately")
  return suspendCoroutineUninterceptedOrReturn { uCont ->
  setupTimeout(TimeoutCoroutine(timeMillis, uCont), block)
@@ -46,8 +51,12 @@ public suspend fun <T> withTimeout(timeMillis: Long, block: suspend CoroutineSco
  * Implementation note: how the time is tracked exactly is an implementation detail of the context's [CoroutineDispatcher].
  */
 @ExperimentalTime
-public suspend fun <T> withTimeout(timeout: Duration, block: suspend CoroutineScope.() -> T): T =
- withTimeout(timeout.toDelayMillis(), block)
+public suspend fun <T> withTimeout(timeout: Duration, block: suspend CoroutineScope.() -> T): T {
+ contract {
+ callsInPlace(block, InvocationKind.EXACTLY_ONCE)
+ }
+ return withTimeout(timeout.toDelayMillis(), block)
+}
 
 /**
  * Runs a given suspending block of code inside a coroutine with a specified [timeout][timeMillis] and returns

kotlinx-coroutines-core/common/src/selects/Select.kt

@@ -1,6 +1,7 @@
 /*
  * Copyright 2016-2020 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license.
  */
+@file:OptIn(ExperimentalContracts::class)
 
 package kotlinx.coroutines.selects
 
@@ -10,6 +11,7 @@ import kotlinx.coroutines.channels.*
 import kotlinx.coroutines.internal.*
 import kotlinx.coroutines.intrinsics.*
 import kotlinx.coroutines.sync.*
+import kotlin.contracts.*
 import kotlin.coroutines.*
 import kotlin.coroutines.intrinsics.*
 import kotlin.jvm.*
@@ -199,8 +201,11 @@ public interface SelectInstance<in R> {
  * Note that this function does not check for cancellation when it is not suspended.
  * Use [yield] or [CoroutineScope.isActive] to periodically check for cancellation in tight loops if needed.
  */
-public suspend inline fun <R> select(crossinline builder: SelectBuilder<R>.() -> Unit): R =
- suspendCoroutineUninterceptedOrReturn { uCont ->
+public suspend inline fun <R> select(crossinline builder: SelectBuilder<R>.() -> Unit): R {
+ contract {
+ callsInPlace(builder, InvocationKind.EXACTLY_ONCE)
+ }
+ return suspendCoroutineUninterceptedOrReturn { uCont ->
  val scope = SelectBuilderImpl(uCont)
  try {
  builder(scope)
@@ -209,6 +214,7 @@ public suspend inline fun <R> select(crossinline builder: SelectBuilder<R>.() ->
  }
  scope.getResult()
  }
+}
 
 
 @SharedImmutable

kotlinx-coroutines-core/common/test/BuilderContractsTest.kt

@@ -0,0 +1,52 @@
+/*
+ * Copyright 2016-2020 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license.
+ */
+
+package kotlinx.coroutines
+
+import kotlinx.coroutines.selects.*
+import kotlin.test.*
+
+class BuilderContractsTest : TestBase() {
+
+ @Test
+ fun testContracts() = runTest {
+ // Coroutine scope
+ val cs: Int
+ coroutineScope {
+ cs = 42
+ }
+ consume(cs)
+
+ // Supervisor scope
+ val svs: Int
+ supervisorScope {
+ svs = 21
+ }
+ consume(svs)
+
+ // with context scope
+ val wctx: Int
+ withContext(Dispatchers.Unconfined) {
+ wctx = 239
+ }
+ consume(wctx)
+
+ val wt: Int
+ withTimeout(Long.MAX_VALUE) {
+ wt = 123
+ }
+ consume(wt)
+
+ val s: Int
+ select<Unit> {
+ s = 42
+ Job().apply { complete() }.onJoin {}
+ }
+ consume(s)
+ }
+
+ private fun consume(a: Int) {
+ a.hashCode() // BE codegen verification
+ }
+}

kotlinx-coroutines-core/jvm/src/Builders.kt

@@ -4,10 +4,12 @@
 
 @file:JvmMultifileClass
 @file:JvmName("BuildersKt")
+@file:OptIn(ExperimentalContracts::class)
 
 package kotlinx.coroutines
 
 import java.util.concurrent.locks.*
+import kotlin.contracts.*
 import kotlin.coroutines.*
 
 /**
@@ -34,6 +36,9 @@ import kotlin.coroutines.*
  */
 @Throws(InterruptedException::class)
 public fun <T> runBlocking(context: CoroutineContext = EmptyCoroutineContext, block: suspend CoroutineScope.() -> T): T {
+ contract {
+ callsInPlace(block, InvocationKind.EXACTLY_ONCE)
+ }
  val currentThread = Thread.currentThread()
  val contextInterceptor = context[ContinuationInterceptor]
  val eventLoop: EventLoop?

kotlinx-coroutines-core/jvm/test/RunBlockingTest.kt

@@ -162,4 +162,13 @@ class RunBlockingTest : TestBase() {
 
  handle.dispose()
  }
+
+ @Test
+ fun testContract() {
+ val rb: Int
+ runBlocking {
+ rb = 42
+ }
+ rb.hashCode() // unused
+ }
 }