Method TestQueryTeardown

be/src/runtime/bufferpool/buffer-pool-test.cc:1549–1601

Test teardown of a query while writes are in flight. This was based on a BufferedBlockMgr regression test for IMPALA-2252 where tear-down of the query's RuntimeStates raced with scratch writes. If write_error is true, force writes to hit errors.

/// query's RuntimeStates raced with scratch writes. If write_error is true,
/// force writes to hit errors.
void BufferPoolTest::TestQueryTeardown(bool write_error) {
  const int64_t TOTAL_BUFFERS = 20;
  const int CLIENT_BUFFERS = 10;
  const int64_t TOTAL_MEM = TOTAL_BUFFERS * TEST_BUFFER_LEN;
  const int64_t CLIENT_MEM = CLIENT_BUFFERS * TEST_BUFFER_LEN;

  // Set up a BufferPool in the TestEnv.
  test_env_.reset(new TestEnv());
  test_env_->SetBufferPoolArgs(TEST_BUFFER_LEN, TOTAL_MEM);
  ASSERT_OK(test_env_->Init());

  BufferPool* pool = test_env_->exec_env()->buffer_pool();
  RuntimeState* state;
  ASSERT_OK(test_env_->CreateQueryState(0, nullptr, &state));

  ClientHandle client;
  ASSERT_OK(pool->RegisterClient("test client", state->query_state()->file_group(),
      state->instance_buffer_reservation(),
      obj_pool_.Add(new MemTracker(-1, "", state->instance_mem_tracker())), CLIENT_MEM,
      NewProfile(), &client));
  ASSERT_TRUE(client.IncreaseReservation(CLIENT_MEM));

  vector<PageHandle> pages;
  CreatePages(pool, &client, TEST_BUFFER_LEN, CLIENT_BUFFERS, &pages);

  if (write_error) {
    UnpinAll(pool, &client, &pages);
    // Allocate more buffers to create memory pressure and force eviction of all the
    // unpinned pages.
    vector<BufferHandle> tmp_buffers;
    AllocateBuffers(pool, &client, TEST_BUFFER_LEN, CLIENT_BUFFERS, &tmp_buffers);
    string tmp_file_path = TmpFilePath(pages.data());
    FreeBuffers(pool, &client, &tmp_buffers);

    ASSERT_OK(PinAll(pool, &client, &pages));
    // Remove temporary file to force future writes to that file to fail.
    DisableBackingFile(tmp_file_path);
  }

  // Unpin will initiate writes. If we triggered a write error earlier, some writes may
  // go down the error path.
  UnpinAll(pool, &client, &pages);

  // Tear down the pages, client, and query in the correct order while writes are in
  // flight.
  DestroyAll(pool, &client, &pages);
  pool->DeregisterClient(&client);
  test_env_->TearDownQueries();

  // All memory should be released from the query.
  EXPECT_EQ(0, test_env_->TotalQueryMemoryConsumption());
  EXPECT_EQ(0, test_env_->exec_env()->buffer_reservation()->GetChildReservations());
}

TEST_F(BufferPoolTest, QueryTeardown) {
  TestQueryTeardown(false);