grpc-swift-nio-transport/IntegrationTests/grpc-performance-tests/Sources/BenchmarkService.swift

/*
 * Copyright 2024, 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 GRPCCore
import Synchronization

import struct Foundation.Data

final class BenchmarkService: Grpc_Testing_BenchmarkService.ServiceProtocol {
  /// Used to check if the server can be streaming responses.
  private let working = Atomic(true)

  /// One request followed by one response.
  /// The server returns a client payload with the size requested by the client.
  func unaryCall(
    request: ServerRequest<Grpc_Testing_SimpleRequest>,
    context: ServerContext
  ) async throws -> ServerResponse<Grpc_Testing_SimpleResponse> {
    // Throw an error if the status is not `ok`. Otherwise, an `ok` status is automatically sent
    // if the request is successful.
    if request.message.responseStatus.isInitialized {
      try self.checkOkStatus(request.message.responseStatus)
    }

    return ServerResponse(
      message: .with {
        $0.payload = Grpc_Testing_Payload.with {
          $0.body = Data(count: Int(request.message.responseSize))
        }
      }
    )
  }

  /// Repeated sequence of one request followed by one response.
  /// The server returns a payload with the size requested by the client for each received message.
  func streamingCall(
    request: StreamingServerRequest<Grpc_Testing_SimpleRequest>,
    context: ServerContext
  ) async throws -> StreamingServerResponse<Grpc_Testing_SimpleResponse> {
    return StreamingServerResponse { writer in
      for try await message in request.messages {
        if message.responseStatus.isInitialized {
          try self.checkOkStatus(message.responseStatus)
        }

        let responseMessage = Grpc_Testing_SimpleResponse.with {
          $0.payload = Grpc_Testing_Payload.with {
            $0.body = Data(count: Int(message.responseSize))
          }
        }

        try await writer.write(responseMessage)
      }

      return [:]
    }
  }

  /// Single-sided unbounded streaming from client to server.
  /// The server returns a payload with the size requested by the client once the client does WritesDone.
  func streamingFromClient(
    request: StreamingServerRequest<Grpc_Testing_SimpleRequest>,
    context: ServerContext
  ) async throws -> ServerResponse<Grpc_Testing_SimpleResponse> {
    var responseSize = 0
    for try await message in request.messages {
      if message.responseStatus.isInitialized {
        try self.checkOkStatus(message.responseStatus)
      }
      responseSize = Int(message.responseSize)
    }

    return ServerResponse(
      message: .with {
        $0.payload = .with {
          $0.body = Data(count: responseSize)
        }
      }
    )
  }

  /// Single-sided unbounded streaming from server to client.
  /// The server repeatedly returns a payload with the size requested by the client.
  func streamingFromServer(
    request: ServerRequest<Grpc_Testing_SimpleRequest>,
    context: ServerContext
  ) async throws -> StreamingServerResponse<Grpc_Testing_SimpleResponse> {
    if request.message.responseStatus.isInitialized {
      try self.checkOkStatus(request.message.responseStatus)
    }

    let response = Grpc_Testing_SimpleResponse.with {
      $0.payload = .with {
        $0.body = Data(count: Int(request.message.responseSize))
      }
    }

    return StreamingServerResponse { writer in
      while self.working.load(ordering: .relaxed) {
        try await writer.write(response)
      }
      return [:]
    }
  }

  /// Two-sided unbounded streaming between server to client.
  /// Both sides send the content of their own choice to the other.
  func streamingBothWays(
    request: StreamingServerRequest<Grpc_Testing_SimpleRequest>,
    context: ServerContext
  ) async throws -> StreamingServerResponse<Grpc_Testing_SimpleResponse> {
    // The 100 size is used by the other implementations as well.
    // We are using the same canned response size for all responses
    // as it is allowed by the spec.
    let response = Grpc_Testing_SimpleResponse.with {
      $0.payload = .with {
        $0.body = Data(count: 100)
      }
    }

    final class InboundStreamingSignal: Sendable {
      private let _isStreaming: Atomic<Bool>

      init() {
        self._isStreaming = Atomic(true)
      }

      var isStreaming: Bool {
        self._isStreaming.load(ordering: .relaxed)
      }

      func stop() {
        self._isStreaming.store(false, ordering: .relaxed)
      }
    }

    // Marks if the inbound streaming is ongoing or finished.
    let inbound = InboundStreamingSignal()

    return StreamingServerResponse { writer in
      try await withThrowingTaskGroup(of: Void.self) { group in
        group.addTask {
          for try await message in request.messages {
            if message.responseStatus.isInitialized {
              try self.checkOkStatus(message.responseStatus)
            }
          }
          inbound.stop()
        }

        group.addTask {
          while inbound.isStreaming && self.working.load(ordering: .acquiring) {
            try await writer.write(response)
          }
        }

        try await group.next()
        group.cancelAll()
        return [:]
      }
    }
  }
}

extension BenchmarkService {
  private func checkOkStatus(_ responseStatus: Grpc_Testing_EchoStatus) throws {
    guard let code = Status.Code(rawValue: Int(responseStatus.code)) else {
      throw RPCError(code: .invalidArgument, message: "The response status code is invalid.")
    }
    if let code = RPCError.Code(code) {
      throw RPCError(code: code, message: responseStatus.message)
    }
  }
}