ray/streaming/src/test/queue_tests_base.h

#include "ray/common/test_util.h"
#include "ray/util/filesystem.h"

namespace ray {
namespace streaming {

ray::ObjectID RandomObjectID() { return ObjectID::FromRandom(); }

static void flushall_redis(void) {
  redisContext *context = redisConnect("127.0.0.1", 6379);
  freeReplyObject(redisCommand(context, "FLUSHALL"));
  freeReplyObject(redisCommand(context, "SET NumRedisShards 1"));
  freeReplyObject(redisCommand(context, "LPUSH RedisShards 127.0.0.1:6380"));
  redisFree(context);
}
/// Base class for real-world tests with streaming queue
class StreamingQueueTestBase : public ::testing::TestWithParam<uint64_t> {
 public:
  StreamingQueueTestBase(int num_nodes, int port)
      : gcs_options_("127.0.0.1", 6379, ""), node_manager_port_(port) {
#ifdef _WIN32
    RAY_CHECK(false) << "port system() calls to Windows before running this test";
#endif
    TestSetupUtil::StartUpRedisServers(std::vector<int>{6379, 6380});

    // flush redis first.
    flushall_redis();

    RAY_CHECK(num_nodes >= 0);
    if (num_nodes > 0) {
      raylet_socket_names_.resize(num_nodes);
      raylet_store_socket_names_.resize(num_nodes);
    }

    // start plasma store.
    for (auto &store_socket : raylet_store_socket_names_) {
      store_socket = TestSetupUtil::StartObjectStore();
    }

    // start gcs server
    gcs_server_socket_name_ = TestSetupUtil::StartGcsServer("127.0.0.1");

    // start raylet on each node. Assign each node with different resources so that
    // a task can be scheduled to the desired node.
    for (int i = 0; i < num_nodes; i++) {
      raylet_socket_names_[i] = TestSetupUtil::StartRaylet(
          raylet_store_socket_names_[i], "127.0.0.1", node_manager_port_ + i, "127.0.0.1",
          "\"CPU,4.0,resource" + std::to_string(i) + ",10\"");
    }
  }

  ~StreamingQueueTestBase() {
    STREAMING_LOG(INFO) << "Stop raylet store and actors";
    for (const auto &raylet_socket_name : raylet_socket_names_) {
      TestSetupUtil::StopRaylet(raylet_socket_name);
    }

    for (const auto &store_socket_name : raylet_store_socket_names_) {
      TestSetupUtil::StopObjectStore(store_socket_name);
    }

    TestSetupUtil::StopGcsServer(gcs_server_socket_name_);
    TestSetupUtil::ShutDownRedisServers();
  }

  JobID NextJobId() const {
    static uint32_t job_counter = 1;
    return JobID::FromInt(job_counter++);
  }

  void InitWorker(ActorID &self_actor_id, ActorID &peer_actor_id,
                  const queue::protobuf::StreamingQueueTestRole role,
                  const std::vector<ObjectID> &queue_ids,
                  const std::vector<ObjectID> &rescale_queue_ids, std::string suite_name,
                  std::string test_name, uint64_t param) {
    auto &driver = CoreWorkerProcess::GetCoreWorker();
    std::string forked_serialized_str;
    ObjectID actor_handle_id;
    Status st = driver.SerializeActorHandle(peer_actor_id, &forked_serialized_str,
                                            &actor_handle_id);
    STREAMING_CHECK(st.ok());
    STREAMING_LOG(INFO) << "forked_serialized_str: " << forked_serialized_str;
    TestInitMessage msg(role, self_actor_id, peer_actor_id, forked_serialized_str,
                        queue_ids, rescale_queue_ids, suite_name, test_name, param);

    std::vector<TaskArg> args;
    args.emplace_back(TaskArg::PassByValue(std::make_shared<RayObject>(
        msg.ToBytes(), nullptr, std::vector<ObjectID>(), true)));
    std::unordered_map<std::string, double> resources;
    TaskOptions options{0, resources};
    std::vector<ObjectID> return_ids;
    RayFunction func{ray::Language::PYTHON,
                     ray::FunctionDescriptorBuilder::BuildPython("", "", "init", "")};

    RAY_CHECK_OK(driver.SubmitActorTask(self_actor_id, func, args, options, &return_ids));
  }

  void SubmitTestToActor(ActorID &actor_id, const std::string test) {
    auto &driver = CoreWorkerProcess::GetCoreWorker();
    uint8_t data[8];
    auto buffer = std::make_shared<LocalMemoryBuffer>(data, 8, true);
    std::vector<TaskArg> args;
    args.emplace_back(TaskArg::PassByValue(
        std::make_shared<RayObject>(buffer, nullptr, std::vector<ObjectID>(), true)));
    std::unordered_map<std::string, double> resources;
    TaskOptions options{0, resources};
    std::vector<ObjectID> return_ids;
    RayFunction func{ray::Language::PYTHON, ray::FunctionDescriptorBuilder::BuildPython(
                                                "", test, "execute_test", "")};

    RAY_CHECK_OK(driver.SubmitActorTask(actor_id, func, args, options, &return_ids));
  }

  bool CheckCurTest(ActorID &actor_id, const std::string test_name) {
    auto &driver = CoreWorkerProcess::GetCoreWorker();
    uint8_t data[8];
    auto buffer = std::make_shared<LocalMemoryBuffer>(data, 8, true);
    std::vector<TaskArg> args;
    args.emplace_back(TaskArg::PassByValue(
        std::make_shared<RayObject>(buffer, nullptr, std::vector<ObjectID>(), true)));
    std::unordered_map<std::string, double> resources;
    TaskOptions options{1, resources};
    std::vector<ObjectID> return_ids;
    RayFunction func{ray::Language::PYTHON, ray::FunctionDescriptorBuilder::BuildPython(
                                                "", "", "check_current_test_status", "")};

    RAY_CHECK_OK(driver.SubmitActorTask(actor_id, func, args, options, &return_ids));

    std::vector<bool> wait_results;
    std::vector<std::shared_ptr<RayObject>> results;
    Status wait_st = driver.Wait(return_ids, 1, 5 * 1000, &wait_results);
    if (!wait_st.ok()) {
      STREAMING_LOG(ERROR) << "Wait fail.";
      return false;
    }
    STREAMING_CHECK(wait_results.size() >= 1);
    if (!wait_results[0]) {
      STREAMING_LOG(WARNING) << "Wait direct call fail.";
      return false;
    }

    Status get_st = driver.Get(return_ids, -1, &results);
    if (!get_st.ok()) {
      STREAMING_LOG(ERROR) << "Get fail.";
      return false;
    }
    STREAMING_CHECK(results.size() >= 1);
    if (results[0]->IsException()) {
      STREAMING_LOG(INFO) << "peer actor may has exceptions.";
      return false;
    }
    STREAMING_CHECK(results[0]->HasData());
    STREAMING_LOG(DEBUG) << "SendForResult result[0] DataSize: " << results[0]->GetSize();

    const std::shared_ptr<ray::Buffer> result_buffer = results[0]->GetData();
    std::shared_ptr<LocalMemoryBuffer> return_buffer =
        std::make_shared<LocalMemoryBuffer>(result_buffer->Data(), result_buffer->Size(),
                                            true);

    uint8_t *bytes = result_buffer->Data();
    uint8_t *p_cur = bytes;
    uint32_t *magic_num = (uint32_t *)p_cur;
    STREAMING_CHECK(*magic_num == Message::MagicNum);

    p_cur += sizeof(Message::MagicNum);
    queue::protobuf::StreamingQueueMessageType *type =
        (queue::protobuf::StreamingQueueMessageType *)p_cur;
    STREAMING_CHECK(*type == queue::protobuf::StreamingQueueMessageType::
                                 StreamingQueueTestCheckStatusRspMsgType);
    std::shared_ptr<TestCheckStatusRspMsg> message =
        TestCheckStatusRspMsg::FromBytes(bytes);
    STREAMING_CHECK(message->TestName() == test_name);
    return message->Status();
  }

  ActorID CreateActorHelper(const std::unordered_map<std::string, double> &resources,
                            bool is_direct_call, int64_t max_restarts) {
    std::unique_ptr<ActorHandle> actor_handle;

    // Test creating actor.
    uint8_t array[] = {1, 2, 3};
    auto buffer = std::make_shared<LocalMemoryBuffer>(array, sizeof(array));

    RayFunction func{ray::Language::PYTHON, ray::FunctionDescriptorBuilder::BuildPython(
                                                "", "", "actor creation task", "")};
    std::vector<TaskArg> args;
    args.emplace_back(TaskArg::PassByValue(
        std::make_shared<RayObject>(buffer, nullptr, std::vector<ObjectID>())));

    std::string name = "";
    ActorCreationOptions actor_options{
        max_restarts,
        /*max_task_retries=*/0,
        /*max_concurrency=*/1,  resources, resources,           {},
        /*is_detached=*/false,  name,      /*is_asyncio=*/false};
    // Create an actor.
    ActorID actor_id;
    RAY_CHECK_OK(CoreWorkerProcess::GetCoreWorker().CreateActor(
        func, args, actor_options, /*extension_data*/ "", &actor_id));
    return actor_id;
  }

  void SubmitTest(uint32_t queue_num, std::string suite_name, std::string test_name,
                  uint64_t timeout_ms) {
    std::vector<ray::ObjectID> queue_id_vec;
    std::vector<ray::ObjectID> rescale_queue_id_vec;
    for (uint32_t i = 0; i < queue_num; ++i) {
      ObjectID queue_id = ray::ObjectID::FromRandom();
      queue_id_vec.emplace_back(queue_id);
    }

    // One scale id
    ObjectID rescale_queue_id = ray::ObjectID::FromRandom();
    rescale_queue_id_vec.emplace_back(rescale_queue_id);

    std::vector<uint64_t> channel_seq_id_vec(queue_num, 0);

    for (size_t i = 0; i < queue_id_vec.size(); ++i) {
      STREAMING_LOG(INFO) << " qid hex => " << queue_id_vec[i].Hex();
    }
    for (auto &qid : rescale_queue_id_vec) {
      STREAMING_LOG(INFO) << " rescale qid hex => " << qid.Hex();
    }
    STREAMING_LOG(INFO) << "Sub process: writer.";

    CoreWorkerOptions options = {
        WorkerType::DRIVER,             // worker_type
        Language::PYTHON,               // langauge
        raylet_store_socket_names_[0],  // store_socket
        raylet_socket_names_[0],        // raylet_socket
        NextJobId(),                    // job_id
        gcs_options_,                   // gcs_options
        true,                           // enable_logging
        "",                             // log_dir
        true,                           // install_failure_signal_handler
        "127.0.0.1",                    // node_ip_address
        node_manager_port_,             // node_manager_port
        "127.0.0.1",                    // raylet_ip_address
        "queue_tests",                  // driver_name
        "",                             // stdout_file
        "",                             // stderr_file
        nullptr,                        // task_execution_callback
        nullptr,                        // check_signals
        nullptr,                        // gc_collect
        nullptr,                        // get_lang_stack
        nullptr,                        // kill_main
        true,                           // ref_counting_enabled
        false,                          // is_local_mode
        1,                              // num_workers
    };
    InitShutdownRAII core_worker_raii(CoreWorkerProcess::Initialize,
                                      CoreWorkerProcess::Shutdown, options);

    // Create writer and reader actors
    std::unordered_map<std::string, double> resources;
    auto actor_id_writer = CreateActorHelper(resources, true, 0);
    auto actor_id_reader = CreateActorHelper(resources, true, 0);

    InitWorker(actor_id_writer, actor_id_reader,
               queue::protobuf::StreamingQueueTestRole::WRITER, queue_id_vec,
               rescale_queue_id_vec, suite_name, test_name, GetParam());
    InitWorker(actor_id_reader, actor_id_writer,
               queue::protobuf::StreamingQueueTestRole::READER, queue_id_vec,
               rescale_queue_id_vec, suite_name, test_name, GetParam());

    std::this_thread::sleep_for(std::chrono::milliseconds(2000));

    SubmitTestToActor(actor_id_writer, test_name);
    SubmitTestToActor(actor_id_reader, test_name);

    uint64_t slept_time_ms = 0;
    while (slept_time_ms < timeout_ms) {
      std::this_thread::sleep_for(std::chrono::milliseconds(5 * 1000));
      STREAMING_LOG(INFO) << "Check test status.";
      if (CheckCurTest(actor_id_writer, test_name) &&
          CheckCurTest(actor_id_reader, test_name)) {
        STREAMING_LOG(INFO) << "Test Success, Exit.";
        return;
      }
      slept_time_ms += 5 * 1000;
    }

    EXPECT_TRUE(false);
    STREAMING_LOG(INFO) << "Test Timeout, Exit.";
  }

  void SetUp() {}

  void TearDown() {}

 protected:
  std::vector<std::string> raylet_socket_names_;
  std::vector<std::string> raylet_store_socket_names_;
  gcs::GcsClientOptions gcs_options_;
  int node_manager_port_;
  std::string gcs_server_socket_name_;
};

}  // namespace streaming
}  // namespace ray