grpc-swift/Sources/GRPC/AsyncAwaitSupport/GRPCAsyncBidirectionalStreamingCall.swift

/*
 * 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.
 */
import NIOCore
import NIOHPACK

/// Async-await variant of ``BidirectionalStreamingCall``.
@available(macOS 10.15, iOS 13, tvOS 13, watchOS 6, *)
public struct GRPCAsyncBidirectionalStreamingCall<Request: Sendable, Response: Sendable>: Sendable {
  private let call: Call<Request, Response>
  private let responseParts: StreamingResponseParts<Response>
  private let responseSource: NIOThrowingAsyncSequenceProducer<
    Response,
    Error,
    NIOAsyncSequenceProducerBackPressureStrategies.HighLowWatermark,
    GRPCAsyncSequenceProducerDelegate
  >.Source
  private let requestSink: AsyncSink<(Request, Compression)>

  /// A request stream writer for sending messages to the server.
  public let requestStream: GRPCAsyncRequestStreamWriter<Request>

  /// The stream of responses from the server.
  public let responseStream: GRPCAsyncResponseStream<Response>

  /// The options used to make the RPC.
  public var options: CallOptions {
    return self.call.options
  }

  /// The path used to make the RPC.
  public var path: String {
    return self.call.path
  }

  /// Cancel this RPC if it hasn't already completed.
  public func cancel() {
    self.call.cancel(promise: nil)
  }

  // MARK: - Response Parts

  private func withRPCCancellation<R: Sendable>(_ fn: () async throws -> R) async rethrows -> R {
    return try await withTaskCancellationHandler(operation: fn) {
      self.cancel()
    }
  }

  /// The initial metadata returned from the server.
  ///
  /// - Important: The initial metadata will only be available when the first response has been
  /// received. However, it is not necessary for the response to have been consumed before reading
  /// this property.
  public var initialMetadata: HPACKHeaders {
    get async throws {
      try await self.withRPCCancellation {
        try await self.responseParts.initialMetadata.get()
      }
    }
  }

  /// The trailing metadata returned from the server.
  ///
  /// - Important: Awaiting this property will suspend until the responses have been consumed.
  public var trailingMetadata: HPACKHeaders {
    get async throws {
      try await self.withRPCCancellation {
        try await self.responseParts.trailingMetadata.get()
      }
    }
  }

  /// The final status of the the RPC.
  ///
  /// - Important: Awaiting this property will suspend until the responses have been consumed.
  public var status: GRPCStatus {
    get async {
      // force-try acceptable because any error is encapsulated in a successful GRPCStatus future.
      await self.withRPCCancellation {
        try! await self.responseParts.status.get()
      }
    }
  }

  private init(call: Call<Request, Response>) {
    self.call = call
    self.responseParts = StreamingResponseParts(on: call.eventLoop) { _ in }

    let sequenceProducer = NIOThrowingAsyncSequenceProducer<
      Response,
      Error,
      NIOAsyncSequenceProducerBackPressureStrategies.HighLowWatermark,
      GRPCAsyncSequenceProducerDelegate
    >.makeSequence(
      backPressureStrategy: .init(lowWatermark: 10, highWatermark: 50),
      delegate: GRPCAsyncSequenceProducerDelegate()
    )

    self.responseSource = sequenceProducer.source
    self.responseStream = .init(sequenceProducer.sequence)
    let (requestStream, requestSink) = call.makeRequestStreamWriter()
    self.requestStream = requestStream
    self.requestSink = AsyncSink(wrapping: requestSink)
  }

  /// We expose this as the only non-private initializer so that the caller
  /// knows that invocation is part of initialisation.
  internal static func makeAndInvoke(call: Call<Request, Response>) -> Self {
    let asyncCall = Self(call: call)

    asyncCall.call.invokeStreamingRequests(
      onStart: {
        asyncCall.requestSink.setWritability(to: true)
      },
      onError: { error in
        asyncCall.responseParts.handleError(error)
        asyncCall.responseSource.finish(error)
        asyncCall.requestSink.finish(error: error)
      },
      onResponsePart: AsyncCall.makeResponsePartHandler(
        responseParts: asyncCall.responseParts,
        responseSource: asyncCall.responseSource,
        requestStream: asyncCall.requestStream
      )
    )

    return asyncCall
  }
}

@available(macOS 10.15, iOS 13, tvOS 13, watchOS 6, *)
internal enum AsyncCall {
  internal static func makeResponsePartHandler<Response, Request>(
    responseParts: StreamingResponseParts<Response>,
    responseSource: NIOThrowingAsyncSequenceProducer<
      Response,
      Error,
      NIOAsyncSequenceProducerBackPressureStrategies.HighLowWatermark,
      GRPCAsyncSequenceProducerDelegate
    >.Source,
    requestStream: GRPCAsyncRequestStreamWriter<Request>?,
    requestType: Request.Type = Request.self
  ) -> (GRPCClientResponsePart<Response>) -> Void {
    return { responsePart in
      // Handle the metadata, trailers and status.
      responseParts.handle(responsePart)

      // Handle the response messages and status.
      switch responsePart {
      case .metadata:
        ()

      case let .message(response):
        // TODO: when we support backpressure we will need to stop ignoring the return value.
        _ = responseSource.yield(response)

      case let .end(status, _):
        if status.isOk {
          responseSource.finish()
        } else {
          responseSource.finish(status)
        }
        requestStream?.finish(status)
      }
    }
  }

  internal static func makeResponsePartHandler<Response, Request>(
    responseParts: UnaryResponseParts<Response>,
    requestStream: GRPCAsyncRequestStreamWriter<Request>?,
    requestType: Request.Type = Request.self,
    responseType: Response.Type = Response.self
  ) -> (GRPCClientResponsePart<Response>) -> Void {
    return { responsePart in
      // Handle (most of) all parts.
      responseParts.handle(responsePart)

      // Handle the status.
      switch responsePart {
      case .metadata, .message:
        ()
      case let .end(status, _):
        requestStream?.finish(status)
      }
    }
  }
}