IO: make buffer release guaranteed for cases

Make byte channel release if it has been closed with
error or if write preparation loop has allocated a buffer
but failed to start writing (for example, joining started
so we have to cancel wriging preparation).

It is important to release allocated/borrowed buffer during
write operation preparation and if CAS loop was not successful
and no write state was set due to some reason (for example, joining).

Author: Sergey Mashkov

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

@@ -131,17 +131,28 @@ internal class ByteBufferChannel(
         var _allocated: ReadWriteBufferState.Initial? = null
         val (old, newState) = updateState { state ->
             when {
-                joining != null -> return null
-                closed != null -> throw closed!!.sendException
+                joining != null -> {
+                    _allocated?.let { releaseBuffer(it) }
+                    return null
+                }
+                closed != null -> {
+                    _allocated?.let { releaseBuffer(it) }
+                    throw closed!!.sendException
+                }
                 state === ReadWriteBufferState.IdleEmpty -> {
                     val allocated = _allocated ?: newBuffer().also { _allocated = it }
                     allocated.startWriting()
                 }
-                state === ReadWriteBufferState.Terminated -> throw closed!!.sendException
+                state === ReadWriteBufferState.Terminated -> {
+                    _allocated?.let { releaseBuffer(it) }
+                    if (joining != null) return null
+                    throw closed!!.sendException
+                }
                 else -> state.startWriting()
             }
         }
 
+//        joining?.let { restoreStateAfterWrite(); return null }
         if (closed != null) {
             restoreStateAfterWrite()
             tryTerminate()
@@ -247,34 +258,44 @@ internal class ByteBufferChannel(
     }
 
     private fun tryCompleteJoining(joined: JoiningState): Boolean {
-        updateState { state ->
-            when {
-                state === ReadWriteBufferState.Terminated -> state
-                state === ReadWriteBufferState.IdleEmpty -> ReadWriteBufferState.Terminated
-                // we don't handle IdleNonEmpty as it should be switched to IdleEmpty in restoreStateAfterRead
-                else -> return false
-            }
-        }
-
+        if (!tryReleaseBuffer()) return false
         ensureClosedJoined(joined)
 
-        ReadOp.getAndSet(this, null)?.resumeWithException(IllegalStateException("Joining is in progress"))
-        WriteOp.getAndSet(this, null)?.resume(Unit)
+        resumeReadOp(IllegalStateException("Joining is in progress"))
+        resumeWriteOp() // here we don't resume it with exception because it should resume and delegate writing
 
         return true
     }
 
     private fun tryTerminate(): Boolean {
-        val closed = closed ?: return false
+        if (closed == null) return false
+
+        if (!tryReleaseBuffer()) return false
+
+        joining?.let { ensureClosedJoined(it) }
+
+        resumeReadOp()
+        resumeWriteOp()
+
+        return true
+    }
 
+    private fun tryReleaseBuffer(): Boolean {
         var toRelease: ReadWriteBufferState.Initial? = null
 
         updateState { state ->
+            toRelease?.let { buffer ->
+                toRelease = null
+                buffer.capacity.resetForWrite()
+                resumeWriteOp()
+            }
+            val closed = closed
+
             when {
                 state === ReadWriteBufferState.Terminated -> return true
                 state === ReadWriteBufferState.IdleEmpty -> ReadWriteBufferState.Terminated
-                closed.cause != null && state is ReadWriteBufferState.IdleNonEmpty -> {
-                    // here we don't need to tryLockForRelease as we already have closed state
+                closed != null && state is ReadWriteBufferState.IdleNonEmpty && (state.capacity.tryLockForRelease() || closed.cause != null) -> {
+                    if (closed.cause != null) state.capacity.forceLockForRelease()
                     toRelease = state.initial
                     ReadWriteBufferState.Terminated
                 }
@@ -288,13 +309,6 @@ internal class ByteBufferChannel(
             }
         }
 
-        joining?.let { ensureClosedJoined(it) }
-
-        WriteOp.getAndSet(this, null)?.resumeWithException(closed.sendException)
-        ReadOp.getAndSet(this, null)?.apply {
-            if (closed.cause != null) resumeWithException(closed.cause) else resume(false)
-        }
-
         return true
     }
 
@@ -1795,7 +1809,17 @@ internal class ByteBufferChannel(
     suspend override fun <A : Appendable> readUTF8LineTo(out: A, limit: Int) = readUTF8LineToAscii(out, limit)
 
     private fun resumeReadOp() {
-        ReadOp.getAndSet(this, null)?.resume(true)
+        ReadOp.getAndSet(this, null)?.apply {
+            val closedCause = closed?.cause
+            when {
+                closedCause != null -> resumeWithException(closedCause)
+                else -> resume(true)
+            }
+        }
+    }
+
+    private fun resumeReadOp(result: Throwable) {
+        ReadOp.getAndSet(this, null)?.resumeWithException(result)
     }
 
     private fun resumeWriteOp() {
@@ -1805,6 +1829,10 @@ internal class ByteBufferChannel(
         }
     }
 
+    private fun resumeWriteOp(cause: Throwable) {
+        WriteOp.getAndSet(this, null)?.resumeWithException(cause)
+    }
+
     private fun resumeClosed(cause: Throwable?) {
         ReadOp.getAndSet(this, null)?.let { c ->
             if (cause != null)

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

@@ -105,6 +105,14 @@ internal class RingBufferCapacity(private val totalCapacity: Int) {
         }
     }
 
+    /**
+     * Make all writers to fail to write any more bytes
+     * Use only during failure termination
+     */
+    fun forceLockForRelease() {
+        AvailableForWrite.getAndSet(this, 0)
+    }
+
     fun isEmpty(): Boolean = availableForWrite == totalCapacity
     fun isFull(): Boolean = availableForWrite == 0
 

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

@@ -897,20 +897,56 @@ class ByteBufferChannelTest {
     }
 
     @Test
-    @Ignore
+    fun writeThenReadStress() = runBlocking<Unit> {
+        for (i in 1..500_000) {
+            val a = ByteBufferChannel(false, pool)
+
+            val w = launch {
+                a.writeLong(1)
+                a.close()
+            }
+            val r = launch {
+                a.readLong()
+            }
+
+            w.join()
+            r.join()
+        }
+
+        ch.close()
+    }
+
+    @Test
+    fun joinToEmptyStress() = runBlocking<Unit> {
+        for (i in 1..500_000) {
+            val a = ByteBufferChannel(false, pool)
+
+            launch(coroutineContext) {
+                a.joinTo(ch, true)
+            }
+
+            yield()
+
+            a.close()
+        }
+    }
+
+    @Test
     fun testJoinToStress() = runBlocking<Unit> {
-        for (i in 1..10) {
-            println("Step $i")
+        for (i in 1..100000) {
             val child = ByteBufferChannel(false, pool)
             val writer = launch {
-                child.writeLong(999)
+                child.writeLong(999 + i.toLong())
                 child.close()
             }
 
             child.joinTo(ch, false)
-            assertEquals(999, ch.readLong())
-            assertTrue { writer.isCompleted }
+            assertEquals(999 + i.toLong(), ch.readLong())
+            writer.join()
         }
+
+        assertEquals(0, ch.availableForRead)
+        ch.close()
     }
 
     @Test