Repository: ray-project/ray
File: src/ray/raylet/tests/local_object_manager_test.cc (C++)

Code before:
    urls.push_back(BuildURL("url", i));
  }
  EXPECT_CALL(worker_pool, PushSpillWorker(_)).Times(1);
  ASSERT_TRUE(worker_pool.io_worker_client->ReplySpillObjects(urls));
  ASSERT_TRUE(owner_client->ReplyUpdateObjectLocationBatch());

  // Spilled all objects
  ASSERT_EQ(GetCurrentSpilledCount(), 40);
  ASSERT_EQ(GetCurrentSpilledBytes(), object_size * 40);
}

TEST_F(LocalObjectManagerMaxFileSizeFusedTest, TestMaxSpillingFileSizeMaxFusionCount) {
  rpc::Address owner_address;
  owner_address.set_worker_id(WorkerID::FromRandom().Binary());

  std::vector<ObjectID> object_ids;
  std::vector<std::unique_ptr<RayObject>> objects;

  // 40 objects * 10 bytes = 400 bytes.
  // max_spilling_file_size=60 -> batches: 6,6,6,6,6,6,4 (total 7 batches).
  const int64_t object_size = 10;
  for (int i = 0; i < 40; i++) {
    ObjectID object_id = ObjectID::FromRandom();
    object_ids.push_back(object_id);
    auto data_buffer = std::make_shared<MockObjectBuffer>(object_size, object_id, unpins);
    auto object = std::make_unique<RayObject>(
        data_buffer, nullptr, std::vector<rpc::ObjectReference>());
    objects.push_back(std::move(object));
  }
  manager.PinObjectsAndWaitForFree(object_ids, std::move(objects), owner_address);

  // Spill all objects and verify the spilled batches.
  const std::vector<int> expected_batches = {6, 6, 6, 6, 6, 6, 4};

  // Spill may take multiple rounds to complete, so we need to loop until all
  // batches are spilled.
  size_t batch_idx = 0;
  while (batch_idx < expected_batches.size()) {
    manager.SpillObjectUptoMaxThroughput();

    // This round may enqueue 0/1/2 pop callbacks (depending on max_io_workers=2).
    while (worker_pool.FlushPopSpillWorkerCallbacks()) {
    }

    // Reply to all in-flight spill RPCs. Each reply consumes one expected batch.
    while (!worker_pool.io_worker_client->callbacks.empty() &&
           batch_idx < expected_batches.size()) {
      const int n = expected_batches[batch_idx];

      std::vector<std::string> urls;
      urls.reserve(n);

Review comment (refactor):
```suggestion
    std::unique_ptr<RayObject> object = std::make_unique<RayObject>(
```
`data_buffer` also needs to be moved if we refactor it to a unique_ptr.

Code after:
  for (int i = 15; i < 25; i++) {
    urls.push_back(BuildURL("url", i));
  }
  EXPECT_CALL(worker_pool, PushSpillWorker(_)).Times(1);
  ASSERT_TRUE(worker_pool.io_worker_client->ReplySpillObjects(urls));
  ASSERT_TRUE(owner_client->ReplyUpdateObjectLocationBatch());

  // Spilled all objects
  ASSERT_EQ(GetCurrentSpilledCount(), 40);
  ASSERT_EQ(GetCurrentSpilledBytes(), object_size * 40);
}

TEST_F(LocalObjectManagerMaxFileSizeFusedTest, TestMaxSpillingFileSizeMaxFusionCount) {
  rpc::Address owner_address;
  owner_address.set_worker_id(WorkerID::FromRandom().Binary());

  std::vector<ObjectID> object_ids;
  std::vector<std::unique_ptr<RayObject>> objects;

  // 40 objects * 10 bytes = 400 bytes.
  // max_spilling_file_size_bytes=60 -> batches: 6,6,6,6,6,6,4 (total 7 batches).
  const int64_t object_size = 10;
  for (int i = 0; i < 40; i++) {
    ObjectID object_id = ObjectID::FromRandom();
    object_ids.push_back(object_id);
    std::shared_ptr<Buffer> data_buffer =
        std::make_shared<MockObjectBuffer>(object_size, object_id, unpins);
    std::unique_ptr<RayObject> object = std::make_unique<RayObject>(
        data_buffer, nullptr, std::vector<rpc::ObjectReference>());
    objects.push_back(std::move(object));
  }
  manager.PinObjectsAndWaitForFree(object_ids, std::move(objects), owner_address);

  // Spill all objects and verify the spilled batches.
  const std::vector<int> expected_batches = {6, 6, 6, 6, 6, 6, 4};

  // Spill may take multiple rounds to complete, so we need to loop until all
  // batches are spilled.
  size_t batch_idx = 0;
  while (batch_idx < expected_batches.size()) {
    manager.SpillObjectUptoMaxThroughput();

    // This round may enqueue 0/1/2 pop callbacks (depending on max_io_workers=2).
    while (worker_pool.FlushPopSpillWorkerCallbacks()) {
    }

    // Reply to all in-flight spill RPCs. Each reply consumes one expected batch.
    while (!worker_pool.io_worker_client->callbacks.empty() &&
           batch_idx < expected_batches.size()) {
      const int n = expected_batches[batch_idx];