#include <stddef.h>
#include <string.h>

#include <cerrno>
#include <cstdint>
#include <string>
#include <vector>

#include "glog/logging.h"
#include "gmock/gmock.h"
#include "gtest/gtest.h"
#include "absl/status/status.h"
#include "absl/status/statusor.h"
#include "absl/time/clock.h"
#include "absl/time/time.h"
#include "absl/types/span.h"
#include "infiniband/verbs.h"
#include "cases/basic_fixture.h"
#include "cases/batch_op_fixture.h"
#include "internal/handle_garble.h"
#include "public/introspection.h"
#include "public/rdma_memblock.h"
#include "public/status_matchers.h"
#include "public/verbs_helper_suite.h"
#include "public/verbs_util.h"

namespace rdma_unit_test {

using ::testing::IsNull;
using ::testing::NotNull;

class CqExTest : public BasicFixture {
 public:
  void SetUp() override {
    if (!Introspection().SupportsExtendedCqs()) {
      GTEST_SKIP() << "NIC does not support extended CQs.";
    }
  }

 protected:
  struct BasicSetup {
    ibv_context* context;
    ibv_pd* pd;
    ibv_comp_channel* channel;
  };

  absl::StatusOr<BasicSetup> CreateBasicSetup() {
    BasicSetup setup;
    ASSIGN_OR_RETURN(setup.context, ibv_.OpenDevice());
    setup.pd = ibv_.AllocPd(setup.context);
    if (!setup.pd) {
      return absl::InternalError("Failed to allocate pd.");
    }
    setup.channel = ibv_.CreateChannel(setup.context);
    if (!setup.channel) {
      return absl::InternalError("Failed to create channel.");
    }
    return setup;
  }
};

TEST_F(CqExTest, Basic) {
  ASSERT_OK_AND_ASSIGN(BasicSetup setup, CreateBasicSetup());
  ibv_cq_init_attr_ex cq_attr{.cqe = 10};
  ibv_cq_ex* cq = ibv_create_cq_ex(setup.context, &cq_attr);
  ASSERT_THAT(cq, NotNull());
  ASSERT_EQ(ibv_destroy_cq(ibv_cq_ex_to_cq(cq)), 0);
}

TEST_F(CqExTest, DestroyInvalidCqEx) {
  ASSERT_OK_AND_ASSIGN(BasicSetup setup, CreateBasicSetup());
  ibv_cq_init_attr_ex attr{.cqe = 10};
  ibv_cq_ex* cq = ibv_.CreateCqEx(setup.context, attr);
  ASSERT_THAT(cq, NotNull());
  HandleGarble garble(cq->handle);
  ASSERT_EQ(ibv_destroy_cq(ibv_cq_ex_to_cq(cq)), ENOENT);
}

TEST_F(CqExTest, ZeroCqe) {
  ASSERT_OK_AND_ASSIGN(BasicSetup setup, CreateBasicSetup());
  ibv_cq_init_attr_ex cq_attr{.cqe = 0};
  ibv_cq_ex* cq = ibv_create_cq_ex(setup.context, &cq_attr);
  ASSERT_THAT(cq, IsNull());
}

TEST_F(CqExTest, MaxCqe) {
  ASSERT_OK_AND_ASSIGN(BasicSetup setup, CreateBasicSetup());
  uint32_t max_cqe = Introspection().device_attr().max_cqe;
  EXPECT_THAT(ibv_.CreateCqEx(setup.context, max_cqe), NotNull());
  EXPECT_THAT(ibv_.CreateCqEx(setup.context, max_cqe + 1), IsNull());
}

TEST_F(CqExTest, WithChannel) {
  ASSERT_OK_AND_ASSIGN(BasicSetup setup, CreateBasicSetup());
  ibv_cq_init_attr_ex cq_attr{.cqe = 10, .channel = setup.channel};
  ibv_cq_ex* cq = ibv_create_cq_ex(setup.context, &cq_attr);
  ASSERT_THAT(cq, NotNull());
  ASSERT_EQ(ibv_destroy_cq(ibv_cq_ex_to_cq(cq)), 0);
}

TEST_F(CqExTest, ShareChannel) {
  ASSERT_OK_AND_ASSIGN(BasicSetup setup, CreateBasicSetup());
  ibv_cq_init_attr_ex cq_attr{.cqe = 10, .channel = setup.channel};
  ibv_cq_ex* cq1 = ibv_create_cq_ex(setup.context, &cq_attr);
  ASSERT_THAT(cq1, NotNull());
  ibv_cq_ex* cq2 = ibv_create_cq_ex(setup.context, &cq_attr);
  ASSERT_THAT(cq2, NotNull());
  ASSERT_EQ(ibv_destroy_cq(ibv_cq_ex_to_cq(cq1)), 0);
  ASSERT_EQ(ibv_destroy_cq(ibv_cq_ex_to_cq(cq2)), 0);
}

TEST_F(CqExTest, ZeroCompVector) {
  ASSERT_OK_AND_ASSIGN(BasicSetup setup, CreateBasicSetup());
  ibv_cq_init_attr_ex cq_attr{.cqe = 10};
  ibv_cq_ex* cq = ibv_create_cq_ex(setup.context, &cq_attr);
  ASSERT_THAT(cq, NotNull());
  ASSERT_EQ(ibv_destroy_cq(ibv_cq_ex_to_cq(cq)), 0);
}

TEST_F(CqExTest, LargeCompVector) {
  ASSERT_OK_AND_ASSIGN(BasicSetup setup, CreateBasicSetup());
  ibv_cq_init_attr_ex cq_attr{.cqe = 10,
                              .comp_vector = static_cast<uint32_t>(
                                  setup.context->num_comp_vectors - 1)};
  ibv_cq_ex* cq = ibv_create_cq_ex(setup.context, &cq_attr);
  ASSERT_THAT(cq, NotNull());
  ASSERT_EQ(ibv_destroy_cq(ibv_cq_ex_to_cq(cq)), 0);
}

TEST_F(CqExTest, AboveMaxCompVector) {
  ASSERT_OK_AND_ASSIGN(BasicSetup setup, CreateBasicSetup());
  ibv_cq_init_attr_ex cq_attr{
      .cqe = 10,
      .comp_vector = static_cast<uint32_t>(setup.context->num_comp_vectors)};
  ibv_cq_ex* cq = ibv_create_cq_ex(setup.context, &cq_attr);
  ASSERT_THAT(cq, IsNull());
}

class CqExOpTest : public BatchOpFixture {
 public:
  void SetUp() override {
    if (!Introspection().SupportsExtendedCqs()) {
      GTEST_SKIP() << "NIC does not support extended CQs.";
    }
  }

 protected:
  struct WcInfo {
    uint64_t wr_id;
    ibv_wc_status status;
    uint32_t qp_num;
    uint64_t timestamp_hcaclock = 0;
    uint64_t timestamp_wallclock = 0;
  };

  absl::StatusOr<BasicSetup> CreateBasicSetup() {
    ASSIGN_OR_RETURN(BasicSetup setup, BatchOpFixture::CreateBasicSetup());

    support_completion_timestamp_hcaclock_ =
        verbs_util::CheckExtendedCompletionHasCapability(
            setup.context, IBV_WC_EX_WITH_COMPLETION_TIMESTAMP);
    support_completion_timestamp_wallclock_ =
        verbs_util::CheckExtendedCompletionHasCapability(
            setup.context, IBV_WC_EX_WITH_COMPLETION_TIMESTAMP_WALLCLOCK);
    return setup;
  }

  uint64_t GetWcFlags() {
    uint64_t wc_flags = IBV_WC_EX_WITH_QP_NUM;
    if (support_completion_timestamp_hcaclock_) {
      wc_flags |= IBV_WC_EX_WITH_COMPLETION_TIMESTAMP;
    }
    if (support_completion_timestamp_wallclock_) {
      wc_flags |= IBV_WC_EX_WITH_COMPLETION_TIMESTAMP_WALLCLOCK;
    }
    return wc_flags;
  }

  WcInfo GetWcInfo(ibv_cq_ex* cq) {
    WcInfo wc{.wr_id = cq->wr_id,
              .status = cq->status,
              .qp_num = ibv_wc_read_qp_num(cq)};
    if (support_completion_timestamp_hcaclock_) {
      wc.timestamp_hcaclock = ibv_wc_read_completion_ts(cq);
    }
    if (support_completion_timestamp_wallclock_) {
      wc.timestamp_wallclock = ibv_wc_read_completion_wallclock_ns(cq);
    }
    return wc;
  }

  void WaitForAndVerifyCompletions(ibv_cq_ex* cq, int count) {
    ASSERT_OK(verbs_util::WaitForPollingExtendedCompletion(cq));
    WcInfo wc = GetWcInfo(cq);
    ASSERT_EQ(wc.status, IBV_WC_SUCCESS);
    if (support_completion_timestamp_hcaclock_) {
      EXPECT_GT(wc.timestamp_hcaclock, 0);
    }
    if (support_completion_timestamp_wallclock_) {
      EXPECT_GT(wc.timestamp_wallclock, 0);
    }

    absl::flat_hash_map<uint32_t, uint64_t> expected_wr_id;
    expected_wr_id.insert({wc.qp_num, wc.wr_id + 1});
    uint64_t last_timestamp_hcaclock = 0;
    uint64_t last_timestamp_wallclock = 0;
    for (int i = 1; i < count; ++i) {
      ASSERT_OK(verbs_util::WaitForNextExtendedCompletion(cq));
      WcInfo wc = GetWcInfo(cq);
      ASSERT_EQ(wc.status, IBV_WC_SUCCESS);
      auto iter = expected_wr_id.find(wc.qp_num);
      if (iter == expected_wr_id.end()) {
        expected_wr_id.insert({wc.qp_num, wc.wr_id + 1});
      } else {
        ASSERT_EQ(wc.wr_id, iter->second++);
      }
      EXPECT_GE(wc.timestamp_hcaclock, last_timestamp_hcaclock);
      EXPECT_GE(wc.timestamp_wallclock, last_timestamp_wallclock);
      last_timestamp_hcaclock = wc.timestamp_hcaclock;
      last_timestamp_wallclock = wc.timestamp_wallclock;
    }
    ibv_end_poll(cq);
  }

  bool support_completion_timestamp_hcaclock_ = false;

  bool support_completion_timestamp_wallclock_ = false;
};

TEST_F(CqExOpTest, BasicPollSendCq) {
  ASSERT_OK_AND_ASSIGN(BasicSetup setup, CreateBasicSetup());
  static constexpr int kMaxQpWr = 1;
  ibv_cq_init_attr_ex cq_attr = {
      .cqe = static_cast<uint32_t>(Introspection().device_attr().max_cqe),
      .wc_flags = GetWcFlags()};
  ibv_cq_ex* send_cq = ibv_.CreateCqEx(setup.context, cq_attr);
  ibv_cq_ex* recv_cq = ibv_.CreateCqEx(setup.context, cq_attr);
  QpPair qp_pair;
  qp_pair.send_qp = ibv_.CreateQp(setup.pd, ibv_cq_ex_to_cq(send_cq),
                                  ibv_cq_ex_to_cq(recv_cq), nullptr, kMaxQpWr,
                                  kMaxQpWr, IBV_QPT_RC, 0);
  ASSERT_THAT(qp_pair.send_qp, NotNull())
      << "Failed to create send qp - " << errno;
  qp_pair.recv_qp = ibv_.CreateQp(setup.pd, ibv_cq_ex_to_cq(send_cq),
                                  ibv_cq_ex_to_cq(recv_cq), nullptr, kMaxQpWr,
                                  kMaxQpWr, IBV_QPT_RC, 0);
  ASSERT_THAT(qp_pair.recv_qp, NotNull())
      << "Failed to create recv qp - " << errno;
  qp_pair.dst_buffer = setup.dst_memblock.subspan(kMaxQpWr, kMaxQpWr);
  ibv_.SetUpLoopbackRcQps(qp_pair.send_qp, qp_pair.recv_qp, setup.port_gid);
  QueueWrite(setup, qp_pair);

  ASSERT_OK(verbs_util::WaitForPollingExtendedCompletion(send_cq));
  ASSERT_EQ(send_cq->status, IBV_WC_SUCCESS);
  if (support_completion_timestamp_hcaclock_) {
    EXPECT_GT(ibv_wc_read_completion_ts(send_cq), 0);
  }
  if (support_completion_timestamp_wallclock_) {
    EXPECT_GT(ibv_wc_read_completion_wallclock_ns(send_cq), 0);
  }
  ibv_end_poll(send_cq);
}

TEST_F(CqExOpTest, BasicPollRecvCq) {
  ASSERT_OK_AND_ASSIGN(BasicSetup setup, CreateBasicSetup());
  static constexpr int kMaxQpWr = 1;
  ibv_cq_init_attr_ex cq_attr = {
      .cqe = static_cast<uint32_t>(Introspection().device_attr().max_cqe),
      .wc_flags = GetWcFlags()};
  ibv_cq_ex* send_cq = ibv_.CreateCqEx(setup.context, cq_attr);
  ibv_cq_ex* recv_cq = ibv_.CreateCqEx(setup.context, cq_attr);
  QpPair qp_pair;
  qp_pair.send_qp = ibv_.CreateQp(setup.pd, ibv_cq_ex_to_cq(send_cq),
                                  ibv_cq_ex_to_cq(recv_cq), nullptr, kMaxQpWr,
                                  kMaxQpWr, IBV_QPT_RC, 0);
  ASSERT_THAT(qp_pair.send_qp, NotNull())
      << "Failed to create send qp - " << errno;
  qp_pair.recv_qp = ibv_.CreateQp(setup.pd, ibv_cq_ex_to_cq(send_cq),
                                  ibv_cq_ex_to_cq(recv_cq), nullptr, kMaxQpWr,
                                  kMaxQpWr, IBV_QPT_RC, 0);
  ASSERT_THAT(qp_pair.recv_qp, NotNull())
      << "Failed to create recv qp - " << errno;
  qp_pair.dst_buffer = setup.dst_memblock.subspan(kMaxQpWr, kMaxQpWr);
  ibv_.SetUpLoopbackRcQps(qp_pair.send_qp, qp_pair.recv_qp, setup.port_gid);
  QueueRecv(setup, qp_pair);
  QueueSend(setup, qp_pair);

  ASSERT_OK(verbs_util::WaitForPollingExtendedCompletion(recv_cq));
  ASSERT_EQ(recv_cq->status, IBV_WC_SUCCESS);
  if (support_completion_timestamp_hcaclock_) {
    EXPECT_GT(ibv_wc_read_completion_ts(recv_cq), 0);
  }
  if (support_completion_timestamp_wallclock_) {
    EXPECT_GT(ibv_wc_read_completion_wallclock_ns(recv_cq), 0);
  }
  ibv_end_poll(recv_cq);
}

TEST_F(CqExOpTest, BatchPollSendCq) {
  ASSERT_OK_AND_ASSIGN(BasicSetup setup, CreateBasicSetup());
  static constexpr int kCqSize = 60;
  static constexpr int kQpPairCount = 20;
  static constexpr int kWritesPerQpPair = kCqSize / kQpPairCount;
  ibv_cq_init_attr_ex cq_attr = {
      .cqe = static_cast<uint32_t>(Introspection().device_attr().max_cqe),
      .wc_flags = GetWcFlags()};
  ibv_cq_ex* send_cq = ibv_.CreateCqEx(setup.context, cq_attr);
  ibv_cq_ex* recv_cq = ibv_.CreateCqEx(setup.context, cq_attr);
  std::vector<QpPair> qp_pairs = CreateTestQpPairs(
      setup, ibv_cq_ex_to_cq(send_cq), ibv_cq_ex_to_cq(recv_cq),
      kWritesPerQpPair, kQpPairCount);

  for (auto& qp_pair : qp_pairs) {
    for (int i = 0; i < kWritesPerQpPair; ++i) {
      QueueWrite(setup, qp_pair);
    }
  }
  WaitForAndVerifyCompletions(send_cq, kWritesPerQpPair * kQpPairCount);
}

TEST_F(CqExOpTest, BatchPollRecvCq) {
  ASSERT_OK_AND_ASSIGN(BasicSetup setup, CreateBasicSetup());
  static constexpr int kCqSize = 60;
  static constexpr int kQpPairCount = 20;
  static constexpr int kSendsPerQpPair = kCqSize / kQpPairCount;
  ibv_cq_init_attr_ex cq_attr = {
      .cqe = static_cast<uint32_t>(Introspection().device_attr().max_cqe),
      .wc_flags = GetWcFlags()};
  ibv_cq_ex* send_cq = ibv_.CreateCqEx(setup.context, cq_attr);
  ibv_cq_ex* recv_cq = ibv_.CreateCqEx(setup.context, cq_attr);
  std::vector<QpPair> qp_pairs = CreateTestQpPairs(
      setup, ibv_cq_ex_to_cq(send_cq), ibv_cq_ex_to_cq(recv_cq),
      kSendsPerQpPair, kQpPairCount);

  for (auto& qp_pair : qp_pairs) {
    for (int i = 0; i < kSendsPerQpPair; ++i) {
      QueueRecv(setup, qp_pair);
    }
  }
  for (auto& qp_pair : qp_pairs) {
    for (int i = 0; i < kSendsPerQpPair; ++i) {
      QueueSend(setup, qp_pair);
    }
  }
  WaitForAndVerifyCompletions(recv_cq, kSendsPerQpPair);
}

class CqExOverflowTest : public CqExOpTest {
 protected:

  static constexpr absl::Duration kPollTime = absl::Seconds(10);

  static constexpr absl::Duration kCompletionWaitTime = absl::Seconds(1);

  void WaitForData(absl::Span<uint8_t> dst_buffer, uint8_t expected_value,
                   absl::Duration poll_timeout = kPollTime) {
    auto read_value = [](uint8_t* target) {
      return reinterpret_cast<volatile std::atomic<uint8_t>*>(target)->load();
    };
    absl::Time stop = absl::Now() + poll_timeout;
    bool completed = false;
    for (absl::Time now = absl::Now(); now < stop; now = absl::Now()) {
      completed = true;
      for (size_t i = 0; i < dst_buffer.size(); ++i) {
        if (read_value(dst_buffer.data() + i) != expected_value) {
          completed = false;
        }
      }
      if (completed == false) {
        absl::SleepFor(absl::Milliseconds(10));
      } else {
        return;
      }
    }
    LOG(FATAL) << "Data failed to land.";  
  }
};

TEST_F(CqExOverflowTest, SendCqOverflow) {
  if (Introspection().FullCqIdlesQp()) {
    GTEST_SKIP() << "This test assumes CQ overflow overwrites completions.";
  }
  ASSERT_OK_AND_ASSIGN(BasicSetup setup, CreateBasicSetup());
  ibv_cq_init_attr_ex cq_attr = {.cqe = verbs_util::kDefaultMaxWr,
                                 .wc_flags = GetWcFlags()};
  ibv_cq_ex* cq = ibv_.CreateCqEx(setup.context, cq_attr);
  const int total_writes = cq->cqe + 10;
  std::vector<QpPair> qp_pairs =
      CreateTestQpPairs(setup, ibv_cq_ex_to_cq(cq), ibv_cq_ex_to_cq(cq),
                        total_writes, 1);
  for (int i = 0; i < total_writes; ++i) {
    QueueWrite(setup, qp_pairs[0]);
  }
  WaitForData(setup.dst_memblock.subspan(0, total_writes), kSrcContent);
  absl::SleepFor(kCompletionWaitTime);
  WaitForAndVerifyCompletions(cq, cq->cqe);
  verbs_util::ExpectNoExtendedCompletion(cq, absl::Seconds(1));
}

TEST_F(CqExOverflowTest, RecvCqOverflow) {
  if (Introspection().FullCqIdlesQp()) {
    GTEST_SKIP() << "This test assumes CQ overflow overwrites completions.";
  }
  ASSERT_OK_AND_ASSIGN(BasicSetup setup, CreateBasicSetup());
  ibv_cq_init_attr_ex cq_attr = {.cqe = verbs_util::kDefaultMaxWr,
                                 .wc_flags = GetWcFlags()};
  ibv_cq_ex* send_cq = ibv_.CreateCqEx(setup.context, cq_attr);
  ibv_cq_ex* recv_cq = ibv_.CreateCqEx(setup.context, cq_attr);
  const int total_sends = recv_cq->cqe + 10;
  std::vector<QpPair> qp_pairs =
      CreateTestQpPairs(setup, ibv_cq_ex_to_cq(send_cq),
                        ibv_cq_ex_to_cq(recv_cq), total_sends, 1);
  for (int i = 0; i < total_sends; ++i) {
    QueueRecv(setup, qp_pairs[0]);
  }
  for (int i = 0; i < total_sends; ++i) {
    QueueSend(setup, qp_pairs[0]);
  }
  WaitForData(setup.dst_memblock.subspan(0, total_sends), kSrcContent);
  absl::SleepFor(kCompletionWaitTime);
  WaitForAndVerifyCompletions(recv_cq, recv_cq->cqe);
  verbs_util::ExpectNoExtendedCompletion(recv_cq, absl::Seconds(1));
}

}
// Copyright 2021 Google LLC
//
// 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 <stddef.h>
#include <string.h>

#include <cerrno>
#include <cstdint>
#include <string>
#include <vector>

#include "glog/logging.h"
#include "gmock/gmock.h"
#include "gtest/gtest.h"
#include "absl/status/status.h"
#include "absl/status/statusor.h"
#include "absl/time/clock.h"
#include "absl/time/time.h"
#include "absl/types/span.h"
#include "infiniband/verbs.h"
#include "cases/basic_fixture.h"
#include "cases/batch_op_fixture.h"
#include "internal/handle_garble.h"
#include "public/introspection.h"
#include "public/rdma_memblock.h"
#include "public/status_matchers.h"
#include "public/verbs_helper_suite.h"
#include "public/verbs_util.h"

namespace rdma_unit_test {

using ::testing::IsNull;
using ::testing::NotNull;

class CqExTest : public BasicFixture {
 public:
  void SetUp() override {
    if (!Introspection().SupportsExtendedCqs()) {
      GTEST_SKIP() << "NIC does not support extended CQs.";
    }
  }

 protected:
  struct BasicSetup {
    ibv_context* context;
    ibv_pd* pd;
    ibv_comp_channel* channel;
  };

  absl::StatusOr<BasicSetup> CreateBasicSetup() {
    BasicSetup setup;
    ASSIGN_OR_RETURN(setup.context, ibv_.OpenDevice());
    setup.pd = ibv_.AllocPd(setup.context);
    if (!setup.pd) {
      return absl::InternalError("Failed to allocate pd.");
    }
    setup.channel = ibv_.CreateChannel(setup.context);
    if (!setup.channel) {
      return absl::InternalError("Failed to create channel.");
    }
    return setup;
  }
};

TEST_F(CqExTest, Basic) {
  ASSERT_OK_AND_ASSIGN(BasicSetup setup, CreateBasicSetup());
  ibv_cq_init_attr_ex cq_attr{.cqe = 10};
  ibv_cq_ex* cq = ibv_create_cq_ex(setup.context, &cq_attr);
  ASSERT_THAT(cq, NotNull());
  ASSERT_EQ(ibv_destroy_cq(ibv_cq_ex_to_cq(cq)), 0);
}

TEST_F(CqExTest, DestroyInvalidCqEx) {
  ASSERT_OK_AND_ASSIGN(BasicSetup setup, CreateBasicSetup());
  ibv_cq_init_attr_ex attr{.cqe = 10};
  ibv_cq_ex* cq = ibv_.CreateCqEx(setup.context, attr);
  ASSERT_THAT(cq, NotNull());
  HandleGarble garble(cq->handle);
  ASSERT_EQ(ibv_destroy_cq(ibv_cq_ex_to_cq(cq)), ENOENT);
}

TEST_F(CqExTest, ZeroCqe) {
  ASSERT_OK_AND_ASSIGN(BasicSetup setup, CreateBasicSetup());
  ibv_cq_init_attr_ex cq_attr{.cqe = 0};
  ibv_cq_ex* cq = ibv_create_cq_ex(setup.context, &cq_attr);
  ASSERT_THAT(cq, IsNull());
}

TEST_F(CqExTest, MaxCqe) {
  ASSERT_OK_AND_ASSIGN(BasicSetup setup, CreateBasicSetup());
  uint32_t max_cqe = Introspection().device_attr().max_cqe;
  EXPECT_THAT(ibv_.CreateCqEx(setup.context, max_cqe), NotNull());
  EXPECT_THAT(ibv_.CreateCqEx(setup.context, max_cqe + 1), IsNull());
}

TEST_F(CqExTest, WithChannel) {
  ASSERT_OK_AND_ASSIGN(BasicSetup setup, CreateBasicSetup());
  ibv_cq_init_attr_ex cq_attr{.cqe = 10, .channel = setup.channel};
  ibv_cq_ex* cq = ibv_create_cq_ex(setup.context, &cq_attr);
  ASSERT_THAT(cq, NotNull());
  ASSERT_EQ(ibv_destroy_cq(ibv_cq_ex_to_cq(cq)), 0);
}

TEST_F(CqExTest, ShareChannel) {
  ASSERT_OK_AND_ASSIGN(BasicSetup setup, CreateBasicSetup());
  ibv_cq_init_attr_ex cq_attr{.cqe = 10, .channel = setup.channel};
  ibv_cq_ex* cq1 = ibv_create_cq_ex(setup.context, &cq_attr);
  ASSERT_THAT(cq1, NotNull());
  ibv_cq_ex* cq2 = ibv_create_cq_ex(setup.context, &cq_attr);
  ASSERT_THAT(cq2, NotNull());
  ASSERT_EQ(ibv_destroy_cq(ibv_cq_ex_to_cq(cq1)), 0);
  ASSERT_EQ(ibv_destroy_cq(ibv_cq_ex_to_cq(cq2)), 0);
}

TEST_F(CqExTest, ZeroCompVector) {
  ASSERT_OK_AND_ASSIGN(BasicSetup setup, CreateBasicSetup());
  ibv_cq_init_attr_ex cq_attr{.cqe = 10};
  ibv_cq_ex* cq = ibv_create_cq_ex(setup.context, &cq_attr);
  ASSERT_THAT(cq, NotNull());
  ASSERT_EQ(ibv_destroy_cq(ibv_cq_ex_to_cq(cq)), 0);
}

TEST_F(CqExTest, LargeCompVector) {
  ASSERT_OK_AND_ASSIGN(BasicSetup setup, CreateBasicSetup());
  ibv_cq_init_attr_ex cq_attr{.cqe = 10,
                              .comp_vector = static_cast<uint32_t>(
                                  setup.context->num_comp_vectors - 1)};
  ibv_cq_ex* cq = ibv_create_cq_ex(setup.context, &cq_attr);
  ASSERT_THAT(cq, NotNull());
  ASSERT_EQ(ibv_destroy_cq(ibv_cq_ex_to_cq(cq)), 0);
}

TEST_F(CqExTest, AboveMaxCompVector) {
  ASSERT_OK_AND_ASSIGN(BasicSetup setup, CreateBasicSetup());
  ibv_cq_init_attr_ex cq_attr{
      .cqe = 10,
      .comp_vector = static_cast<uint32_t>(setup.context->num_comp_vectors)};
  ibv_cq_ex* cq = ibv_create_cq_ex(setup.context, &cq_attr);
  ASSERT_THAT(cq, IsNull());
}

class CqExOpTest : public BatchOpFixture {
 public:
  void SetUp() override {
    if (!Introspection().SupportsExtendedCqs()) {
      GTEST_SKIP() << "NIC does not support extended CQs.";
    }
  }

 protected:
  struct WcInfo {
    uint64_t wr_id;
    ibv_wc_status status;
    uint32_t qp_num;
    uint64_t timestamp_hcaclock = 0;
    uint64_t timestamp_wallclock = 0;
  };

  absl::StatusOr<BasicSetup> CreateBasicSetup() {
    ASSIGN_OR_RETURN(BasicSetup setup, BatchOpFixture::CreateBasicSetup());
    // Set the supported timestamp types.
    support_completion_timestamp_hcaclock_ =
        verbs_util::CheckExtendedCompletionHasCapability(
            setup.context, IBV_WC_EX_WITH_COMPLETION_TIMESTAMP);
    support_completion_timestamp_wallclock_ =
        verbs_util::CheckExtendedCompletionHasCapability(
            setup.context, IBV_WC_EX_WITH_COMPLETION_TIMESTAMP_WALLCLOCK);
    return setup;
  }

  uint64_t GetWcFlags() {
    uint64_t wc_flags = IBV_WC_EX_WITH_QP_NUM;
    if (support_completion_timestamp_hcaclock_) {
      wc_flags |= IBV_WC_EX_WITH_COMPLETION_TIMESTAMP;
    }
    if (support_completion_timestamp_wallclock_) {
      wc_flags |= IBV_WC_EX_WITH_COMPLETION_TIMESTAMP_WALLCLOCK;
    }
    return wc_flags;
  }

  WcInfo GetWcInfo(ibv_cq_ex* cq) {
    WcInfo wc{.wr_id = cq->wr_id,
              .status = cq->status,
              .qp_num = ibv_wc_read_qp_num(cq)};
    if (support_completion_timestamp_hcaclock_) {
      wc.timestamp_hcaclock = ibv_wc_read_completion_ts(cq);
    }
    if (support_completion_timestamp_wallclock_) {
      wc.timestamp_wallclock = ibv_wc_read_completion_wallclock_ns(cq);
    }
    return wc;
  }

  void WaitForAndVerifyCompletions(ibv_cq_ex* cq, int count) {
    ASSERT_OK(verbs_util::WaitForPollingExtendedCompletion(cq));
    WcInfo wc = GetWcInfo(cq);
    ASSERT_EQ(wc.status, IBV_WC_SUCCESS);
    if (support_completion_timestamp_hcaclock_) {
      EXPECT_GT(wc.timestamp_hcaclock, 0);
    }
    if (support_completion_timestamp_wallclock_) {
      EXPECT_GT(wc.timestamp_wallclock, 0);
    }

    // Maps qp_num to the next wr_id.
    absl::flat_hash_map<uint32_t, uint64_t> expected_wr_id;
    expected_wr_id.insert({wc.qp_num, wc.wr_id + 1});
    uint64_t last_timestamp_hcaclock = 0;
    uint64_t last_timestamp_wallclock = 0;
    for (int i = 1; i < count; ++i) {
      ASSERT_OK(verbs_util::WaitForNextExtendedCompletion(cq));
      WcInfo wc = GetWcInfo(cq);
      ASSERT_EQ(wc.status, IBV_WC_SUCCESS);
      auto iter = expected_wr_id.find(wc.qp_num);
      if (iter == expected_wr_id.end()) {
        expected_wr_id.insert({wc.qp_num, wc.wr_id + 1});
      } else {
        ASSERT_EQ(wc.wr_id, iter->second++);
      }
      EXPECT_GE(wc.timestamp_hcaclock, last_timestamp_hcaclock);
      EXPECT_GE(wc.timestamp_wallclock, last_timestamp_wallclock);
      last_timestamp_hcaclock = wc.timestamp_hcaclock;
      last_timestamp_wallclock = wc.timestamp_wallclock;
    }
    ibv_end_poll(cq);
  }

  // Whether WCs support HCA clock timestamps or not.
  bool support_completion_timestamp_hcaclock_ = false;
  // Whether WCs support wallclock timestamps or not.
  bool support_completion_timestamp_wallclock_ = false;
};

TEST_F(CqExOpTest, BasicPollSendCq) {
  ASSERT_OK_AND_ASSIGN(BasicSetup setup, CreateBasicSetup());
  static constexpr int kMaxQpWr = 1;
  ibv_cq_init_attr_ex cq_attr = {
      .cqe = static_cast<uint32_t>(Introspection().device_attr().max_cqe),
      .wc_flags = GetWcFlags()};
  ibv_cq_ex* send_cq = ibv_.CreateCqEx(setup.context, cq_attr);
  ibv_cq_ex* recv_cq = ibv_.CreateCqEx(setup.context, cq_attr);
  QpPair qp_pair;
  qp_pair.send_qp = ibv_.CreateQp(setup.pd, ibv_cq_ex_to_cq(send_cq),
                                  ibv_cq_ex_to_cq(recv_cq), nullptr, kMaxQpWr,
                                  kMaxQpWr, IBV_QPT_RC, /*sig_all=*/0);
  ASSERT_THAT(qp_pair.send_qp, NotNull())
      << "Failed to create send qp - " << errno;
  qp_pair.recv_qp = ibv_.CreateQp(setup.pd, ibv_cq_ex_to_cq(send_cq),
                                  ibv_cq_ex_to_cq(recv_cq), nullptr, kMaxQpWr,
                                  kMaxQpWr, IBV_QPT_RC, /*sig_all=*/0);
  ASSERT_THAT(qp_pair.recv_qp, NotNull())
      << "Failed to create recv qp - " << errno;
  qp_pair.dst_buffer = setup.dst_memblock.subspan(kMaxQpWr, kMaxQpWr);
  ibv_.SetUpLoopbackRcQps(qp_pair.send_qp, qp_pair.recv_qp, setup.port_gid);
  QueueWrite(setup, qp_pair);

  // Wait for completion and verify timestamp.
  ASSERT_OK(verbs_util::WaitForPollingExtendedCompletion(send_cq));
  ASSERT_EQ(send_cq->status, IBV_WC_SUCCESS);
  if (support_completion_timestamp_hcaclock_) {
    EXPECT_GT(ibv_wc_read_completion_ts(send_cq), 0);
  }
  if (support_completion_timestamp_wallclock_) {
    EXPECT_GT(ibv_wc_read_completion_wallclock_ns(send_cq), 0);
  }
  ibv_end_poll(send_cq);
}

TEST_F(CqExOpTest, BasicPollRecvCq) {
  ASSERT_OK_AND_ASSIGN(BasicSetup setup, CreateBasicSetup());
  static constexpr int kMaxQpWr = 1;
  ibv_cq_init_attr_ex cq_attr = {
      .cqe = static_cast<uint32_t>(Introspection().device_attr().max_cqe),
      .wc_flags = GetWcFlags()};
  ibv_cq_ex* send_cq = ibv_.CreateCqEx(setup.context, cq_attr);
  ibv_cq_ex* recv_cq = ibv_.CreateCqEx(setup.context, cq_attr);
  QpPair qp_pair;
  qp_pair.send_qp = ibv_.CreateQp(setup.pd, ibv_cq_ex_to_cq(send_cq),
                                  ibv_cq_ex_to_cq(recv_cq), nullptr, kMaxQpWr,
                                  kMaxQpWr, IBV_QPT_RC, /*sig_all=*/0);
  ASSERT_THAT(qp_pair.send_qp, NotNull())
      << "Failed to create send qp - " << errno;
  qp_pair.recv_qp = ibv_.CreateQp(setup.pd, ibv_cq_ex_to_cq(send_cq),
                                  ibv_cq_ex_to_cq(recv_cq), nullptr, kMaxQpWr,
                                  kMaxQpWr, IBV_QPT_RC, /*sig_all=*/0);
  ASSERT_THAT(qp_pair.recv_qp, NotNull())
      << "Failed to create recv qp - " << errno;
  qp_pair.dst_buffer = setup.dst_memblock.subspan(kMaxQpWr, kMaxQpWr);
  ibv_.SetUpLoopbackRcQps(qp_pair.send_qp, qp_pair.recv_qp, setup.port_gid);
  QueueRecv(setup, qp_pair);
  QueueSend(setup, qp_pair);

  // Wait for completion and verify timestamp.
  ASSERT_OK(verbs_util::WaitForPollingExtendedCompletion(recv_cq));
  ASSERT_EQ(recv_cq->status, IBV_WC_SUCCESS);
  if (support_completion_timestamp_hcaclock_) {
    EXPECT_GT(ibv_wc_read_completion_ts(recv_cq), 0);
  }
  if (support_completion_timestamp_wallclock_) {
    EXPECT_GT(ibv_wc_read_completion_wallclock_ns(recv_cq), 0);
  }
  ibv_end_poll(recv_cq);
}

TEST_F(CqExOpTest, BatchPollSendCq) {
  ASSERT_OK_AND_ASSIGN(BasicSetup setup, CreateBasicSetup());
  static constexpr int kCqSize = 60;
  static constexpr int kQpPairCount = 20;
  static constexpr int kWritesPerQpPair = kCqSize / kQpPairCount;
  ibv_cq_init_attr_ex cq_attr = {
      .cqe = static_cast<uint32_t>(Introspection().device_attr().max_cqe),
      .wc_flags = GetWcFlags()};
  ibv_cq_ex* send_cq = ibv_.CreateCqEx(setup.context, cq_attr);
  ibv_cq_ex* recv_cq = ibv_.CreateCqEx(setup.context, cq_attr);
  std::vector<QpPair> qp_pairs = CreateTestQpPairs(
      setup, ibv_cq_ex_to_cq(send_cq), ibv_cq_ex_to_cq(recv_cq),
      kWritesPerQpPair, kQpPairCount);

  for (auto& qp_pair : qp_pairs) {
    for (int i = 0; i < kWritesPerQpPair; ++i) {
      QueueWrite(setup, qp_pair);
    }
  }
  WaitForAndVerifyCompletions(send_cq, kWritesPerQpPair * kQpPairCount);
}

TEST_F(CqExOpTest, BatchPollRecvCq) {
  ASSERT_OK_AND_ASSIGN(BasicSetup setup, CreateBasicSetup());
  static constexpr int kCqSize = 60;
  static constexpr int kQpPairCount = 20;
  static constexpr int kSendsPerQpPair = kCqSize / kQpPairCount;
  ibv_cq_init_attr_ex cq_attr = {
      .cqe = static_cast<uint32_t>(Introspection().device_attr().max_cqe),
      .wc_flags = GetWcFlags()};
  ibv_cq_ex* send_cq = ibv_.CreateCqEx(setup.context, cq_attr);
  ibv_cq_ex* recv_cq = ibv_.CreateCqEx(setup.context, cq_attr);
  std::vector<QpPair> qp_pairs = CreateTestQpPairs(
      setup, ibv_cq_ex_to_cq(send_cq), ibv_cq_ex_to_cq(recv_cq),
      kSendsPerQpPair, kQpPairCount);

  for (auto& qp_pair : qp_pairs) {
    for (int i = 0; i < kSendsPerQpPair; ++i) {
      QueueRecv(setup, qp_pair);
    }
  }
  for (auto& qp_pair : qp_pairs) {
    for (int i = 0; i < kSendsPerQpPair; ++i) {
      QueueSend(setup, qp_pair);
    }
  }
  WaitForAndVerifyCompletions(recv_cq, kSendsPerQpPair);
}

class CqExOverflowTest : public CqExOpTest {
 protected:
  // Maximum amount of time for waiting for data to land in destination buffer.
  static constexpr absl::Duration kPollTime = absl::Seconds(10);
  // There is a time discrepancy between the data landing on the remote buffer
  // and the completion being generated in the local queue pair. Waiting a
  // small amount of time after WaitForData() returns before polling the
  // completion queue helps to reduce the race.
  static constexpr absl::Duration kCompletionWaitTime = absl::Seconds(1);

  // Wait for send/write data to land on the remote buffer. The function will
  // block until all bytes on |dst_buffer| is of |expected_value|. Notice there
  // is still a time discrepancy between the data landing on the desination
  // buffer and the completion entry being generated (but not pushed to
  // completion queue). Use to detect the completion of an op whose completion
  // entries is dropped due to CQ overflow.
  void WaitForData(absl::Span<uint8_t> dst_buffer, uint8_t expected_value,
                   absl::Duration poll_timeout = kPollTime) {
    auto read_value = [](uint8_t* target) {
      return reinterpret_cast<volatile std::atomic<uint8_t>*>(target)->load();
    };
    absl::Time stop = absl::Now() + poll_timeout;
    bool completed = false;
    for (absl::Time now = absl::Now(); now < stop; now = absl::Now()) {
      completed = true;
      for (size_t i = 0; i < dst_buffer.size(); ++i) {
        if (read_value(dst_buffer.data() + i) != expected_value) {
          completed = false;
        }
      }
      if (completed == false) {
        absl::SleepFor(absl::Milliseconds(10));
      } else {
        return;
      }
    }
    LOG(FATAL) << "Data failed to land.";  // Crash ok;
  }
};

TEST_F(CqExOverflowTest, SendCqOverflow) {
  if (Introspection().FullCqIdlesQp()) {
    GTEST_SKIP() << "This test assumes CQ overflow overwrites completions.";
  }
  ASSERT_OK_AND_ASSIGN(BasicSetup setup, CreateBasicSetup());
  ibv_cq_init_attr_ex cq_attr = {.cqe = verbs_util::kDefaultMaxWr,
                                 .wc_flags = GetWcFlags()};
  ibv_cq_ex* cq = ibv_.CreateCqEx(setup.context, cq_attr);
  const int total_writes = cq->cqe + 10;
  std::vector<QpPair> qp_pairs =
      CreateTestQpPairs(setup, ibv_cq_ex_to_cq(cq), ibv_cq_ex_to_cq(cq),
                        total_writes, /*count=*/1);
  for (int i = 0; i < total_writes; ++i) {
    QueueWrite(setup, qp_pairs[0]);
  }
  WaitForData(setup.dst_memblock.subspan(0, total_writes), kSrcContent);
  absl::SleepFor(kCompletionWaitTime);
  WaitForAndVerifyCompletions(cq, cq->cqe);
  verbs_util::ExpectNoExtendedCompletion(cq, absl::Seconds(1));
}

TEST_F(CqExOverflowTest, RecvCqOverflow) {
  if (Introspection().FullCqIdlesQp()) {
    GTEST_SKIP() << "This test assumes CQ overflow overwrites completions.";
  }
  ASSERT_OK_AND_ASSIGN(BasicSetup setup, CreateBasicSetup());
  ibv_cq_init_attr_ex cq_attr = {.cqe = verbs_util::kDefaultMaxWr,
                                 .wc_flags = GetWcFlags()};
  ibv_cq_ex* send_cq = ibv_.CreateCqEx(setup.context, cq_attr);
  ibv_cq_ex* recv_cq = ibv_.CreateCqEx(setup.context, cq_attr);
  const int total_sends = recv_cq->cqe + 10;
  std::vector<QpPair> qp_pairs =
      CreateTestQpPairs(setup, ibv_cq_ex_to_cq(send_cq),
                        ibv_cq_ex_to_cq(recv_cq), total_sends, /*count=*/1);
  for (int i = 0; i < total_sends; ++i) {
    QueueRecv(setup, qp_pairs[0]);
  }
  for (int i = 0; i < total_sends; ++i) {
    QueueSend(setup, qp_pairs[0]);
  }
  WaitForData(setup.dst_memblock.subspan(0, total_sends), kSrcContent);
  absl::SleepFor(kCompletionWaitTime);
  WaitForAndVerifyCompletions(recv_cq, recv_cq->cqe);
  verbs_util::ExpectNoExtendedCompletion(recv_cq, absl::Seconds(1));
}

}  // namespace rdma_unit_test
