//
//
// Copyright 2017 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 <string.h>

#include <atomic>

#include <benchmark/benchmark.h>

#include <grpc/grpc.h>
#include <grpc/support/alloc.h>
#include <grpc/support/log.h>

#include "src/core/lib/gprpp/crash.h"
#include "src/core/lib/gprpp/time.h"
#include "src/core/lib/iomgr/ev_posix.h"
#include "src/core/lib/iomgr/port.h"
#include "src/core/lib/surface/completion_queue.h"
#include "test/core/util/test_config.h"
#include "test/cpp/microbenchmarks/helpers.h"
#include "test/cpp/util/test_config.h"

struct grpc_pollset {
  gpr_mu mu;
};

static gpr_mu g_mu;
static gpr_cv g_cv;
static int g_threads_active;
static bool g_active;

namespace grpc {
namespace testing {
static grpc_completion_queue* g_cq;

static void pollset_shutdown(grpc_pollset* /*ps*/, grpc_closure* closure) {
  grpc_core::ExecCtx::Run(DEBUG_LOCATION, closure, absl::OkStatus());
}

static void pollset_init(grpc_pollset* ps, gpr_mu** mu) {
  gpr_mu_init(&ps->mu);
  *mu = &ps->mu;
}

static void pollset_destroy(grpc_pollset* ps) { gpr_mu_destroy(&ps->mu); }

static grpc_error_handle pollset_kick(grpc_pollset* /*p*/,
                                      grpc_pollset_worker* /*worker*/) {
  return absl::OkStatus();
}

// Callback when the tag is dequeued from the completion queue. Does nothing
static void cq_done_cb(void* /*done_arg*/, grpc_cq_completion* cq_completion) {
  gpr_free(cq_completion);
}

// Queues a completion tag if deadline is > 0.
// Does nothing if deadline is 0 (i.e gpr_time_0(GPR_CLOCK_MONOTONIC))
static grpc_error_handle pollset_work(grpc_pollset* ps,
                                      grpc_pollset_worker** /*worker*/,
                                      grpc_core::Timestamp deadline) {
  if (deadline == grpc_core::Timestamp::ProcessEpoch()) {
    gpr_log(GPR_DEBUG, "no-op");
    return absl::OkStatus();
  }

  gpr_mu_unlock(&ps->mu);

  void* tag = reinterpret_cast<void*>(10);  // Some random number
  GPR_ASSERT(grpc_cq_begin_op(g_cq, tag));
  grpc_cq_end_op(
      g_cq, tag, absl::OkStatus(), cq_done_cb, nullptr,
      static_cast<grpc_cq_completion*>(gpr_malloc(sizeof(grpc_cq_completion))));
  grpc_core::ExecCtx::Get()->Flush();
  gpr_mu_lock(&ps->mu);
  return absl::OkStatus();
}

static grpc_event_engine_vtable make_engine_vtable(const char* name) {
  grpc_event_engine_vtable vtable;
  memset(&vtable, 0, sizeof(vtable));

  vtable.pollset_size = sizeof(grpc_pollset);
  vtable.pollset_init = pollset_init;
  vtable.pollset_shutdown = pollset_shutdown;
  vtable.pollset_destroy = pollset_destroy;
  vtable.pollset_work = pollset_work;
  vtable.pollset_kick = pollset_kick;
  vtable.is_any_background_poller_thread = [] { return false; };
  vtable.add_closure_to_background_poller = [](grpc_closure* /*closure*/,
                                               grpc_error_handle /*error*/) {
    return false;
  };
  vtable.shutdown_background_closure = [] {};
  vtable.shutdown_engine = [] {};
  vtable.check_engine_available = [](bool) { return true; };
  vtable.init_engine = [] {};
  vtable.name = name;

  return vtable;
}

static void setup() {
  grpc_init();
  GPR_ASSERT(strcmp(grpc_get_poll_strategy_name(), "none") == 0 ||
             strcmp(grpc_get_poll_strategy_name(), "bm_cq_multiple_threads") ==
                 0);

  g_cq = grpc_completion_queue_create_for_next(nullptr);
}

static void teardown() {
  grpc_completion_queue_shutdown(g_cq);

  // Drain any events
  gpr_timespec deadline = gpr_time_0(GPR_CLOCK_MONOTONIC);
  while (grpc_completion_queue_next(g_cq, deadline, nullptr).type !=
         GRPC_QUEUE_SHUTDOWN) {
    // Do nothing
  }

  grpc_completion_queue_destroy(g_cq);
  grpc_shutdown();
}

// A few notes about Multi-threaded benchmarks:

// Setup:
// The benchmark framework ensures that none of the threads proceed beyond the
// state.KeepRunning() call unless all the threads have called state.keepRunning
// at least once.  So it is safe to do the initialization in one of the threads
// before state.KeepRunning() is called.

// Teardown:
// The benchmark framework also ensures that no thread is running the benchmark
// code (i.e the code between two successive calls of state.KeepRunning()) if
// state.KeepRunning() returns false. So it is safe to do the teardown in one
// of the threads after state.keepRunning() returns false.

// However, our use requires synchronization because we do additional work at
// each thread that requires specific ordering (TrackCounters must be
// constructed after grpc_init because it needs the number of cores, initialized
// by grpc, and its Finish call must take place before grpc_shutdown so that it
// can use grpc_stats).
//
static void BM_Cq_Throughput(benchmark::State& state) {
  gpr_timespec deadline = gpr_inf_future(GPR_CLOCK_MONOTONIC);
  auto thd_idx = state.thread_index();

  gpr_mu_lock(&g_mu);
  g_threads_active++;
  if (thd_idx == 0) {
    setup();
    g_active = true;
    gpr_cv_broadcast(&g_cv);
  } else {
    while (!g_active) {
      gpr_cv_wait(&g_cv, &g_mu, deadline);
    }
  }
  gpr_mu_unlock(&g_mu);

  for (auto _ : state) {
    GPR_ASSERT(grpc_completion_queue_next(g_cq, deadline, nullptr).type ==
               GRPC_OP_COMPLETE);
  }

  state.SetItemsProcessed(state.iterations());

  gpr_mu_lock(&g_mu);
  g_threads_active--;
  if (g_threads_active == 0) {
    gpr_cv_broadcast(&g_cv);
  } else {
    while (g_threads_active > 0) {
      gpr_cv_wait(&g_cv, &g_mu, deadline);
    }
  }
  gpr_mu_unlock(&g_mu);

  if (thd_idx == 0) {
    teardown();
    g_active = false;
  }
}

BENCHMARK(BM_Cq_Throughput)->ThreadRange(1, 16)->UseRealTime();

namespace {
const grpc_event_engine_vtable g_none_vtable =
    grpc::testing::make_engine_vtable("none");
const grpc_event_engine_vtable g_bm_vtable =
    grpc::testing::make_engine_vtable("bm_cq_multiple_threads");
}  // namespace

}  // namespace testing
}  // namespace grpc

// Some distros have RunSpecifiedBenchmarks under the benchmark namespace,
// and others do not. This allows us to support both modes.
namespace benchmark {
void RunTheBenchmarksNamespaced() { RunSpecifiedBenchmarks(); }
}  // namespace benchmark

int main(int argc, char** argv) {
  // This test should only ever be run with a non or any polling engine
  // Override the polling engine for the non-polling engine
  // and add a custom polling engine
  grpc_register_event_engine_factory(&grpc::testing::g_none_vtable, false);
  grpc_register_event_engine_factory(&grpc::testing::g_bm_vtable, true);
  grpc::testing::TestEnvironment env(&argc, argv);
  gpr_mu_init(&g_mu);
  gpr_cv_init(&g_cv);
  ::benchmark::Initialize(&argc, argv);
  grpc::testing::InitTest(&argc, &argv, false);
  benchmark::RunTheBenchmarksNamespaced();
  return 0;
}