Add an internal pausable writer

Sources/GRPC/AsyncAwaitSupport/AsyncWriter.swift

@@ -0,0 +1,312 @@
+/*
+ * Copyright 2021, gRPC Authors All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#if compiler(>=5.5)
+import NIOCore
+
+/// An asynchronous writer which forwards messages to a delegate.
+///
+/// Forwarding of messages to the delegate may be paused and resumed by controlling the writability
+/// of the writer. This may be controlled by calls to ``toggleWritability()``. When the writer is
+/// paused (by becoming unwritable) calls to ``write(_:)`` may suspend. When the writer is resumed
+/// (by becoming writable) any calls which are suspended may be resumed.
+///
+/// The writer must also be "finished" with a final value: as for writing, calls to ``finish(_:)``
+/// may suspend if the writer has been paused.
+@available(macOS 12, iOS 15, tvOS 15, watchOS 8, *)
+@usableFromInline
+internal final actor AsyncWriter<Delegate: AsyncWriterDelegate> {
+  @usableFromInline
+  internal typealias Element = Delegate.Element
+
+  @usableFromInline
+  internal typealias End = Delegate.End
+
+  /// A value pending a write.
+  @usableFromInline
+  internal struct _Pending<Value> {
+    @usableFromInline
+    var value: Value
+
+    @usableFromInline
+    var continuation: CheckedContinuation<Void, Error>
+
+    @inlinable
+    internal init(_ value: Value, continuation: CheckedContinuation<Void, Error>) {
+      self.value = value
+      self.continuation = continuation
+    }
+  }
+
+  @usableFromInline
+  typealias PendingElement = _Pending<Element>
+
+  @usableFromInline
+  typealias PendingEnd = _Pending<End>
+
+  @usableFromInline
+  internal enum _CompletionState {
+    /// Finish hasn't been called yet. May move to `pending` or `completed`.
+    case incomplete
+    /// Finish has been called but the writer is paused. May move to `completed`.
+    case pending(PendingEnd)
+    /// The completion message has been sent to the delegate. This is a terminal state.
+    case completed
+
+    /// Move from `pending` to `completed` and return the `PendingCompletion`. Returns `nil` if
+    /// the state was not `pending`.
+    @inlinable
+    mutating func completeIfPending() -> PendingEnd? {
+      switch self {
+      case let .pending(pending):
+        self = .completed
+        return pending
+      case .incomplete, .completed:
+        return nil
+      }
+    }
+
+    @usableFromInline
+    var isPendingOrCompleted: Bool {
+      switch self {
+      case .incomplete:
+        return false
+      case .pending, .completed:
+        return true
+      }
+    }
+  }
+
+  /// The maximum number of pending elements. `pendingElements` must not grow beyond this limit.
+  @usableFromInline
+  internal let _maxPendingElements: Int
+
+  /// The maximum number of writes to the delegate made in `resume` before yielding to allow other
+  /// values to be queued.
+  @usableFromInline
+  internal let _maxWritesBeforeYield: Int
+
+  /// Elements and continuations which have been buffered but are awaiting consumption by the
+  /// delegate.
+  @usableFromInline
+  internal var _pendingElements: CircularBuffer<PendingElement>
+
+  /// The completion state of the writer.
+  @usableFromInline
+  internal var _completionState: _CompletionState
+
+  /// Whether the writer is paused.
+  @usableFromInline
+  internal var _isPaused: Bool = false
+
+  /// The delegate to process elements. By convention we call the delegate before resuming any
+  /// continuation.
+  @usableFromInline
+  internal let _delegate: Delegate
+
+  @inlinable
+  internal init(
+    maxPendingElements: Int = 16,
+    maxWritesBeforeYield: Int = 5,
+    delegate: Delegate
+  ) {
+    self._maxPendingElements = maxPendingElements
+    self._maxWritesBeforeYield = maxWritesBeforeYield
+    self._pendingElements = CircularBuffer(initialCapacity: maxPendingElements)
+    self._completionState = .incomplete
+    self._delegate = delegate
+  }
+
+  deinit {
+    switch self._completionState {
+    case .completed:
+      ()
+    case .incomplete, .pending:
+      assertionFailure("writer has not completed is pending completion")
+    }
+  }
+
+  /// As ``toggleWritability()`` but executed asynchronously.
+  @usableFromInline
+  internal nonisolated func toggleWritabilityAsynchronously() {
+    Task {
+      await self.toggleWritability()
+    }
+  }
+
+  /// Toggles whether the writer is writable or not. The writer is initially writable.
+  ///
+  /// If the writer becomes writable then it may resume writes to the delegate. If it becomes
+  /// unwritable then calls to `write` may suspend until the writability changes again.
+  ///
+  /// This API does not offer explicit control over the writability state so the caller must ensure
+  /// calls to this function correspond with changes in writability. The reason for this is that the
+  /// underlying type is an `actor` and updating its state is therefore asynchronous. However,
+  /// this functions is not called from an asynchronous context so it is not possible to `await`
+  /// state updates to complete. Instead, changing the state is via a `nonisolated` function on
+  /// the `actor` which spawns a new task. If this or a similar API allowed the writability to be
+  /// explicitly set then calls to that API are not guaranteed to be ordered which may lead to
+  /// deadlock.
+  @usableFromInline
+  internal func toggleWritability() async {
+    if self._isPaused {
+      self._isPaused = false
+      await self.resumeWriting()
+    } else {
+      self._isPaused = true
+    }
+  }
+
+  private func resumeWriting() async {
+    var writes = 0
+
+    while !self._isPaused {
+      if let pendingElement = self._pendingElements.popFirst() {
+        self._delegate.write(pendingElement.value)
+        pendingElement.continuation.resume()
+      } else if let pendingCompletion = self._completionState.completeIfPending() {
+        self._delegate.writeEnd(pendingCompletion.value)
+        pendingCompletion.continuation.resume()
+      } else {
+        break
+      }
+
+      // `writes` will never exceed `maxWritesBeforeYield` so unchecked arithmetic is okay here.
+      writes &+= 1
+      if writes == self._maxWritesBeforeYield {
+        writes = 0
+        // We yield every so often to let the delegate (i.e. 'NIO.Channel') catch up since it may
+        // decide it is no longer writable.
+        await Task.yield()
+      }
+    }
+  }
+
+  /// As ``cancel()`` but executed asynchronously.
+  @usableFromInline
+  internal nonisolated func cancelAsynchronously() {
+    Task {
+      await self.cancel()
+    }
+  }
+
+  /// Cancel all pending writes.
+  ///
+  /// Any pending writes will be dropped and their continuations will be resumed with
+  /// a `CancellationError`. Any writes after cancellation has completed will also fail.
+  @usableFromInline
+  internal func cancel() {
+    // If there's an end we should fail that last.
+    let pendingEnd: PendingEnd?
+
+    // Mark our state as completed before resuming any continuations (any future writes should fail
+    // immediately).
+    switch self._completionState {
+    case .incomplete:
+      pendingEnd = nil
+      self._completionState = .completed
+
+    case let .pending(pending):
+      pendingEnd = pending
+      self._completionState = .completed
+
+    case .completed:
+      pendingEnd = nil
+    }
+
+    let cancellationError = CancellationError()
+
+    while let pending = self._pendingElements.popFirst() {
+      pending.continuation.resume(throwing: cancellationError)
+    }
+
+    pendingEnd?.continuation.resume(throwing: cancellationError)
+  }
+
+  /// Write an `element`.
+  ///
+  /// The call may be suspend if the writer is paused.
+  ///
+  /// Throws: ``AsyncWriterError`` if the writer has already been finished or too many write tasks
+  ///   have been suspended.
+  @inlinable
+  internal func write(_ element: Element) async throws {
+    try await withCheckedThrowingContinuation { continuation in
+      self._write(element, continuation: continuation)
+    }
+  }
+
+  @inlinable
+  internal func _write(_ element: Element, continuation: CheckedContinuation<Void, Error>) {
+    // There are three outcomes of writing:
+    // - write the element directly (if the writer isn't paused and no writes are pending)
+    // - queue the element (the writer is paused or there are writes already pending)
+    // - error (the writer is complete or the queue is full).
+
+    if self._completionState.isPendingOrCompleted {
+      continuation.resume(throwing: AsyncWriterError.alreadyFinished)
+    } else if !self._isPaused, self._pendingElements.isEmpty {
+      self._delegate.write(element)
+      continuation.resume()
+    } else if self._pendingElements.count < self._maxPendingElements {
+      // The continuation will be resumed later.
+      self._pendingElements.append(PendingElement(element, continuation: continuation))
+    } else {
+      continuation.resume(throwing: AsyncWriterError.tooManyPendingWrites)
+    }
+  }
+
+  /// Write the final element
+  @inlinable
+  internal func finish(_ end: End) async throws {
+    try await withCheckedThrowingContinuation { continuation in
+      self._finish(end, continuation: continuation)
+    }
+  }
+
+  @inlinable
+  internal func _finish(_ end: End, continuation: CheckedContinuation<Void, Error>) {
+    if self._completionState.isPendingOrCompleted {
+      continuation.resume(throwing: AsyncWriterError.alreadyFinished)
+    } else if !self._isPaused, self._pendingElements.isEmpty {
+      self._completionState = .completed
+      self._delegate.writeEnd(end)
+      continuation.resume()
+    } else {
+      // Either we're paused or there are pending writes which must be consumed first.
+      self._completionState = .pending(PendingEnd(end, continuation: continuation))
+    }
+  }
+}
+
+@usableFromInline
+internal enum AsyncWriterError: Error, Hashable {
+  case tooManyPendingWrites
+  case alreadyFinished
+}
+
+@usableFromInline
+internal protocol AsyncWriterDelegate: AnyObject {
+  associatedtype Element
+  associatedtype End
+
+  @inlinable
+  func write(_ element: Element)
+
+  @inlinable
+  func writeEnd(_ end: End)
+}
+
+#endif // compiler(>=5.5)