Let's see how it works :)

Author: ndkoval

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

@@ -6,6 +6,7 @@
 
 package kotlinx.coroutines
 
+import kotlinx.coroutines.internal.*
 import kotlin.coroutines.*
 import kotlin.test.*
 
@@ -159,4 +160,31 @@ class CancellableContinuationHandlersTest : TestBase() {
         }
         finish(3)
     }
+
+    @Test
+    fun testSegmentAsHandler() = runTest {
+        class MySegment : Segment<MySegment>(0, null, 0) {
+            override val maxSlots: Int get() = 0
+
+            var invokeOnCancellationCalled = false
+            override fun invokeOnCancellation(index: Int, cause: Throwable?) {
+                invokeOnCancellationCalled = true
+            }
+        }
+        val s = MySegment()
+        expect(1)
+        try {
+            suspendCancellableCoroutine<Unit> { c ->
+                expect(2)
+                c as CancellableContinuationImpl<*>
+                c.invokeOnCancellation(s, 0)
+                c.cancel()
+            }
+        } catch (e: CancellationException) {
+            expect(3)
+        }
+        expect(4)
+        check(s.invokeOnCancellationCalled)
+        finish(5)
+    }
 }