text: /**
   * Copyright (c) 2018-present, Facebook, Inc.
   * All rights reserved.
   *
   * This source code is licensed under the BSD-style license found in the
   * LICENSE file in the root directory of this source tree.
   */
  
  #include "gloo/test/base_test.h"
  
  #include <algorithm>
  #include <array>
  #include <unordered_set>
  
  #include "gloo/transport/tcp/unbound_buffer.h"
  
  namespace gloo {
  namespace test {
  namespace {
  
  // Test parameterization (transport, context size, buffer size).
  using Param = std::tuple<Transport, int, int>;
  
  // Test fixture.
  class SendRecvTest : public BaseTest,
                       public ::testing::WithParamInterface<Param> {};
  
  TEST_P(SendRecvTest, AllToAll) {
    const auto transport = std::get<0>(GetParam());
    const auto contextSize = std::get<1>(GetParam());
  
    spawn(transport, contextSize, [&](std::shared_ptr<Context> context) {
      using buffer_ptr = std::unique_ptr<::gloo::transport::UnboundBuffer>;
      std::vector<int> input(contextSize);
      std::vector<int> output(contextSize);
      std::vector<buffer_ptr> inputBuffers(contextSize);
      std::vector<buffer_ptr> outputBuffers(contextSize);
  
      // Initialize
      for (auto i = 0; i < context->size; i++) {
        input[i] = context->rank;
        output[i] = -1;
        inputBuffers[i] =
            context->createUnboundBuffer(&input[i], sizeof(input[i]));
        outputBuffers[i] =
            context->createUnboundBuffer(&output[i], sizeof(output[i]));
      }
  
      // Send a message with the local rank to every other rank
      for (auto i = 0; i < context->size; i++) {
        if (i == context->rank) {
          continue;
        }
        inputBuffers[i]->send(i, context->rank);
      }
  
      // Receive message from every other rank
      for (auto i = 0; i < context->size; i++) {
        if (i == context->rank) {
          continue;
        }
        outputBuffers[i]->recv(i, i);
      }
  
      // Wait for send and recv to complete
      for (auto i = 0; i < context->size; i++) {
        if (i == context->rank) {
          continue;
        }
        inputBuffers[i]->waitSend();
        outputBuffers[i]->waitRecv();
      }
  
      // Verify output
      for (auto i = 0; i < context->size; i++) {
        if (i == context->rank) {
          continue;
        }
        ASSERT_EQ(i, output[i]) << "Mismatch at index " << i;
      }
    });
  }
  
  TEST_P(SendRecvTest, AllToAllOffset) {
    const auto transport = std::get<0>(GetParam());
    const auto contextSize = std::get<1>(GetParam());
  
    spawn(transport, contextSize, [&](std::shared_ptr<Context> context) {
      auto elementSize = sizeof(int);
      std::vector<int> input(contextSize);
      std::vector<int> output(contextSize);
      auto inputBuffer =
          context->createUnboundBuffer(input.data(), input.size() * elementSize);
      auto outputBuffer = context->createUnboundBuffer(
          output.data(), output.size() * elementSize);
  
      // Initialize
      for (auto i = 0; i < context->size; i++) {
        input[i] = i;
        output[i] = -1;
      }
  
      // Send a message with the local rank to every other rank
      for (auto i = 0; i < context->size; i++) {
        if (i == context->rank) {
          continue;
        }
        inputBuffer->send(i, 0, context->rank * elementSize, elementSize);
      }
  
      // Receive message from every other rank
      for (auto i = 0; i < context->size; i++) {
        if (i == context->rank) {
          continue;
        }
        outputBuffer->recv(i, 0, i * elementSize, elementSize);
      }
  
      // Wait for send and recv to complete
      for (auto i = 0; i < context->size; i++) {
        if (i == context->rank) {
          continue;
        }
        inputBuffer->waitSend();
        outputBuffer->waitRecv();
      }
  
      // Verify output
      for (auto i = 0; i < context->size; i++) {
        if (i == context->rank) {
          continue;
        }
        ASSERT_EQ(i, output[i]) << "Mismatch at index " << i;
      }
    });
  }
  
  TEST_P(SendRecvTest, AllToAllEmptyThenNonEmpty) {
    const auto transport = std::get<0>(GetParam());
    const auto contextSize = std::get<1>(GetParam());
  
    spawn(transport, contextSize, [&](std::shared_ptr<Context> context) {
      auto elementSize = sizeof(int);
      std::vector<int> input(contextSize);
      std::vector<int> output(contextSize);
      using buffer_ptr = std::unique_ptr<::gloo::transport::UnboundBuffer>;
      std::vector<buffer_ptr> emptyInputBuffers(contextSize);
      std::vector<buffer_ptr> emptyOutputBuffers(contextSize);
      std::vector<buffer_ptr> nonEmptyInputBuffers(contextSize);
      std::vector<buffer_ptr> nonEmptyOutputBuffers(contextSize);
  
      // Initialize
      for (auto i = 0; i < context->size; i++) {
        input[i] = context->rank;
        output[i] = -1;
        emptyInputBuffers[i] = context->createUnboundBuffer(nullptr, 0);
        emptyOutputBuffers[i] = context->createUnboundBuffer(nullptr, 0);
        nonEmptyInputBuffers[i] =
            context->createUnboundBuffer(&input[i], elementSize);
        nonEmptyOutputBuffers[i] =
            context->createUnboundBuffer(&output[i], elementSize);
      }
  
      // Kick off all sends and receives
      for (auto i = 0; i < context->size; i++) {
        if (i == context->rank) {
          continue;
        }
        emptyOutputBuffers[i]->recv(i, i);
        emptyInputBuffers[i]->send(i, context->rank);
        nonEmptyOutputBuffers[i]->recv(i, i);
        nonEmptyInputBuffers[i]->send(i, context->rank);
      }
  
      // Wait for send and recv to complete
      for (auto i = 0; i < context->size; i++) {
        if (i == context->rank) {
          continue;
        }
        emptyInputBuffers[i]->waitSend();
        emptyOutputBuffers[i]->waitRecv();
        nonEmptyInputBuffers[i]->waitSend();
        nonEmptyOutputBuffers[i]->waitRecv();
      }
  
      // Verify output
      for (auto i = 0; i < context->size; i++) {
        if (i == context->rank) {
          continue;
        }
        ASSERT_EQ(i, output[i]) << "Mismatch at index " << i;
      }
    });
  }
  
  TEST_P(SendRecvTest, RecvFromAny) {
    const auto transport = std::get<0>(GetParam());
    const auto contextSize = std::get<1>(GetParam());
  
    spawn(transport, contextSize, [&](std::shared_ptr<Context> context) {
      constexpr uint64_t slot = 0x1337;
      if (context->rank == 0) {
        std::unordered_set<int> outputData;
        std::unordered_set<int> outputRanks;
        int tmp;
        auto buf = context->createUnboundBuffer(&tmp, sizeof(tmp));
  
        // Compile vector of ranks to receive from
        std::vector<int> ranks;
        for (auto i = 1; i < context->size; i++) {
          ranks.push_back(i);
        }
  
        // Receive from N-1 peers
        for (auto i = 1; i < context->size; i++) {
          int srcRank = -1;
          buf->recv(ranks, slot);
          buf->waitRecv(&srcRank);
          outputData.insert(tmp);
          outputRanks.insert(srcRank);
        }
  
        // Verify result
        for (auto i = 1; i < context->size; i++) {
          ASSERT_EQ(1, outputData.count(i)) << "Missing output " << i;
          ASSERT_EQ(1, outputRanks.count(i)) << "Missing rank " << i;
        }
      } else {
        // Send to rank 0
        int tmp = context->rank;
        auto buf = context->createUnboundBuffer(&tmp, sizeof(tmp));
        buf->send(0, slot);
        buf->waitSend();
      }
    });
  }
  
  TEST_P(SendRecvTest, RecvFromAnyOffset) {
    const auto transport = std::get<0>(GetParam());
    const auto contextSize = std::get<1>(GetParam());
  
    spawn(transport, contextSize, [&](std::shared_ptr<Context> context) {
      auto elementSize = sizeof(int);
      constexpr uint64_t slot = 0x1337;
      if (context->rank == 0) {
        std::unordered_set<int> outputData;
        std::unordered_set<int> outputRanks;
        std::array<int, 2> tmp;
        auto buf =
            context->createUnboundBuffer(tmp.data(), tmp.size() * elementSize);
  
        // Compile vector of ranks to receive from
        std::vector<int> ranks;
        for (auto i = 1; i < context->size; i++) {
          ranks.push_back(i);
        }
  
        // Receive from N-1 peers
        for (auto i = 1; i < context->size; i++) {
          int srcRank = -1;
          buf->recv(ranks, slot, (i % tmp.size()) * elementSize, elementSize);
          buf->waitRecv(&srcRank);
          outputData.insert(tmp[i % tmp.size()]);
          outputRanks.insert(srcRank);
        }
  
        // Verify result
        for (auto i = 1; i < context->size; i++) {
          ASSERT_EQ(1, outputData.count(i)) << "Missing output " << i;
          ASSERT_EQ(1, outputRanks.count(i)) << "Missing rank " << i;
        }
      } else {
        // Send to rank 0
        int tmp = context->rank;
        auto buf = context->createUnboundBuffer(&tmp, sizeof(tmp));
        buf->send(0, slot);
        buf->waitSend();
      }
    });
  }
  
  TEST_P(SendRecvTest, RecvFromAnyPipeline) {
    const auto transport = std::get<0>(GetParam());
    const auto contextSize = std::get<1>(GetParam());
  
    spawn(transport, contextSize, [&](std::shared_ptr<Context> context) {
      constexpr uint64_t slot = 0x1337;
  
      if (context->rank == 0) {
        std::vector<int> output;
        std::array<int, 2> tmp;
        auto buf0 = context->createUnboundBuffer(&tmp[0], sizeof(tmp[0]));
        auto buf1 = context->createUnboundBuffer(&tmp[1], sizeof(tmp[1]));
  
        // Compile vector of ranks to receive from
        std::vector<int> ranks;
        for (auto i = 0; i < context->size; i++) {
          if (i == context->rank) {
            continue;
          }
          ranks.push_back(i);
        }
  
        // Receive twice per peer
        for (auto i = 0; i < context->size - 1; i++) {
          buf0->recv(ranks, slot);
          buf1->recv(ranks, slot);
          buf0->waitRecv();
          buf1->waitRecv();
          output.push_back(tmp[0]);
          output.push_back(tmp[1]);
        }
  
        // Verify output
        std::sort(output.begin(), output.end());
        for (auto i = 1; i < context->size; i++) {
          ASSERT_EQ(i, output[(i - 1) * 2 + 0]) << "Mismatch at " << i;
          ASSERT_EQ(i, output[(i - 1) * 2 + 1]) << "Mismatch at " << i;
        }
      } else {
        // Send twice to rank 0 on the same slot
        std::array<int, 2> tmp;
        tmp[0] = context->rank;
        tmp[1] = context->rank;
        auto buf0 = context->createUnboundBuffer(&tmp[0], sizeof(tmp[0]));
        auto buf1 = context->createUnboundBuffer(&tmp[1], sizeof(tmp[1]));
        buf0->send(0, slot);
        buf1->send(0, slot);
        buf0->waitSend();
        buf1->waitSend();
      }
    });
  }
  
  // This test stresses the pattern that is commonly used for some sort
  // of RPC on top of send/recv primitives. The server executes an
  // indirect recv (that can be fulfilled any some subset of peers) to
  // receive the message size and rank of a client, followed by an
  // immediate recv from that rank to get the actual message.
  TEST_P(SendRecvTest, RecvFromAnyRPC) {
    const auto transport = std::get<0>(GetParam());
    const auto contextSize = std::get<1>(GetParam());
  
    spawn(transport, contextSize, [&](std::shared_ptr<Context> context) {
      constexpr uint64_t slot = 0x1337;
      constexpr auto niters = 10;
      size_t tmp0;
      size_t tmp1;
      auto buf0 = context->createUnboundBuffer(&tmp0, sizeof(tmp0));
      auto buf1 = context->createUnboundBuffer(&tmp1, sizeof(tmp1));
  
      if (context->rank == 0) {
        // Keep number of recvs per rank
        std::unordered_map<int, int> counts;
        // Compile vector of ranks to receive from
        std::vector<int> allRanks;
        for (auto i = 0; i < context->size; i++) {
          if (i == context->rank) {
            continue;
          }
          allRanks.push_back(i);
        }
        // Receive twice from every peer, niters times
        for (auto i = 0; i < (niters * (context->size - 1)); i++) {
          int rank;
          // Receive from any peer first
          buf0->recv(allRanks, slot);
          buf0->waitRecv(&rank);
          counts[rank]++;
          // Then receive from the same peer again
          buf1->recv(rank, slot);
          buf1->waitRecv();
        }
        // Verify result
        for (auto i = 0; i < context->size; i++) {
          if (i == context->rank) {
            continue;
          }
          ASSERT_EQ(niters, counts[i]) << "recv mismatch for rank " << i;
        }
      } else {
        for (auto i = 0; i < niters; i++) {
          buf0->send(0, slot);
          buf0->waitSend();
          buf1->send(0, slot);
          buf1->waitSend();
        }
      }
    });
  }
  
  TEST_P(SendRecvTest, RecvFromAnyRPCEmptyThenNonEmpty) {
    const auto transport = std::get<0>(GetParam());
    const auto contextSize = std::get<1>(GetParam());
  
    spawn(transport, contextSize, [&](std::shared_ptr<Context> context) {
      constexpr uint64_t slot = 0x1337;
      constexpr auto niters = 10;
      size_t tmp0;
      size_t tmp1;
      auto buf0 = context->createUnboundBuffer(&tmp0, sizeof(tmp0));
      auto buf1 = context->createUnboundBuffer(&tmp1, sizeof(tmp1));
  
      if (context->rank == 0) {
        // Compile vector of ranks to receive from
        std::vector<int> allRanks;
        for (auto i = 0; i < context->size; i++) {
          if (i == context->rank) {
            continue;
          }
          allRanks.push_back(i);
        }
        // Receive twice from every peer, niters times
        for (auto i = 0; i < (niters * (context->size - 1)); i++) {
          int rank;
          // Receive from any peer first (0 bytes)
          tmp0 = 0xdeadbeef;
          buf0->recv(allRanks, slot, /* offset= */ 0, /* nbytes= */ 0);
          buf0->waitRecv(&rank);
          ASSERT_EQ(tmp0, 0xdeadbeef);
          // Then receive from the same peer again
          tmp1 = 0xdeadbeef;
          buf1->recv(rank, slot);
          buf1->waitRecv();
          GLOO_ENFORCE_EQ(tmp1, rank);
        }
      } else {
        for (auto i = 0; i < niters; i++) {
          tmp0 = context->rank;
          buf0->send(0, slot, /* offset= */ 0, /* nbytes= */ 0);
          buf0->waitSend();
          tmp1 = context->rank;
          buf1->send(0, slot);
          buf1->waitSend();
        }
      }
    });
  }
  
  INSTANTIATE_TEST_CASE_P(
      SendRecvDefault,
      SendRecvTest,
      ::testing::Combine(
          ::testing::Values(Transport::TCP, Transport::UV),
          ::testing::Values(2, 3, 4, 5, 6, 7, 8),
          ::testing::Values(1)));
  
  } // namespace
  } // namespace test
  } // namespace gloo
  