Address review feedback: bound buffered writes and scope backpressure to yamux

- Reinstate maxBufferedConnectionWrites (10 MiB default) and its public API on
  StreamMuxerProtocol.getYamux. The limit is now enforced against the actual
  bytes pending in the child channel outbound buffers, covering both
  send-window stalls and parent TCP backpressure stalls (the latter was never
  bounded before). On overflow the offending stream is reset, its pending
  writes are failed and the stream channel is closed, releasing the buffers.
- Move all writability/backpressure logic (parent writability gating, write
  buffer watermarks, user-defined writability flags) from the common classes
  into YamuxHandler. AbstractMuxHandler and MuxChannel keep only neutral
  extension points, so Mplex behaves exactly as before.
- Remove the redundant retainedPendingWrite ref-counting in MuxChannel: the
  ChannelOutboundBuffer owns pending messages and releases them on
  remove/fail/close on every path.
- Clamp send window arithmetic so cumulative window updates from a
  misbehaving remote cannot overflow the window counter.
- Fix the deferred-disconnect counter not being decremented when a pending
  write fails, which could leave the FIN frame unsent.

Author: StefanBratanov

libp2p/src/main/kotlin/io/libp2p/core/mux/StreamMuxerProtocol.kt

@@ -4,6 +4,7 @@ import io.libp2p.core.multistream.MultistreamProtocol
 import io.libp2p.core.multistream.ProtocolBinding
 import io.libp2p.mux.mplex.MplexStreamMuxer
 import io.libp2p.mux.yamux.DEFAULT_ACK_BACKLOG_LIMIT
+import io.libp2p.mux.yamux.DEFAULT_MAX_BUFFERED_CONNECTION_WRITES
 import io.libp2p.mux.yamux.YamuxStreamMuxer
 
 fun interface StreamMuxerProtocol {
@@ -22,17 +23,24 @@ fun interface StreamMuxerProtocol {
         }
 
         /**
+         * @param maxBufferedConnectionWrites the maximum amount of bytes which may be buffered (pending)
+         * across the write buffers of all streams of a single connection. When the limit is overflowed,
+         * the stream attempting to write is reset
          * @param ackBacklogLimit the maximum amount of opened streams per connection which have not been acknowledged
          */
         @JvmStatic
         @JvmOverloads
-        fun getYamux(ackBacklogLimit: Int = DEFAULT_ACK_BACKLOG_LIMIT): StreamMuxerProtocol {
+        fun getYamux(
+            maxBufferedConnectionWrites: Int = DEFAULT_MAX_BUFFERED_CONNECTION_WRITES,
+            ackBacklogLimit: Int = DEFAULT_ACK_BACKLOG_LIMIT
+        ): StreamMuxerProtocol {
             return StreamMuxerProtocol { multistreamProtocol, protocols ->
                 YamuxStreamMuxer(
                     multistreamProtocol.createMultistream(
                         protocols
                     ).toStreamHandler(),
                     multistreamProtocol,
+                    maxBufferedConnectionWrites,
                     ackBacklogLimit
                 )
             }

libp2p/src/main/kotlin/io/libp2p/etc/util/netty/mux/AbstractMuxHandler.kt

@@ -17,11 +17,6 @@ private val log = LoggerFactory.getLogger(AbstractMuxHandler::class.java)
 abstract class AbstractMuxHandler<TData>() :
     ChannelInboundHandlerAdapter() {
 
-    private companion object {
-        const val PARENT_CHANNEL_WRITABILITY_INDEX = 1
-        const val MUXER_WRITABILITY_INDEX = 2
-    }
-
     private val streamMap: MutableMap<MuxId, MuxChannel<TData>> = mutableMapOf()
     var ctx: ChannelHandlerContext? = null
     private val activeFuture = CompletableFuture<Void>()
@@ -89,11 +84,12 @@ abstract class AbstractMuxHandler<TData>() :
      */
     abstract fun releaseMessage(msg: TData)
 
-    protected open fun isChildWritable(child: MuxChannel<TData>): Boolean = true
-
-    fun canWriteChild(child: MuxChannel<TData>): Boolean =
-        getChannelHandlerContext().channel().isWritable && isChildWritable(child)
-
+    /**
+     * Writes the child channel data to the underlying connection.
+     * Returns the number of consumed bytes which may be less than the readable bytes of [data]
+     * (including 0) when the muxer can't write at the moment (e.g. flow control window exhausted).
+     * Unconsumed data remains buffered in the child channel and is written again on [retryChildWrite]
+     */
     abstract fun onChildWrite(child: MuxChannel<TData>, data: TData): Int
 
     fun flushChildWrites() {
@@ -108,32 +104,22 @@ abstract class AbstractMuxHandler<TData>() :
         streamMap.values.forEach { it.retryWrite() }
     }
 
-    protected fun refreshChildWritability(id: MuxId) {
-        streamMap[id]?.let { refreshChildWritability(it) }
+    protected fun setChildMuxWritability(id: MuxId, index: Int, writable: Boolean) {
+        streamMap[id]?.setMuxWritability(index, writable)
     }
 
-    protected open fun onChildRegistered(child: MuxChannel<TData>) {
-        refreshChildWritability(child)
+    protected fun setAllChildMuxWritability(index: Int, writable: Boolean) {
+        streamMap.values.forEach { it.setMuxWritability(index, writable) }
     }
 
-    private fun refreshChildWritability(child: MuxChannel<TData>) {
-        child.setMuxWritability(
-            PARENT_CHANNEL_WRITABILITY_INDEX,
-            getChannelHandlerContext().channel().isWritable
-        )
-        child.setMuxWritability(MUXER_WRITABILITY_INDEX, isChildWritable(child))
-    }
+    /**
+     * Total bytes (including a small per-message accounting overhead) buffered
+     * in the outbound buffers of all child channels
+     */
+    protected fun totalChildPendingWriteBytes(): Long =
+        streamMap.values.sumOf { it.pendingWriteBytes() }
 
-    override fun channelWritabilityChanged(ctx: ChannelHandlerContext) {
-        val parentWritable = ctx.channel().isWritable
-        streamMap.values.forEach {
-            it.setMuxWritability(PARENT_CHANNEL_WRITABILITY_INDEX, parentWritable)
-        }
-        if (parentWritable) {
-            retryAllChildWrites()
-        }
-        super.channelWritabilityChanged(ctx)
-    }
+    protected open fun onChildRegistered(child: MuxChannel<TData>) {}
 
     protected fun onRemoteOpen(id: MuxId) {
         val initializer = inboundInitializer

libp2p/src/main/kotlin/io/libp2p/etc/util/netty/mux/MuxChannel.kt

@@ -5,7 +5,6 @@ import io.libp2p.etc.util.netty.AbstractChildChannel
 import io.netty.buffer.ByteBuf
 import io.netty.channel.ChannelMetadata
 import io.netty.channel.ChannelOutboundBuffer
-import io.netty.util.ReferenceCountUtil
 import java.net.SocketAddress
 
 /**
@@ -21,8 +20,7 @@ class MuxChannel<TData>(
     var remoteDisconnected = false
     var localDisconnected = false
     private var localDisconnectSent = false
-    private var disconnectFlushedMessages = 0
-    private var retainedPendingWrite: Any? = null
+    private var pendingDisconnectWrites = 0
 
     override fun metadata(): ChannelMetadata = ChannelMetadata(true)
     override fun localAddress0() =
@@ -46,77 +44,77 @@ class MuxChannel<TData>(
         unsafe().outboundBuffer()?.setUserDefinedWritability(index, writable)
     }
 
+    fun pendingWriteBytes(): Long = unsafe().outboundBuffer()?.totalPendingWriteBytes() ?: 0L
+
     @Suppress("SwallowedException")
     override fun doWrite(buf: ChannelOutboundBuffer) {
         while (true) {
             val msg = buf.current() ?: break
-            if (localDisconnected && disconnectFlushedMessages == 0) {
+            if (localDisconnected && pendingDisconnectWrites == 0) {
                 // Must not throw from doWrite — exceptions escape uncaught to the Netty event loop.
                 // Wrap buf.remove() defensively: in some Netty versions promise listeners triggered
                 // by buf.remove() can propagate back through it.
                 try {
                     buf.remove(ConnectionClosedException("The stream was closed for writing locally: $id"))
                 } catch (e: Throwable) { }
-                releaseRetainedPendingWrite(msg)
                 continue
             }
-            if (!parent.canWriteChild(this)) {
-                retainPendingWrite(msg)
-                break
-            }
             try {
                 @Suppress("UNCHECKED_CAST")
-                val consumedBytes = parent.onChildWrite(this, msg as TData)
-                if (consumedBytes <= 0) {
-                    retainPendingWrite(msg)
-                    break
-                }
-                val currentMessageComplete = msg !is ByteBuf || consumedBytes >= msg.readableBytes()
-                if (!currentMessageComplete) {
-                    retainPendingWrite(msg)
-                }
-                if (msg is ByteBuf) {
+                msg as TData
+                val complete: Boolean
+                if (msg is ByteBuf && msg.isReadable) {
+                    val consumedBytes = parent.onChildWrite(this, msg)
+                    if (consumedBytes <= 0) {
+                        // the muxer can't write at the moment: the message stays buffered
+                        // (the outbound buffer owns it) until the parent retries this write
+                        break
+                    }
+                    complete = consumedBytes >= msg.readableBytes()
                     buf.removeBytes(consumedBytes.toLong())
                 } else {
+                    parent.onChildWrite(this, msg)
+                    complete = true
                     buf.remove()
                 }
                 parent.flushChildWrites()
-                if (currentMessageComplete) {
-                    releaseRetainedPendingWrite(msg)
-                }
-                if (localDisconnected && currentMessageComplete) {
-                    disconnectFlushedMessages--
-                    if (disconnectFlushedMessages == 0) {
-                        finishLocalDisconnect()
-                        continue
-                    }
+                if (!complete) {
+                    // partially written: wait for a retry before sending the rest
+                    break
                 }
+                onPendingDisconnectWriteCompleted()
             } catch (cause: Throwable) {
                 buf.remove(cause)
-                releaseRetainedPendingWrite(msg)
+                onPendingDisconnectWriteCompleted()
             }
         }
     }
 
     override fun doDisconnect() {
         localDisconnected = true
-        disconnectFlushedMessages = unsafe().outboundBuffer()?.size() ?: 0
-        if (disconnectFlushedMessages == 0) {
+        // delay the local disconnect until all writes flushed before the disconnect are sent
+        pendingDisconnectWrites = unsafe().outboundBuffer()?.size() ?: 0
+        if (pendingDisconnectWrites == 0) {
             finishLocalDisconnect()
         }
     }
 
+    private fun onPendingDisconnectWriteCompleted() {
+        if (localDisconnected && pendingDisconnectWrites > 0) {
+            pendingDisconnectWrites--
+            if (pendingDisconnectWrites == 0) {
+                finishLocalDisconnect()
+            }
+        }
+    }
+
     fun onRemoteDisconnected() {
         pipeline().fireUserEventTriggered(RemoteWriteClosed)
         remoteDisconnected = true
         closeIfBothDisconnected()
     }
 
     override fun doClose() {
-        retainedPendingWrite?.let {
-            retainedPendingWrite = null
-            ReferenceCountUtil.release(it)
-        }
         super.doClose()
         parent.onClosed(this)
     }
@@ -138,20 +136,6 @@ class MuxChannel<TData>(
         deactivate()
         closeIfBothDisconnected()
     }
-
-    private fun retainPendingWrite(msg: Any) {
-        if (retainedPendingWrite !== msg) {
-            retainedPendingWrite = msg
-            ReferenceCountUtil.retain(msg)
-        }
-    }
-
-    private fun releaseRetainedPendingWrite(msg: Any) {
-        if (retainedPendingWrite === msg) {
-            retainedPendingWrite = null
-            ReferenceCountUtil.release(msg)
-        }
-    }
 }
 
 data class MultiplexSocketAddress(val parentAddress: SocketAddress, val streamId: MuxId) : SocketAddress() {

libp2p/src/main/kotlin/io/libp2p/mux/yamux/YamuxHandler.kt

@@ -18,6 +18,7 @@ import java.util.concurrent.atomic.AtomicBoolean
 import java.util.concurrent.atomic.AtomicInteger
 import kotlin.properties.Delegates
 
+const val DEFAULT_MAX_BUFFERED_CONNECTION_WRITES = 10 * 1024 * 1024 // 10 MiB
 const val DEFAULT_MAX_BUFFERED_STREAM_WRITES = 64 * 1024
 const val DEFAULT_ACK_BACKLOG_LIMIT = 256
 
@@ -29,10 +30,16 @@ open class YamuxHandler(
     ready: CompletableFuture<StreamMuxer.Session>?,
     inboundStreamHandler: StreamHandler<*>,
     private val connectionInitiator: Boolean,
+    private val maxBufferedConnectionWrites: Int,
     private val ackBacklogLimit: Int,
     private val initialWindowSize: Int = INITIAL_WINDOW_SIZE
 ) : MuxHandler(ready, inboundStreamHandler) {
 
+    private companion object {
+        const val PARENT_WRITABILITY_INDEX = 1
+        const val SEND_WINDOW_WRITABILITY_INDEX = 2
+    }
+
     private val childWriteBufferWaterMark =
         WriteBufferWaterMark(DEFAULT_MAX_BUFFERED_STREAM_WRITES / 2, DEFAULT_MAX_BUFFERED_STREAM_WRITES)
 
@@ -45,9 +52,9 @@ open class YamuxHandler(
         val receiveWindowSize = AtomicInteger(initialWindowSize)
         var closedForWriting by Delegates.writeOnce(false)
 
-        fun dispose() {}
-
-        fun isWritable(): Boolean = sendWindowSize.get() > 0
+        fun refreshSendWindowWritability() {
+            setChildMuxWritability(id, SEND_WINDOW_WRITABILITY_INDEX, sendWindowSize.get() > 0)
+        }
 
         fun handleFrameRead(msg: YamuxFrame) {
             handleFlags(msg)
@@ -78,8 +85,13 @@ open class YamuxHandler(
 
         private fun handleWindowUpdate(msg: YamuxFrame) {
             val delta = msg.length.toInt()
-            sendWindowSize.addAndGet(delta)
-            refreshChildWritability(id)
+            // clamp the result so a misbehaving remote can't overflow the window counter
+            sendWindowSize.getAndUpdate { current ->
+                (current.toLong() + delta)
+                    .coerceIn(Int.MIN_VALUE.toLong(), Int.MAX_VALUE.toLong())
+                    .toInt()
+            }
+            refreshSendWindowWritability()
             retryChildWrite(id)
         }
 
@@ -129,7 +141,7 @@ open class YamuxHandler(
                 ctx.write(YamuxFrame(id, YamuxType.DATA, YamuxFlag.NONE, length.toLong(), slicedData))
                 written += length
             }
-            refreshChildWritability(id)
+            refreshSendWindowWritability()
             return written
         }
 
@@ -173,8 +185,6 @@ open class YamuxHandler(
         }
 
     override fun channelUnregistered(ctx: ChannelHandlerContext?) {
-        streamHandlers.values.forEach { it.dispose() }
-
         if (!goAwayPromise.isDone) {
             goAwayPromise.completeExceptionally(ConnectionClosedException("Connection was closed without Go Away message"))
         }
@@ -215,16 +225,46 @@ open class YamuxHandler(
         }
     }
 
-    override fun isChildWritable(child: MuxChannel<ByteBuf>): Boolean =
-        streamHandlers[child.id]?.isWritable() == true
-
     override fun onChildRegistered(child: MuxChannel<ByteBuf>) {
         child.config().setWriteBufferWaterMark(childWriteBufferWaterMark)
-        super.onChildRegistered(child)
+        child.setMuxWritability(PARENT_WRITABILITY_INDEX, getChannelHandlerContext().channel().isWritable)
+        streamHandlers[child.id]?.refreshSendWindowWritability()
+    }
+
+    override fun channelWritabilityChanged(ctx: ChannelHandlerContext) {
+        val parentWritable = ctx.channel().isWritable
+        setAllChildMuxWritability(PARENT_WRITABILITY_INDEX, parentWritable)
+        if (parentWritable) {
+            retryAllChildWrites()
+        }
+        super.channelWritabilityChanged(ctx)
     }
 
     override fun onChildWrite(child: MuxChannel<ByteBuf>, data: ByteBuf): Int {
-        return getStreamHandlerOrThrow(child.id).sendData(data)
+        val streamHandler = getStreamHandlerOrThrow(child.id)
+        if (!getChannelHandlerContext().channel().isWritable) {
+            // propagate the connection backpressure to the stream:
+            // don't queue more frames on the parent channel
+            verifyBufferedWritesLimit(child, 0)
+            return 0
+        }
+        val written = streamHandler.sendData(data)
+        if (written < data.readableBytes()) {
+            verifyBufferedWritesLimit(child, written)
+        }
+        return written
+    }
+
+    private fun verifyBufferedWritesLimit(child: MuxChannel<ByteBuf>, justWrittenBytes: Int) {
+        val bufferedBytes = totalChildPendingWriteBytes() - justWrittenBytes
+        if (bufferedBytes > maxBufferedConnectionWrites) {
+            onLocalClose(child)
+            // close the channel after the pending writes have been failed by the throw below
+            child.eventLoop().execute { child.closeImpl() }
+            throw WriteBufferOverflowMuxerException(
+                "Overflowed send buffer ($bufferedBytes/$maxBufferedConnectionWrites). Last stream attempting to write: ${child.id}"
+            )
+        }
     }
 
     override fun onLocalOpen(child: MuxChannel<ByteBuf>) {
@@ -261,7 +301,7 @@ open class YamuxHandler(
     }
 
     override fun onChildClosed(child: MuxChannel<ByteBuf>) {
-        streamHandlers.remove(child.id)?.dispose()
+        streamHandlers.remove(child.id)
     }
 
     private fun handlePing(msg: YamuxFrame) {