Add new ChannelOption to get the amount of buffered outbound data in the Channel

Sources/NIOCore/BSDSocketAPI.swift

@@ -405,6 +405,9 @@ extension NIOBSDSocket.Option {
 
  /// Specifies the total per-socket buffer space reserved for receives.
  public static let so_rcvbuf = Self(rawValue: SO_RCVBUF)
+
+ /// Specifies the total per-socket buffer space reserved for sends.
+ public static let so_sndbuf = Self(rawValue: SO_SNDBUF)
 
  /// Specifies the receive timeout.
  public static let so_rcvtimeo = Self(rawValue: SO_RCVTIMEO)

Sources/NIOCore/ChannelOption.swift

@@ -287,6 +287,13 @@ extension ChannelOptions {
  public typealias Value = Bool
  public init() {}
  }
+
+ /// `BufferedWritableBytesOption` allows users to know the number of writable bytes currently buffered in the `Channel`.
+ public struct BufferedWritableBytesOption: ChannelOption, Sendable {
+ public typealias Value = Int
+
+ public init() {}
+ }
  }
 }
 
@@ -358,6 +365,9 @@ public struct ChannelOptions: Sendable {
 
  /// - seealso: `ReceivePacketInfo`
  public static let receivePacketInfo = Types.ReceivePacketInfo()
+
+ /// - seealso: `BufferedWritableBytesOption`
+ public static let bufferedWritableBytes = Types.BufferedWritableBytesOption()
 }
 
 /// - seealso: `SocketOption`.
@@ -451,6 +461,11 @@ extension ChannelOption where Self == ChannelOptions.Types.ReceivePacketInfo {
  public static var receivePacketInfo: Self { .init() }
 }
 
+/// - seealso: `BufferedWritableBytesOption`
+extension ChannelOption where Self == ChannelOptions.Types.BufferedWritableBytesOption {
+ public static var bufferedWritableBytes: Self { .init() }
+}
+
 extension ChannelOptions {
  /// A type-safe storage facility for `ChannelOption`s. You will only ever need this if you implement your own
  /// `Channel` that needs to store `ChannelOption`s.

Sources/NIOEmbedded/AsyncTestingChannel.swift

@@ -581,6 +581,14 @@ public final class NIOAsyncTestingChannel: Channel {
  if option is ChannelOptions.Types.AllowRemoteHalfClosureOption {
  return self.allowRemoteHalfClosure as! Option.Value
  }
+ if option is ChannelOptions.Types.BufferedWritableBytesOption {
+ let result = self.channelcore.pendingOutboundBuffer.reduce(0) { partialResult, dataAndPromise in
+ let buffer = self.channelcore.unwrapData(dataAndPromise.0, as: ByteBuffer.self)
+ return partialResult + buffer.readableBytes
+ }
+
+ return result as! Option.Value
+ }
  fatalError("option \(option) not supported")
  }
 

Sources/NIOEmbedded/Embedded.swift

@@ -863,6 +863,14 @@ public final class EmbeddedChannel: Channel {
  if option is ChannelOptions.Types.AllowRemoteHalfClosureOption {
  return self.allowRemoteHalfClosure as! Option.Value
  }
+ if option is ChannelOptions.Types.BufferedWritableBytesOption {
+ let result = self.channelcore.pendingOutboundBuffer.reduce(0) { partialResult, dataAndPromise in
+ let buffer = self.channelcore.unwrapData(dataAndPromise.0, as: ByteBuffer.self)
+ return partialResult + buffer.readableBytes
+ }
+
+ return result as! Option.Value
+ }
  fatalError("option \(option) not supported")
  }
 

Sources/NIOPosix/BaseStreamSocketChannel.swift

@@ -76,6 +76,8 @@ class BaseStreamSocketChannel<Socket: SocketProtocol>: BaseSocketChannel<Socket>
  return self.pendingWrites.writeSpinCount as! Option.Value
  case _ as ChannelOptions.Types.WriteBufferWaterMarkOption:
  return self.pendingWrites.waterMark as! Option.Value
+ case _ as ChannelOptions.Types.BufferedWritableBytesOption:
+ return Int(self.pendingWrites.bufferedBytes) as! Option.Value
  default:
  return try super.getOption0(option)
  }

Sources/NIOPosix/PendingDatagramWritesManager.swift

@@ -450,6 +450,10 @@ final class PendingDatagramWritesManager: PendingWritesManager {
  var isEmpty: Bool {
  self.state.isEmpty
  }
+
+ var bufferedBytes: Int64 {
+ return self.state.bytes
+ }
 
  private func add(_ pendingWrite: PendingDatagramWrite) -> Bool {
  assert(self.isOpen)

Sources/NIOPosix/PendingWritesManager.swift

@@ -323,6 +323,10 @@ final class PendingStreamWritesManager: PendingWritesManager {
  var isEmpty: Bool {
  self.state.isEmpty
  }
+
+ var bufferedBytes: Int64 {
+ return self.state.bytes
+ }
 
  /// Add a pending write alongside its promise.
  ///

Sources/NIOPosix/SocketChannel.swift

@@ -676,6 +676,8 @@ final class DatagramChannel: BaseSocketChannel<Socket> {
  throw ChannelError._operationUnsupported
  }
  return try self.socket.getUDPReceiveOffload() as! Option.Value
+ case _ as ChannelOptions.Types.BufferedWritableBytesOption:
+ return Int(self.pendingWrites.bufferedBytes) as! Option.Value
  default:
  return try super.getOption0(option)
  }

Tests/NIOEmbeddedTests/AsyncTestingChannelTests.swift

@@ -551,6 +551,107 @@ class AsyncTestingChannelTests: XCTestCase {
  _ = try await channel.finish()
  await XCTAsyncAssertThrowsError(try await channel.finish())
  }
+
+ func testWriteOutboundEmptyBufferedByte() async throws {
+ let channel = NIOAsyncTestingChannel()
+ var buffered: ChannelOptions.Types.BufferedWritableBytesOption.Value = try await channel.getOption(.bufferedWritableBytes)
+ XCTAssertEqual(0, buffered)
+
+ let buf = channel.allocator.buffer(capacity: 10)
+
+ channel.write(buf, promise: nil)
+ buffered = try await channel.getOption(.bufferedWritableBytes)
+ XCTAssertEqual(0, buffered)
+
+ channel.flush()
+ buffered = try await channel.getOption(.bufferedWritableBytes)
+ XCTAssertEqual(0, buffered)
+
+ try await XCTAsyncAssertEqual(buf, try await channel.waitForOutboundWrite(as: ByteBuffer.self))
+ try await XCTAsyncAssertTrue(try await channel.finish().isClean) 
+ }
+
+ func testWriteOutboundBufferedBytesSingleWrite() async throws {
+ let channel = NIOAsyncTestingChannel()
+ var buf = channel.allocator.buffer(capacity: 10)
+ buf.writeString("hello")
+
+ channel.write(buf, promise: nil)
+ var buffered: ChannelOptions.Types.BufferedWritableBytesOption.Value = try await channel.getOption(.bufferedWritableBytes)
+ XCTAssertEqual(buf.readableBytes, buffered)
+ channel.flush()
+
+ buffered = try await channel.getOption(.bufferedWritableBytes).get()
+ XCTAssertEqual(0, buffered)
+ try await XCTAsyncAssertEqual(buf, try await channel.waitForOutboundWrite(as: ByteBuffer.self))
+ try await XCTAsyncAssertTrue(try await channel.finish().isClean)
+ }
+
+ func testWriteOuboundBufferedBytesMultipleWrites() async throws {
+ let channel = NIOAsyncTestingChannel()
+ var buf = channel.allocator.buffer(capacity: 10)
+ buf.writeString("hello")
+ let totalCount = 5
+ for _ in 0..<totalCount {
+ channel.write(buf, promise: nil)
+ }
+ var buffered: ChannelOptions.Types.BufferedWritableBytesOption.Value = try await channel.getOption(.bufferedWritableBytes)
+ XCTAssertEqual(buf.readableBytes * totalCount, buffered)
+
+ channel.flush()
+ buffered = try await channel.getOption(.bufferedWritableBytes)
+ XCTAssertEqual(0, buffered)
+
+ for _ in 0..<totalCount {
+ try await XCTAsyncAssertEqual(buf, try await channel.waitForOutboundWrite(as: ByteBuffer.self))
+ }
+
+ try await XCTAsyncAssertTrue(try await channel.finish().isClean)
+ }
+
+ func testWriteOuboundBufferedBytesWriteAndFlushInterleaved() async throws {
+ let channel = NIOAsyncTestingChannel()
+ var buf = channel.allocator.buffer(capacity: 10)
+ buf.writeString("hello")
+
+ channel.write(buf, promise: nil)
+ channel.write(buf, promise: nil)
+ channel.write(buf, promise: nil)
+ var buffered: ChannelOptions.Types.BufferedWritableBytesOption.Value = try await channel.getOption(.bufferedWritableBytes)
+ XCTAssertEqual(buf.readableBytes * 3, buffered)
+
+ channel.flush()
+ buffered = try await channel.getOption(.bufferedWritableBytes)
+ XCTAssertEqual(0, buffered)
+
+ channel.write(buf, promise: nil)
+ channel.write(buf, promise: nil)
+ buffered = try await channel.getOption(.bufferedWritableBytes)
+ XCTAssertEqual(buf.readableBytes * 2, buffered)
+ channel.flush()
+ buffered = try await channel.getOption(.bufferedWritableBytes)
+ XCTAssertEqual(0, buffered)
+
+
+ for _ in 0..<5 {
+ try await XCTAsyncAssertEqual(buf, try await channel.waitForOutboundWrite(as: ByteBuffer.self))
+ }
+
+ try await XCTAsyncAssertTrue(try await channel.finish().isClean)
+ }
+
+ func testWriteOutboundBufferedBytesWriteAndFlush() async throws {
+ let channel = NIOAsyncTestingChannel()
+ var buf = channel.allocator.buffer(capacity: 10)
+ buf.writeString("hello")
+
+ try await XCTAsyncAssertTrue(await channel.writeOutbound(buf).isFull)
+ let buffered: ChannelOptions.Types.BufferedWritableBytesOption.Value = try await channel.getOption(.bufferedWritableBytes)
+ XCTAssertEqual(0, buffered)
+
+ try await XCTAsyncAssertEqual(buf, try await channel.waitForOutboundWrite(as: ByteBuffer.self))
+ try await XCTAsyncAssertTrue(try await channel.finish().isClean)
+ }
 }
 
 @available(macOS 10.15, iOS 13.0, watchOS 6.0, tvOS 13.0, *)

Tests/NIOEmbeddedTests/EmbeddedChannelTest.swift

@@ -559,4 +559,104 @@ class EmbeddedChannelTest: XCTestCase {
 
  XCTAssertEqual(try channel._channelCore.remoteAddress0(), remoteAddress)
  }
+
+ func testWriteOutboundEmptyBufferedByte() throws {
+ let channel = EmbeddedChannel()
+ var buffered = try channel.getOption(ChannelOptions.bufferedWritableBytes).wait()
+ XCTAssertEqual(0, buffered)
+
+ let buf = channel.allocator.buffer(capacity: 10)
+
+ channel.write(buf, promise: nil)
+ buffered = try channel.getOption(ChannelOptions.bufferedWritableBytes).wait()
+ XCTAssertEqual(0, buffered)
+
+ channel.flush()
+ buffered = try channel.getOption(ChannelOptions.bufferedWritableBytes).wait()
+ XCTAssertEqual(0, buffered)
+ XCTAssertNoThrow(XCTAssertEqual(buf, try channel.readOutbound()))
+ XCTAssertNoThrow(XCTAssertTrue(try channel.finish().isClean))
+ }
+
+ func testWriteOutboundBufferedByteSingleWrite() throws {
+ let channel = EmbeddedChannel()
+ var buf = channel.allocator.buffer(capacity: 10)
+ buf.writeString("hello")
+
+ channel.write(buf, promise: nil)
+ var buffered = try channel.getOption(ChannelOptions.bufferedWritableBytes).wait()
+ XCTAssertEqual(buf.readableBytes, buffered)
+ channel.flush()
+
+ buffered = try channel.getOption(ChannelOptions.bufferedWritableBytes).wait()
+ XCTAssertEqual(0, buffered)
+ XCTAssertNoThrow(XCTAssertEqual(buf, try channel.readOutbound()))
+ XCTAssertNoThrow(XCTAssertTrue(try channel.finish().isClean))
+ }
+
+ func testWriteOuboundBufferedBytesMultipleWrites() throws {
+ let channel = EmbeddedChannel()
+ var buf = channel.allocator.buffer(capacity: 10)
+ buf.writeString("hello")
+ let totalCount = 5
+ for _ in 0..<totalCount {
+ channel.write(buf, promise: nil)
+ }
+
+ var buffered = try channel.getOption(ChannelOptions.bufferedWritableBytes).wait()
+ XCTAssertEqual(buf.readableBytes * totalCount, buffered)
+
+ channel.flush()
+ buffered = try channel.getOption(ChannelOptions.bufferedWritableBytes).wait()
+ XCTAssertEqual(0, buffered)
+ for _ in 0..<totalCount {
+ XCTAssertNoThrow(XCTAssertEqual(buf, try channel.readOutbound()))
+ }
+
+ XCTAssertNoThrow(XCTAssertTrue(try channel.finish().isClean))
+ }
+
+ func testWriteOuboundBufferedBytesWriteAndFlushInterleaved() throws {
+ let channel = EmbeddedChannel()
+ var buf = channel.allocator.buffer(capacity: 10)
+ buf.writeString("hello")
+
+ channel.write(buf, promise: nil)
+ channel.write(buf, promise: nil)
+ channel.write(buf, promise: nil)
+ var buffered = try channel.getOption(ChannelOptions.bufferedWritableBytes).wait()
+ XCTAssertEqual(buf.readableBytes * 3, buffered)
+
+ channel.flush()
+ buffered = try channel.getOption(ChannelOptions.bufferedWritableBytes).wait()
+ XCTAssertEqual(0, buffered)
+
+ channel.write(buf, promise: nil)
+ channel.write(buf, promise: nil)
+ buffered = try channel.getOption(ChannelOptions.bufferedWritableBytes).wait()
+ XCTAssertEqual(buf.readableBytes * 2, buffered)
+ channel.flush()
+ buffered = try channel.getOption(ChannelOptions.bufferedWritableBytes).wait()
+ XCTAssertEqual(0, buffered)
+
+
+ try (0..<5).forEach { _ in
+ XCTAssertNoThrow(XCTAssertEqual(buf, try channel.readOutbound()))
+ }
+
+ XCTAssertNoThrow(XCTAssertTrue(try channel.finish().isClean))
+ }
+
+ func testWriteOutboundBufferedBytesWriteAndFlush() throws {
+ let channel = EmbeddedChannel()
+ var buf = channel.allocator.buffer(capacity: 10)
+ buf.writeString("hello")
+
+ try XCTAssertTrue(channel.writeOutbound(buf).isFull)
+ let buffered = try channel.getOption(.bufferedWritableBytes).wait()
+ XCTAssertEqual(0, buffered)
+
+ XCTAssertEqual(buf, try channel.readOutbound(as: ByteBuffer.self))
+ XCTAssertTrue(try channel.finish().isClean)
+ }
 }