Introduce CoroutineContext.job extension (#2312)

Fixes #2159

Author: qwwdfsad

Diff for: 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 for: 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 for: kotlinx-coroutines-core/common/test/EnsureActiveTest.kt renamed to 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 for: 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 for: 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, "")
         }