IO: fix continuous writing session interfered with joining

Author: Sergey Mashkov

core/kotlinx-coroutines-io/src/main/kotlin/kotlinx/coroutines/experimental/io/ByteBufferChannel.kt

@@ -1757,47 +1757,74 @@ internal class ByteBufferChannel(
     }
 
     override suspend fun writeSuspendSession(visitor: suspend WriterSuspendSession.() -> Unit) {
-        writing { byteBuffer, ringBufferCapacity ->
-            var locked = 0
-
-            val session = object : WriterSuspendSession {
-                override fun request(min: Int): ByteBuffer? {
-                    locked += ringBufferCapacity.tryWriteAtLeast(0)
-                    if (locked < min) return null
-                    byteBuffer.prepareBuffer(writeByteOrder, writePosition, locked)
-                    if (byteBuffer.remaining() < min) return null
-                    if (joining != null) return null
+        var locked = 0
 
-                    return byteBuffer
-                }
+        var current = joining?.let { resolveDelegation(this, it) } ?: this
+        var byteBuffer = current.setupStateForWrite() ?: return writeSuspendSession(visitor)
+        var ringBufferCapacity = current.state.capacity
 
-                override fun written(n: Int) {
-                    require(n >= 0)
-                    if (n > locked) throw IllegalStateException()
-                    locked -= n
-                    byteBuffer.bytesWritten(ringBufferCapacity, n)
-                }
+        val session = object : WriterSuspendSession {
+            override fun request(min: Int): ByteBuffer? {
+                locked += ringBufferCapacity.tryWriteAtLeast(0)
+                if (locked < min) return null
+                byteBuffer.prepareBuffer(writeByteOrder, writePosition, locked)
+                if (byteBuffer.remaining() < min) return null
+                if (current.joining != null) return null
 
-                override suspend fun tryAwait(n: Int) {
-                    if (locked >= n) return
-                    if (locked > 0) {
-                        ringBufferCapacity.completeRead(locked)
-                        locked = 0
-                    }
+                return byteBuffer
+            }
+
+            override fun written(n: Int) {
+                require(n >= 0)
+                if (n > locked) throw IllegalStateException()
+                locked -= n
+                byteBuffer.bytesWritten(ringBufferCapacity, n)
+            }
 
-                    return tryWriteSuspend(n)
+            override suspend fun tryAwait(n: Int) {
+                val joining = current.joining
+                if (joining != null) {
+                    return tryAwaitJoinSwitch(n, joining)
                 }
+
+                if (locked >= n) return
+                if (locked > 0) {
+                    ringBufferCapacity.completeRead(locked)
+                    locked = 0
+                }
+
+                return tryWriteSuspend(n)
             }
 
-            try {
-                visitor(session)
-            } finally {
+            private suspend fun tryAwaitJoinSwitch(n: Int, joining: JoiningState) {
                 if (locked > 0) {
                     ringBufferCapacity.completeRead(locked)
                     locked = 0
                 }
+                flush()
+                restoreStateAfterWrite()
+                tryTerminate()
+
+                do {
+                    current.tryWriteSuspend(n)
+                    current = resolveDelegation(current, joining) ?: continue
+                    byteBuffer = current.setupStateForWrite() ?: continue
+                    ringBufferCapacity = current.state.capacity
+                } while (false)
             }
         }
+
+        try {
+            visitor(session)
+        } finally {
+            if (locked > 0) {
+                ringBufferCapacity.completeRead(locked)
+                locked = 0
+            }
+
+            current.restoreStateAfterWrite()
+            current.tryTerminate()
+        }
     }
 
     override fun consumed(n: Int) {

core/kotlinx-coroutines-io/src/test/kotlin/kotlinx/coroutines/experimental/io/ByteBufferChannelTest.kt

@@ -1546,6 +1546,78 @@ class ByteBufferChannelTest : TestBase() {
         sub.join()
     }
 
+    @Test
+    fun testWriteSuspendSessionSmokeTest() = runTest {
+        ch.writeSuspendSession {
+            val buffer = request(1)
+            assertNotNull(buffer)
+        }
+
+        ch.writeSuspendSession {
+            val buffer = request(1)!!
+            buffer.put(0x11)
+            written(1)
+        }
+
+        assertEquals(0, ch.availableForRead)
+        ch.flush()
+        assertEquals(1, ch.availableForRead)
+        assertEquals(0x11, ch.readByte())
+    }
+
+    @Test
+    fun testWriteSuspendSessionJoined() = runTest {
+        val next = ByteChannel()
+        launch(Unconfined) {
+            ch.joinTo(next, true)
+        }
+
+        yield()
+
+        ch.writeSuspendSession {
+            val buffer = request(1)
+            assertNotNull(buffer)
+            buffer!!.put(0x11)
+            written(1)
+        }
+
+        assertEquals(0, next.availableForRead)
+        ch.flush()
+        assertEquals(1, next.availableForRead)
+        assertEquals(0x11, next.readByte())
+    }
+
+    @Test
+    fun testWriteSuspendSessionJoinDuringWrite() = runTest {
+        val next = ByteChannel()
+
+        ch.writeSuspendSession {
+            var buffer = request(1)
+            assertNotNull(buffer)
+            buffer!!.put(0x11)
+            written(1)
+
+            launch(Unconfined) {
+                ch.joinTo(next, true)
+            }
+
+            yield()
+
+            assertNull(request(1))
+            tryAwait(1)
+            buffer = request(1)
+            assertNotNull(buffer)
+            buffer!!.put(0x22)
+            written(1)
+        }
+
+        ch.flush()
+
+        assertEquals(2, next.availableForRead)
+        assertEquals(0x11, next.readByte())
+        assertEquals(0x22, next.readByte())
+    }
+
     private inline fun buildPacket(block: ByteWritePacket.() -> Unit): ByteReadPacket {
         val builder = BytePacketBuilder(0, pktPool)
         try {