From 9eaa9c647c6b10911b34ab2cd10a9aa4bc08b713 Mon Sep 17 00:00:00 2001
From: Vsevolod Tolstopyatov <qwwdfsad@gmail.com>
Date: Mon, 19 Oct 2020 10:17:16 -0700
Subject: [PATCH] Introduce CoroutineContext.job extension (#2312)

Fixes #2159
---
 .../api/kotlinx-coroutines-core.api                 |  1 +
 kotlinx-coroutines-core/common/src/Job.kt           |  9 +++++++++
 .../{EnsureActiveTest.kt => JobExtensionsTest.kt}   | 13 ++++++++++++-
 kotlinx-coroutines-core/jvm/src/Interruptible.kt    |  3 +--
 .../jvm/src/debug/internal/DebugProbesImpl.kt       |  2 +-
 5 files changed, 24 insertions(+), 4 deletions(-)
 rename kotlinx-coroutines-core/common/test/{EnsureActiveTest.kt => JobExtensionsTest.kt} (81%)

diff --git a/kotlinx-coroutines-core/api/kotlinx-coroutines-core.api b/kotlinx-coroutines-core/api/kotlinx-coroutines-core.api
index bb1c0f36ab..c3eddb98b0 100644
--- a/kotlinx-coroutines-core/api/kotlinx-coroutines-core.api
+++ b/kotlinx-coroutines-core/api/kotlinx-coroutines-core.api
@@ -394,6 +394,7 @@ public final class kotlinx/coroutines/JobKt {
 	public static final fun cancelFutureOnCompletion (Lkotlinx/coroutines/Job;Ljava/util/concurrent/Future;)Lkotlinx/coroutines/DisposableHandle;
 	public static final fun ensureActive (Lkotlin/coroutines/CoroutineContext;)V
 	public static final fun ensureActive (Lkotlinx/coroutines/Job;)V
+	public static final fun getJob (Lkotlin/coroutines/CoroutineContext;)Lkotlinx/coroutines/Job;
 	public static final fun isActive (Lkotlin/coroutines/CoroutineContext;)Z
 }
 
diff --git a/kotlinx-coroutines-core/common/src/Job.kt b/kotlinx-coroutines-core/common/src/Job.kt
index 754fa43ece..2e05635a29 100644
--- a/kotlinx-coroutines-core/common/src/Job.kt
+++ b/kotlinx-coroutines-core/common/src/Job.kt
@@ -634,6 +634,15 @@ public fun CoroutineContext.cancelChildren(cause: CancellationException? = null)
 @Deprecated(level = DeprecationLevel.HIDDEN, message = "Since 1.2.0, binary compatibility with versions <= 1.1.x")
 public fun CoroutineContext.cancelChildren(): Unit = cancelChildren(null)
 
+/**
+ * Retrieves the current [Job] instance from the given [CoroutineContext] or
+ * throws [IllegalStateException] if no job is present in the context.
+ *
+ * This method is a short-cut for `coroutineContext[Job]!!` and should be used only when it is known in advance that
+ * the context does have instance of the job in it.
+ */
+public val CoroutineContext.job: Job get() = get(Job) ?: error("Current context doesn't contain Job in it: $this")
+
 /**
  * @suppress This method has bad semantics when cause is not a [CancellationException]. Use [CoroutineContext.cancelChildren].
  */
diff --git a/kotlinx-coroutines-core/common/test/EnsureActiveTest.kt b/kotlinx-coroutines-core/common/test/JobExtensionsTest.kt
similarity index 81%
rename from kotlinx-coroutines-core/common/test/EnsureActiveTest.kt
rename to kotlinx-coroutines-core/common/test/JobExtensionsTest.kt
index 89e749cae0..b335926b1f 100644
--- a/kotlinx-coroutines-core/common/test/EnsureActiveTest.kt
+++ b/kotlinx-coroutines-core/common/test/JobExtensionsTest.kt
@@ -4,9 +4,10 @@
 
 package kotlinx.coroutines
 
+import kotlin.coroutines.*
 import kotlin.test.*
 
-class EnsureActiveTest : TestBase() {
+class JobExtensionsTest : TestBase() {
 
     private val job = Job()
     private val scope = CoroutineScope(job + CoroutineExceptionHandler { _, _ ->  })
@@ -81,4 +82,14 @@ class EnsureActiveTest : TestBase() {
         assertTrue(exception is JobCancellationException)
         assertTrue(exception.cause is TestException)
     }
+
+    @Test
+    fun testJobExtension() = runTest {
+        assertSame(coroutineContext[Job]!!, coroutineContext.job)
+        assertSame(NonCancellable, NonCancellable.job)
+        assertSame(job, job.job)
+        assertFailsWith<IllegalStateException> { EmptyCoroutineContext.job }
+        assertFailsWith<IllegalStateException> { Dispatchers.Default.job }
+        assertFailsWith<IllegalStateException> { (Dispatchers.Default + CoroutineName("")).job }
+    }
 }
diff --git a/kotlinx-coroutines-core/jvm/src/Interruptible.kt b/kotlinx-coroutines-core/jvm/src/Interruptible.kt
index f50e0936d5..070aa62497 100644
--- a/kotlinx-coroutines-core/jvm/src/Interruptible.kt
+++ b/kotlinx-coroutines-core/jvm/src/Interruptible.kt
@@ -40,8 +40,7 @@ public suspend fun <T> runInterruptible(
 
 private fun <T> runInterruptibleInExpectedContext(coroutineContext: CoroutineContext, block: () -> T): T {
     try {
-        val job = coroutineContext[Job]!! // withContext always creates a job
-        val threadState = ThreadState(job)
+        val threadState = ThreadState(coroutineContext.job)
         threadState.setup()
         try {
             return block()
diff --git a/kotlinx-coroutines-core/jvm/src/debug/internal/DebugProbesImpl.kt b/kotlinx-coroutines-core/jvm/src/debug/internal/DebugProbesImpl.kt
index 4b7c09b347..83bc02c6d6 100644
--- a/kotlinx-coroutines-core/jvm/src/debug/internal/DebugProbesImpl.kt
+++ b/kotlinx-coroutines-core/jvm/src/debug/internal/DebugProbesImpl.kt
@@ -111,7 +111,7 @@ internal object DebugProbesImpl {
         check(isInstalled) { "Debug probes are not installed" }
         val jobToStack = capturedCoroutines
             .filter { it.delegate.context[Job] != null }
-            .associateBy({ it.delegate.context[Job]!! }, { it.info })
+            .associateBy({ it.delegate.context.job }, { it.info })
         return buildString {
             job.build(jobToStack, this, "")
         }