/*
 *
 * Copyright 2015 gRPC authors.
 *
 * 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.
 *
 */

#include <forward_list>
#include <functional>
#include <list>
#include <memory>
#include <mutex>
#include <sstream>
#include <string>
#include <thread>
#include <utility>
#include <vector>

#include <grpc/grpc.h>
#include <grpc/support/cpu.h>
#include <grpc/support/log.h>
#include <grpcpp/alarm.h>
#include <grpcpp/channel.h>
#include <grpcpp/client_context.h>
#include <grpcpp/generic/generic_stub.h>

#include "absl/memory/memory.h"

#include "src/core/lib/surface/completion_queue.h"
#include "src/proto/grpc/testing/benchmark_service.grpc.pb.h"
#include "test/cpp/qps/client.h"
#include "test/cpp/qps/usage_timer.h"
#include "test/cpp/util/create_test_channel.h"

namespace grpc {
namespace testing {

class ClientRpcContext {
 public:
  ClientRpcContext() {}
  virtual ~ClientRpcContext() {}
  // next state, return false if done. Collect stats when appropriate
  virtual bool RunNextState(bool, HistogramEntry* entry) = 0;
  virtual void StartNewClone(CompletionQueue* cq) = 0;
  static void* tag(ClientRpcContext* c) { return static_cast<void*>(c); }
  static ClientRpcContext* detag(void* t) {
    return static_cast<ClientRpcContext*>(t);
  }

  virtual void Start(CompletionQueue* cq, const ClientConfig& config) = 0;
  virtual void TryCancel() = 0;
};

template <class RequestType, class ResponseType>
class ClientRpcContextUnaryImpl : public ClientRpcContext {
 public:
  ClientRpcContextUnaryImpl(
      BenchmarkService::Stub* stub, const RequestType& req,
      std::function<gpr_timespec()> next_issue,
      std::function<
          std::unique_ptr<grpc::ClientAsyncResponseReader<ResponseType>>(
              BenchmarkService::Stub*, grpc::ClientContext*, const RequestType&,
              CompletionQueue*)>
          prepare_req,
      std::function<void(grpc::Status, ResponseType*, HistogramEntry*)> on_done)
      : context_(),
        stub_(stub),
        cq_(nullptr),
        req_(req),
        response_(),
        next_state_(State::READY),
        callback_(on_done),
        next_issue_(std::move(next_issue)),
        prepare_req_(prepare_req) {}
  ~ClientRpcContextUnaryImpl() override {}
  void Start(CompletionQueue* cq, const ClientConfig& config) override {
    GPR_ASSERT(!config.use_coalesce_api());  // not supported.
    StartInternal(cq);
  }
  bool RunNextState(bool /*ok*/, HistogramEntry* entry) override {
    switch (next_state_) {
      case State::READY:
        start_ = UsageTimer::Now();
        response_reader_ = prepare_req_(stub_, &context_, req_, cq_);
        response_reader_->StartCall();
        next_state_ = State::RESP_DONE;
        response_reader_->Finish(&response_, &status_,
                                 ClientRpcContext::tag(this));
        return true;
      case State::RESP_DONE:
        if (status_.ok()) {
          entry->set_value((UsageTimer::Now() - start_) * 1e9);
        }
        callback_(status_, &response_, entry);
        next_state_ = State::INVALID;
        return false;
      default:
        GPR_ASSERT(false);
        return false;
    }
  }
  void StartNewClone(CompletionQueue* cq) override {
    auto* clone = new ClientRpcContextUnaryImpl(stub_, req_, next_issue_,
                                                prepare_req_, callback_);
    clone->StartInternal(cq);
  }
  void TryCancel() override { context_.TryCancel(); }

 private:
  grpc::ClientContext context_;
  BenchmarkService::Stub* stub_;
  CompletionQueue* cq_;
  std::unique_ptr<Alarm> alarm_;
  const RequestType& req_;
  ResponseType response_;
  enum State { INVALID, READY, RESP_DONE };
  State next_state_;
  std::function<void(grpc::Status, ResponseType*, HistogramEntry*)> callback_;
  std::function<gpr_timespec()> next_issue_;
  std::function<std::unique_ptr<grpc::ClientAsyncResponseReader<ResponseType>>(
      BenchmarkService::Stub*, grpc::ClientContext*, const RequestType&,
      CompletionQueue*)>
      prepare_req_;
  grpc::Status status_;
  double start_;
  std::unique_ptr<grpc::ClientAsyncResponseReader<ResponseType>>
      response_reader_;

  void StartInternal(CompletionQueue* cq) {
    cq_ = cq;
    if (!next_issue_) {  // ready to issue
      RunNextState(true, nullptr);
    } else {  // wait for the issue time
      alarm_.reset(new Alarm);
      alarm_->Set(cq_, next_issue_(), ClientRpcContext::tag(this));
    }
  }
};

template <class StubType, class RequestType>
class AsyncClient : public ClientImpl<StubType, RequestType> {
  // Specify which protected members we are using since there is no
  // member name resolution until the template types are fully resolved
 public:
  using Client::closed_loop_;
  using Client::NextIssuer;
  using Client::SetupLoadTest;
  using ClientImpl<StubType, RequestType>::cores_;
  using ClientImpl<StubType, RequestType>::channels_;
  using ClientImpl<StubType, RequestType>::request_;
  AsyncClient(const ClientConfig& config,
              std::function<ClientRpcContext*(
                  StubType*, std::function<gpr_timespec()> next_issue,
                  const RequestType&)>
                  setup_ctx,
              std::function<std::unique_ptr<StubType>(std::shared_ptr<Channel>)>
                  create_stub)
      : ClientImpl<StubType, RequestType>(config, create_stub),
        num_async_threads_(NumThreads(config)) {
    SetupLoadTest(config, num_async_threads_);

    int tpc = std::max(1, config.threads_per_cq());      // 1 if unspecified
    int num_cqs = (num_async_threads_ + tpc - 1) / tpc;  // ceiling operator
    for (int i = 0; i < num_cqs; i++) {
      cli_cqs_.emplace_back(new CompletionQueue);
    }

    for (int i = 0; i < num_async_threads_; i++) {
      cq_.emplace_back(i % cli_cqs_.size());
      next_issuers_.emplace_back(NextIssuer(i));
      shutdown_state_.emplace_back(new PerThreadShutdownState());
    }

    int t = 0;
    for (int ch = 0; ch < config.client_channels(); ch++) {
      for (int i = 0; i < config.outstanding_rpcs_per_channel(); i++) {
        auto* cq = cli_cqs_[t].get();
        auto ctx =
            setup_ctx(channels_[ch].get_stub(), next_issuers_[t], request_);
        ctx->Start(cq, config);
      }
      t = (t + 1) % cli_cqs_.size();
    }
  }
  ~AsyncClient() override {
    for (auto cq = cli_cqs_.begin(); cq != cli_cqs_.end(); cq++) {
      void* got_tag;
      bool ok;
      while ((*cq)->Next(&got_tag, &ok)) {
        delete ClientRpcContext::detag(got_tag);
      }
    }
  }

  int GetPollCount() override {
    int count = 0;
    for (auto cq = cli_cqs_.begin(); cq != cli_cqs_.end(); cq++) {
      count += grpc_get_cq_poll_num((*cq)->cq());
    }
    return count;
  }

 protected:
  const int num_async_threads_;

 private:
  struct PerThreadShutdownState {
    mutable std::mutex mutex;
    bool shutdown;
    PerThreadShutdownState() : shutdown(false) {}
  };

  int NumThreads(const ClientConfig& config) {
    int num_threads = config.async_client_threads();
    if (num_threads <= 0) {  // Use dynamic sizing
      num_threads = cores_;
      gpr_log(GPR_INFO, "Sizing async client to %d threads", num_threads);
    }
    return num_threads;
  }
  void DestroyMultithreading() final {
    for (auto ss = shutdown_state_.begin(); ss != shutdown_state_.end(); ++ss) {
      std::lock_guard<std::mutex> lock((*ss)->mutex);
      (*ss)->shutdown = true;
    }
    for (auto cq = cli_cqs_.begin(); cq != cli_cqs_.end(); cq++) {
      (*cq)->Shutdown();
    }
    this->EndThreads();  // this needed for resolution
  }

  ClientRpcContext* ProcessTag(size_t thread_idx, void* tag) {
    ClientRpcContext* ctx = ClientRpcContext::detag(tag);
    if (shutdown_state_[thread_idx]->shutdown) {
      ctx->TryCancel();
      delete ctx;
      bool ok;
      while (cli_cqs_[cq_[thread_idx]]->Next(&tag, &ok)) {
        ctx = ClientRpcContext::detag(tag);
        ctx->TryCancel();
        delete ctx;
      }
      return nullptr;
    }
    return ctx;
  }

  void ThreadFunc(size_t thread_idx, Client::Thread* t) final {
    void* got_tag;
    bool ok;

    HistogramEntry entry;
    HistogramEntry* entry_ptr = &entry;
    if (!cli_cqs_[cq_[thread_idx]]->Next(&got_tag, &ok)) {
      return;
    }
    std::mutex* shutdown_mu = &shutdown_state_[thread_idx]->mutex;
    shutdown_mu->lock();
    ClientRpcContext* ctx = ProcessTag(thread_idx, got_tag);
    if (ctx == nullptr) {
      shutdown_mu->unlock();
      return;
    }
    while (cli_cqs_[cq_[thread_idx]]->DoThenAsyncNext(
        [&, ctx, ok, entry_ptr, shutdown_mu]() {
          if (!ctx->RunNextState(ok, entry_ptr)) {
            // The RPC and callback are done, so clone the ctx
            // and kickstart the new one
            ctx->StartNewClone(cli_cqs_[cq_[thread_idx]].get());
            delete ctx;
          }
          shutdown_mu->unlock();
        },
        &got_tag, &ok, gpr_inf_future(GPR_CLOCK_REALTIME))) {
      t->UpdateHistogram(entry_ptr);
      entry = HistogramEntry();
      shutdown_mu->lock();
      ctx = ProcessTag(thread_idx, got_tag);
      if (ctx == nullptr) {
        shutdown_mu->unlock();
        return;
      }
    }
  }

  std::vector<std::unique_ptr<CompletionQueue>> cli_cqs_;
  std::vector<int> cq_;
  std::vector<std::function<gpr_timespec()>> next_issuers_;
  std::vector<std::unique_ptr<PerThreadShutdownState>> shutdown_state_;
};

static std::unique_ptr<BenchmarkService::Stub> BenchmarkStubCreator(
    const std::shared_ptr<Channel>& ch) {
  return BenchmarkService::NewStub(ch);
}

class AsyncUnaryClient final
    : public AsyncClient<BenchmarkService::Stub, SimpleRequest> {
 public:
  explicit AsyncUnaryClient(const ClientConfig& config)
      : AsyncClient<BenchmarkService::Stub, SimpleRequest>(
            config, SetupCtx, BenchmarkStubCreator) {
    StartThreads(num_async_threads_);
  }
  ~AsyncUnaryClient() override {}

 private:
  static void CheckDone(const grpc::Status& s, SimpleResponse* /*response*/,
                        HistogramEntry* entry) {
    entry->set_status(s.error_code());
  }
  static std::unique_ptr<grpc::ClientAsyncResponseReader<SimpleResponse>>
  PrepareReq(BenchmarkService::Stub* stub, grpc::ClientContext* ctx,
             const SimpleRequest& request, CompletionQueue* cq) {
    return stub->PrepareAsyncUnaryCall(ctx, request, cq);
  };
  static ClientRpcContext* SetupCtx(BenchmarkService::Stub* stub,
                                    std::function<gpr_timespec()> next_issue,
                                    const SimpleRequest& req) {
    return new ClientRpcContextUnaryImpl<SimpleRequest, SimpleResponse>(
        stub, req, std::move(next_issue), AsyncUnaryClient::PrepareReq,
        AsyncUnaryClient::CheckDone);
  }
};

template <class RequestType, class ResponseType>
class ClientRpcContextStreamingPingPongImpl : public ClientRpcContext {
 public:
  ClientRpcContextStreamingPingPongImpl(
      BenchmarkService::Stub* stub, const RequestType& req,
      std::function<gpr_timespec()> next_issue,
      std::function<std::unique_ptr<
          grpc::ClientAsyncReaderWriter<RequestType, ResponseType>>(
          BenchmarkService::Stub*, grpc::ClientContext*, CompletionQueue*)>
          prepare_req,
      std::function<void(grpc::Status, ResponseType*)> on_done)
      : context_(),
        stub_(stub),
        cq_(nullptr),
        req_(req),
        response_(),
        next_state_(State::INVALID),
        callback_(on_done),
        next_issue_(std::move(next_issue)),
        prepare_req_(prepare_req),
        coalesce_(false) {}
  ~ClientRpcContextStreamingPingPongImpl() override {}
  void Start(CompletionQueue* cq, const ClientConfig& config) override {
    StartInternal(cq, config.messages_per_stream(), config.use_coalesce_api());
  }
  bool RunNextState(bool ok, HistogramEntry* entry) override {
    while (true) {
      switch (next_state_) {
        case State::STREAM_IDLE:
          if (!next_issue_) {  // ready to issue
            next_state_ = State::READY_TO_WRITE;
          } else {
            next_state_ = State::WAIT;
          }
          break;  // loop around, don't return
        case State::WAIT:
          next_state_ = State::READY_TO_WRITE;
          alarm_.reset(new Alarm);
          alarm_->Set(cq_, next_issue_(), ClientRpcContext::tag(this));
          return true;
        case State::READY_TO_WRITE:
          if (!ok) {
            return false;
          }
          start_ = UsageTimer::Now();
          next_state_ = State::WRITE_DONE;
          if (coalesce_ && messages_issued_ == messages_per_stream_ - 1) {
            stream_->WriteLast(req_, WriteOptions(),
                               ClientRpcContext::tag(this));
          } else {
            stream_->Write(req_, ClientRpcContext::tag(this));
          }
          return true;
        case State::WRITE_DONE:
          if (!ok) {
            return false;
          }
          next_state_ = State::READ_DONE;
          stream_->Read(&response_, ClientRpcContext::tag(this));
          return true;
          break;
        case State::READ_DONE:
          entry->set_value((UsageTimer::Now() - start_) * 1e9);
          callback_(status_, &response_);
          if ((messages_per_stream_ != 0) &&
              (++messages_issued_ >= messages_per_stream_)) {
            next_state_ = State::WRITES_DONE_DONE;
            if (coalesce_) {
              // WritesDone should have been called on the last Write.
              // loop around to call Finish.
              break;
            }
            stream_->WritesDone(ClientRpcContext::tag(this));
            return true;
          }
          next_state_ = State::STREAM_IDLE;
          break;  // loop around
        case State::WRITES_DONE_DONE:
          next_state_ = State::FINISH_DONE;
          stream_->Finish(&status_, ClientRpcContext::tag(this));
          return true;
        case State::FINISH_DONE:
          next_state_ = State::INVALID;
          return false;
          break;
        default:
          GPR_ASSERT(false);
          return false;
      }
    }
  }
  void StartNewClone(CompletionQueue* cq) override {
    auto* clone = new ClientRpcContextStreamingPingPongImpl(
        stub_, req_, next_issue_, prepare_req_, callback_);
    clone->StartInternal(cq, messages_per_stream_, coalesce_);
  }
  void TryCancel() override { context_.TryCancel(); }

 private:
  grpc::ClientContext context_;
  BenchmarkService::Stub* stub_;
  CompletionQueue* cq_;
  std::unique_ptr<Alarm> alarm_;
  const RequestType& req_;
  ResponseType response_;
  enum State {
    INVALID,
    STREAM_IDLE,
    WAIT,
    READY_TO_WRITE,
    WRITE_DONE,
    READ_DONE,
    WRITES_DONE_DONE,
    FINISH_DONE
  };
  State next_state_;
  std::function<void(grpc::Status, ResponseType*)> callback_;
  std::function<gpr_timespec()> next_issue_;
  std::function<
      std::unique_ptr<grpc::ClientAsyncReaderWriter<RequestType, ResponseType>>(
          BenchmarkService::Stub*, grpc::ClientContext*, CompletionQueue*)>
      prepare_req_;
  grpc::Status status_;
  double start_;
  std::unique_ptr<grpc::ClientAsyncReaderWriter<RequestType, ResponseType>>
      stream_;

  // Allow a limit on number of messages in a stream
  int messages_per_stream_;
  int messages_issued_;
  // Whether to use coalescing API.
  bool coalesce_;

  void StartInternal(CompletionQueue* cq, int messages_per_stream,
                     bool coalesce) {
    cq_ = cq;
    messages_per_stream_ = messages_per_stream;
    messages_issued_ = 0;
    coalesce_ = coalesce;
    if (coalesce_) {
      GPR_ASSERT(messages_per_stream_ != 0);
      context_.set_initial_metadata_corked(true);
    }
    stream_ = prepare_req_(stub_, &context_, cq);
    next_state_ = State::STREAM_IDLE;
    stream_->StartCall(ClientRpcContext::tag(this));
    if (coalesce_) {
      // When the initial metadata is corked, the tag will not come back and we
      // need to manually drive the state machine.
      RunNextState(true, nullptr);
    }
  }
};

class AsyncStreamingPingPongClient final
    : public AsyncClient<BenchmarkService::Stub, SimpleRequest> {
 public:
  explicit AsyncStreamingPingPongClient(const ClientConfig& config)
      : AsyncClient<BenchmarkService::Stub, SimpleRequest>(
            config, SetupCtx, BenchmarkStubCreator) {
    StartThreads(num_async_threads_);
  }

  ~AsyncStreamingPingPongClient() override {}

 private:
  static void CheckDone(const grpc::Status& /*s*/,
                        SimpleResponse* /*response*/) {}
  static std::unique_ptr<
      grpc::ClientAsyncReaderWriter<SimpleRequest, SimpleResponse>>
  PrepareReq(BenchmarkService::Stub* stub, grpc::ClientContext* ctx,
             CompletionQueue* cq) {
    auto stream = stub->PrepareAsyncStreamingCall(ctx, cq);
    return stream;
  };
  static ClientRpcContext* SetupCtx(BenchmarkService::Stub* stub,
                                    std::function<gpr_timespec()> next_issue,
                                    const SimpleRequest& req) {
    return new ClientRpcContextStreamingPingPongImpl<SimpleRequest,
                                                     SimpleResponse>(
        stub, req, std::move(next_issue),
        AsyncStreamingPingPongClient::PrepareReq,
        AsyncStreamingPingPongClient::CheckDone);
  }
};

template <class RequestType, class ResponseType>
class ClientRpcContextStreamingFromClientImpl : public ClientRpcContext {
 public:
  ClientRpcContextStreamingFromClientImpl(
      BenchmarkService::Stub* stub, const RequestType& req,
      std::function<gpr_timespec()> next_issue,
      std::function<std::unique_ptr<grpc::ClientAsyncWriter<RequestType>>(
          BenchmarkService::Stub*, grpc::ClientContext*, ResponseType*,
          CompletionQueue*)>
          prepare_req,
      std::function<void(grpc::Status, ResponseType*)> on_done)
      : context_(),
        stub_(stub),
        cq_(nullptr),
        req_(req),
        response_(),
        next_state_(State::INVALID),
        callback_(on_done),
        next_issue_(std::move(next_issue)),
        prepare_req_(prepare_req) {}
  ~ClientRpcContextStreamingFromClientImpl() override {}
  void Start(CompletionQueue* cq, const ClientConfig& config) override {
    GPR_ASSERT(!config.use_coalesce_api());  // not supported yet.
    StartInternal(cq);
  }
  bool RunNextState(bool ok, HistogramEntry* entry) override {
    while (true) {
      switch (next_state_) {
        case State::STREAM_IDLE:
          if (!next_issue_) {  // ready to issue
            next_state_ = State::READY_TO_WRITE;
          } else {
            next_state_ = State::WAIT;
          }
          break;  // loop around, don't return
        case State::WAIT:
          alarm_.reset(new Alarm);
          alarm_->Set(cq_, next_issue_(), ClientRpcContext::tag(this));
          next_state_ = State::READY_TO_WRITE;
          return true;
        case State::READY_TO_WRITE:
          if (!ok) {
            return false;
          }
          start_ = UsageTimer::Now();
          next_state_ = State::WRITE_DONE;
          stream_->Write(req_, ClientRpcContext::tag(this));
          return true;
        case State::WRITE_DONE:
          if (!ok) {
            return false;
          }
          entry->set_value((UsageTimer::Now() - start_) * 1e9);
          next_state_ = State::STREAM_IDLE;
          break;  // loop around
        default:
          GPR_ASSERT(false);
          return false;
      }
    }
  }
  void StartNewClone(CompletionQueue* cq) override {
    auto* clone = new ClientRpcContextStreamingFromClientImpl(
        stub_, req_, next_issue_, prepare_req_, callback_);
    clone->StartInternal(cq);
  }
  void TryCancel() override { context_.TryCancel(); }

 private:
  grpc::ClientContext context_;
  BenchmarkService::Stub* stub_;
  CompletionQueue* cq_;
  std::unique_ptr<Alarm> alarm_;
  const RequestType& req_;
  ResponseType response_;
  enum State {
    INVALID,
    STREAM_IDLE,
    WAIT,
    READY_TO_WRITE,
    WRITE_DONE,
  };
  State next_state_;
  std::function<void(grpc::Status, ResponseType*)> callback_;
  std::function<gpr_timespec()> next_issue_;
  std::function<std::unique_ptr<grpc::ClientAsyncWriter<RequestType>>(
      BenchmarkService::Stub*, grpc::ClientContext*, ResponseType*,
      CompletionQueue*)>
      prepare_req_;
  grpc::Status status_;
  double start_;
  std::unique_ptr<grpc::ClientAsyncWriter<RequestType>> stream_;

  void StartInternal(CompletionQueue* cq) {
    cq_ = cq;
    stream_ = prepare_req_(stub_, &context_, &response_, cq);
    next_state_ = State::STREAM_IDLE;
    stream_->StartCall(ClientRpcContext::tag(this));
  }
};

class AsyncStreamingFromClientClient final
    : public AsyncClient<BenchmarkService::Stub, SimpleRequest> {
 public:
  explicit AsyncStreamingFromClientClient(const ClientConfig& config)
      : AsyncClient<BenchmarkService::Stub, SimpleRequest>(
            config, SetupCtx, BenchmarkStubCreator) {
    StartThreads(num_async_threads_);
  }

  ~AsyncStreamingFromClientClient() override {}

 private:
  static void CheckDone(const grpc::Status& /*s*/,
                        SimpleResponse* /*response*/) {}
  static std::unique_ptr<grpc::ClientAsyncWriter<SimpleRequest>> PrepareReq(
      BenchmarkService::Stub* stub, grpc::ClientContext* ctx,
      SimpleResponse* resp, CompletionQueue* cq) {
    auto stream = stub->PrepareAsyncStreamingFromClient(ctx, resp, cq);
    return stream;
  };
  static ClientRpcContext* SetupCtx(BenchmarkService::Stub* stub,
                                    std::function<gpr_timespec()> next_issue,
                                    const SimpleRequest& req) {
    return new ClientRpcContextStreamingFromClientImpl<SimpleRequest,
                                                       SimpleResponse>(
        stub, req, std::move(next_issue),
        AsyncStreamingFromClientClient::PrepareReq,
        AsyncStreamingFromClientClient::CheckDone);
  }
};

template <class RequestType, class ResponseType>
class ClientRpcContextStreamingFromServerImpl : public ClientRpcContext {
 public:
  ClientRpcContextStreamingFromServerImpl(
      BenchmarkService::Stub* stub, const RequestType& req,
      std::function<gpr_timespec()> next_issue,
      std::function<std::unique_ptr<grpc::ClientAsyncReader<ResponseType>>(
          BenchmarkService::Stub*, grpc::ClientContext*, const RequestType&,
          CompletionQueue*)>
          prepare_req,
      std::function<void(grpc::Status, ResponseType*)> on_done)
      : context_(),
        stub_(stub),
        cq_(nullptr),
        req_(req),
        response_(),
        next_state_(State::INVALID),
        callback_(on_done),
        next_issue_(std::move(next_issue)),
        prepare_req_(prepare_req) {}
  ~ClientRpcContextStreamingFromServerImpl() override {}
  void Start(CompletionQueue* cq, const ClientConfig& config) override {
    GPR_ASSERT(!config.use_coalesce_api());  // not supported
    StartInternal(cq);
  }
  bool RunNextState(bool ok, HistogramEntry* entry) override {
    while (true) {
      switch (next_state_) {
        case State::STREAM_IDLE:
          if (!ok) {
            return false;
          }
          start_ = UsageTimer::Now();
          next_state_ = State::READ_DONE;
          stream_->Read(&response_, ClientRpcContext::tag(this));
          return true;
        case State::READ_DONE:
          if (!ok) {
            return false;
          }
          entry->set_value((UsageTimer::Now() - start_) * 1e9);
          callback_(status_, &response_);
          next_state_ = State::STREAM_IDLE;
          break;  // loop around
        default:
          GPR_ASSERT(false);
          return false;
      }
    }
  }
  void StartNewClone(CompletionQueue* cq) override {
    auto* clone = new ClientRpcContextStreamingFromServerImpl(
        stub_, req_, next_issue_, prepare_req_, callback_);
    clone->StartInternal(cq);
  }
  void TryCancel() override { context_.TryCancel(); }

 private:
  grpc::ClientContext context_;
  BenchmarkService::Stub* stub_;
  CompletionQueue* cq_;
  std::unique_ptr<Alarm> alarm_;
  const RequestType& req_;
  ResponseType response_;
  enum State { INVALID, STREAM_IDLE, READ_DONE };
  State next_state_;
  std::function<void(grpc::Status, ResponseType*)> callback_;
  std::function<gpr_timespec()> next_issue_;
  std::function<std::unique_ptr<grpc::ClientAsyncReader<ResponseType>>(
      BenchmarkService::Stub*, grpc::ClientContext*, const RequestType&,
      CompletionQueue*)>
      prepare_req_;
  grpc::Status status_;
  double start_;
  std::unique_ptr<grpc::ClientAsyncReader<ResponseType>> stream_;

  void StartInternal(CompletionQueue* cq) {
    // TODO(vjpai): Add support to rate-pace this
    cq_ = cq;
    stream_ = prepare_req_(stub_, &context_, req_, cq);
    next_state_ = State::STREAM_IDLE;
    stream_->StartCall(ClientRpcContext::tag(this));
  }
};

class AsyncStreamingFromServerClient final
    : public AsyncClient<BenchmarkService::Stub, SimpleRequest> {
 public:
  explicit AsyncStreamingFromServerClient(const ClientConfig& config)
      : AsyncClient<BenchmarkService::Stub, SimpleRequest>(
            config, SetupCtx, BenchmarkStubCreator) {
    StartThreads(num_async_threads_);
  }

  ~AsyncStreamingFromServerClient() override {}

 private:
  static void CheckDone(const grpc::Status& /*s*/,
                        SimpleResponse* /*response*/) {}
  static std::unique_ptr<grpc::ClientAsyncReader<SimpleResponse>> PrepareReq(
      BenchmarkService::Stub* stub, grpc::ClientContext* ctx,
      const SimpleRequest& req, CompletionQueue* cq) {
    auto stream = stub->PrepareAsyncStreamingFromServer(ctx, req, cq);
    return stream;
  };
  static ClientRpcContext* SetupCtx(BenchmarkService::Stub* stub,
                                    std::function<gpr_timespec()> next_issue,
                                    const SimpleRequest& req) {
    return new ClientRpcContextStreamingFromServerImpl<SimpleRequest,
                                                       SimpleResponse>(
        stub, req, std::move(next_issue),
        AsyncStreamingFromServerClient::PrepareReq,
        AsyncStreamingFromServerClient::CheckDone);
  }
};

class ClientRpcContextGenericStreamingImpl : public ClientRpcContext {
 public:
  ClientRpcContextGenericStreamingImpl(
      grpc::GenericStub* stub, const ByteBuffer& req,
      std::function<gpr_timespec()> next_issue,
      std::function<std::unique_ptr<grpc::GenericClientAsyncReaderWriter>(
          grpc::GenericStub*, grpc::ClientContext*,
          const std::string& method_name, CompletionQueue*)>
          prepare_req,
      std::function<void(grpc::Status, ByteBuffer*)> on_done)
      : context_(),
        stub_(stub),
        cq_(nullptr),
        req_(req),
        response_(),
        next_state_(State::INVALID),
        callback_(std::move(on_done)),
        next_issue_(std::move(next_issue)),
        prepare_req_(std::move(prepare_req)) {}
  ~ClientRpcContextGenericStreamingImpl() override {}
  void Start(CompletionQueue* cq, const ClientConfig& config) override {
    GPR_ASSERT(!config.use_coalesce_api());  // not supported yet.
    StartInternal(cq, config.messages_per_stream());
  }
  bool RunNextState(bool ok, HistogramEntry* entry) override {
    while (true) {
      switch (next_state_) {
        case State::STREAM_IDLE:
          if (!next_issue_) {  // ready to issue
            next_state_ = State::READY_TO_WRITE;
          } else {
            next_state_ = State::WAIT;
          }
          break;  // loop around, don't return
        case State::WAIT:
          next_state_ = State::READY_TO_WRITE;
          alarm_ = absl::make_unique<Alarm>();
          alarm_->Set(cq_, next_issue_(), ClientRpcContext::tag(this));
          return true;
        case State::READY_TO_WRITE:
          if (!ok) {
            return false;
          }
          start_ = UsageTimer::Now();
          next_state_ = State::WRITE_DONE;
          stream_->Write(req_, ClientRpcContext::tag(this));
          return true;
        case State::WRITE_DONE:
          if (!ok) {
            return false;
          }
          next_state_ = State::READ_DONE;
          stream_->Read(&response_, ClientRpcContext::tag(this));
          return true;
          break;
        case State::READ_DONE:
          entry->set_value((UsageTimer::Now() - start_) * 1e9);
          callback_(status_, &response_);
          if ((messages_per_stream_ != 0) &&
              (++messages_issued_ >= messages_per_stream_)) {
            next_state_ = State::WRITES_DONE_DONE;
            stream_->WritesDone(ClientRpcContext::tag(this));
            return true;
          }
          next_state_ = State::STREAM_IDLE;
          break;  // loop around
        case State::WRITES_DONE_DONE:
          next_state_ = State::FINISH_DONE;
          stream_->Finish(&status_, ClientRpcContext::tag(this));
          return true;
        case State::FINISH_DONE:
          next_state_ = State::INVALID;
          return false;
          break;
        default:
          GPR_ASSERT(false);
          return false;
      }
    }
  }
  void StartNewClone(CompletionQueue* cq) override {
    auto* clone = new ClientRpcContextGenericStreamingImpl(
        stub_, req_, next_issue_, prepare_req_, callback_);
    clone->StartInternal(cq, messages_per_stream_);
  }
  void TryCancel() override { context_.TryCancel(); }

 private:
  grpc::ClientContext context_;
  grpc::GenericStub* stub_;
  CompletionQueue* cq_;
  std::unique_ptr<Alarm> alarm_;
  ByteBuffer req_;
  ByteBuffer response_;
  enum State {
    INVALID,
    STREAM_IDLE,
    WAIT,
    READY_TO_WRITE,
    WRITE_DONE,
    READ_DONE,
    WRITES_DONE_DONE,
    FINISH_DONE
  };
  State next_state_;
  std::function<void(grpc::Status, ByteBuffer*)> callback_;
  std::function<gpr_timespec()> next_issue_;
  std::function<std::unique_ptr<grpc::GenericClientAsyncReaderWriter>(
      grpc::GenericStub*, grpc::ClientContext*, const std::string&,
      CompletionQueue*)>
      prepare_req_;
  grpc::Status status_;
  double start_;
  std::unique_ptr<grpc::GenericClientAsyncReaderWriter> stream_;

  // Allow a limit on number of messages in a stream
  int messages_per_stream_;
  int messages_issued_;

  void StartInternal(CompletionQueue* cq, int messages_per_stream) {
    cq_ = cq;
    const std::string kMethodName(
        "/grpc.testing.BenchmarkService/StreamingCall");
    messages_per_stream_ = messages_per_stream;
    messages_issued_ = 0;
    stream_ = prepare_req_(stub_, &context_, kMethodName, cq);
    next_state_ = State::STREAM_IDLE;
    stream_->StartCall(ClientRpcContext::tag(this));
  }
};

static std::unique_ptr<grpc::GenericStub> GenericStubCreator(
    const std::shared_ptr<Channel>& ch) {
  return absl::make_unique<grpc::GenericStub>(ch);
}

class GenericAsyncStreamingClient final
    : public AsyncClient<grpc::GenericStub, ByteBuffer> {
 public:
  explicit GenericAsyncStreamingClient(const ClientConfig& config)
      : AsyncClient<grpc::GenericStub, ByteBuffer>(config, SetupCtx,
                                                   GenericStubCreator) {
    StartThreads(num_async_threads_);
  }

  ~GenericAsyncStreamingClient() override {}

 private:
  static void CheckDone(const grpc::Status& /*s*/, ByteBuffer* /*response*/) {}
  static std::unique_ptr<grpc::GenericClientAsyncReaderWriter> PrepareReq(
      grpc::GenericStub* stub, grpc::ClientContext* ctx,
      const std::string& method_name, CompletionQueue* cq) {
    auto stream = stub->PrepareCall(ctx, method_name, cq);
    return stream;
  };
  static ClientRpcContext* SetupCtx(grpc::GenericStub* stub,
                                    std::function<gpr_timespec()> next_issue,
                                    const ByteBuffer& req) {
    return new ClientRpcContextGenericStreamingImpl(
        stub, req, std::move(next_issue),
        GenericAsyncStreamingClient::PrepareReq,
        GenericAsyncStreamingClient::CheckDone);
  }
};

std::unique_ptr<Client> CreateAsyncClient(const ClientConfig& config) {
  switch (config.rpc_type()) {
    case UNARY:
      return std::unique_ptr<Client>(new AsyncUnaryClient(config));
    case STREAMING:
      return std::unique_ptr<Client>(new AsyncStreamingPingPongClient(config));
    case STREAMING_FROM_CLIENT:
      return std::unique_ptr<Client>(
          new AsyncStreamingFromClientClient(config));
    case STREAMING_FROM_SERVER:
      return std::unique_ptr<Client>(
          new AsyncStreamingFromServerClient(config));
    case STREAMING_BOTH_WAYS:
      // TODO(vjpai): Implement this
      assert(false);
      return nullptr;
    default:
      assert(false);
      return nullptr;
  }
}
std::unique_ptr<Client> CreateGenericAsyncStreamingClient(
    const ClientConfig& config) {
  return std::unique_ptr<Client>(new GenericAsyncStreamingClient(config));
}

}  // namespace testing
}  // namespace grpc