Skip to content

Commit

Permalink
Improve the sanity checks in (Multi)GetEntity and friends (#12630)
Browse files Browse the repository at this point in the history 
Summary:
Pull Request resolved: #12630

The patch cleans up, improves, and brings into sync (to the extent possible without API signature changes) the sanity checks around the `GetEntity` / `MultiGetEntity` family of APIs, including the read-your-own-writes (`WriteBatchWithIndex`) and transaction layers. The checks are centralized in two main sets of entry points, namely in `DB(Impl)` and the "main" `GetEntityFromBatchAndDB` / `MultiGetEntityFromBatchAndDB` overloads in `WriteBatchWithIndex`. This eliminates the need to duplicate the checks in the transaction classes.

Reviewed By: jaykorean

Differential Revision: D57125741

fbshipit-source-id: 4dd059ef644a9b173fbba767538943397e4cc6cd
  • Loading branch information
@ltamasi @facebook-github-bot
ltamasi authored and facebook-github-bot committed May 9, 2024
1 parent 1a33576 commit 97e7090
Show file tree
Hide file tree
Showing 9 changed files with 542 additions and 114 deletions.
119 changes: 105 additions & 14 deletions db/db_impl/db_impl.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2099,7 +2099,7 @@ Status DBImpl::GetEntity(const ReadOptions& _read_options,
if (_read_options.io_activity != Env::IOActivity::kUnknown &&
_read_options.io_activity != Env::IOActivity::kGetEntity) {
return Status::InvalidArgument(
"Cannot call GetEntity with `ReadOptions::io_activity` != "
"Can only call GetEntity with `ReadOptions::io_activity` set to "
"`Env::IOActivity::kUnknown` or `Env::IOActivity::kGetEntity`");
}
ReadOptions read_options(_read_options);
Expand All @@ -2126,7 +2126,7 @@ Status DBImpl::GetEntity(const ReadOptions& _read_options, const Slice& key,
if (_read_options.io_activity != Env::IOActivity::kUnknown &&
_read_options.io_activity != Env::IOActivity::kGetEntity) {
s = Status::InvalidArgument(
"Cannot call GetEntity with `ReadOptions::io_activity` != "
"Can only call GetEntity with `ReadOptions::io_activity` set to "
"`Env::IOActivity::kUnknown` or `Env::IOActivity::kGetEntity`");
for (size_t i = 0; i < num_column_families; ++i) {
(*result)[i].SetStatus(s);
Expand Down Expand Up @@ -3185,22 +3185,55 @@ void DBImpl::MultiGetEntity(const ReadOptions& _read_options, size_t num_keys,
ColumnFamilyHandle** column_families,
const Slice* keys, PinnableWideColumns* results,
Status* statuses, bool sorted_input) {
assert(statuses);

if (!column_families) {
const Status s = Status::InvalidArgument(
"Cannot call MultiGetEntity without column families");
for (size_t i = 0; i < num_keys; ++i) {
statuses[i] = s;
}

return;
}

if (!keys) {
const Status s =
Status::InvalidArgument("Cannot call MultiGetEntity without keys");
for (size_t i = 0; i < num_keys; ++i) {
statuses[i] = s;
}

return;
}

if (!results) {
const Status s = Status::InvalidArgument(
"Cannot call MultiGetEntity without PinnableWideColumns objects");
for (size_t i = 0; i < num_keys; ++i) {
statuses[i] = s;
}

return;
}

if (_read_options.io_activity != Env::IOActivity::kUnknown &&
_read_options.io_activity != Env::IOActivity::kMultiGetEntity) {
Status s = Status::InvalidArgument(
"Can only call MultiGetEntity with `ReadOptions::io_activity` is "
const Status s = Status::InvalidArgument(
"Can only call MultiGetEntity with `ReadOptions::io_activity` set to "
"`Env::IOActivity::kUnknown` or `Env::IOActivity::kMultiGetEntity`");
for (size_t i = 0; i < num_keys; ++i) {
if (statuses[i].ok()) {
statuses[i] = s;
}
statuses[i] = s;
}

return;
}

ReadOptions read_options(_read_options);
if (read_options.io_activity == Env::IOActivity::kUnknown) {
read_options.io_activity = Env::IOActivity::kMultiGetEntity;
}

MultiGetCommon(read_options, num_keys, column_families, keys,
/* values */ nullptr, results, /* timestamps */ nullptr,
statuses, sorted_input);
Expand All @@ -3210,22 +3243,54 @@ void DBImpl::MultiGetEntity(const ReadOptions& _read_options,
ColumnFamilyHandle* column_family, size_t num_keys,
const Slice* keys, PinnableWideColumns* results,
Status* statuses, bool sorted_input) {
assert(statuses);

if (!column_family) {
const Status s = Status::InvalidArgument(
"Cannot call MultiGetEntity without a column family handle");
for (size_t i = 0; i < num_keys; ++i) {
statuses[i] = s;
}

return;
}

if (!keys) {
const Status s =
Status::InvalidArgument("Cannot call MultiGetEntity without keys");
for (size_t i = 0; i < num_keys; ++i) {
statuses[i] = s;
}

return;
}

if (!results) {
const Status s = Status::InvalidArgument(
"Cannot call MultiGetEntity without PinnableWideColumns objects");
for (size_t i = 0; i < num_keys; ++i) {
statuses[i] = s;
}

return;
}

if (_read_options.io_activity != Env::IOActivity::kUnknown &&
_read_options.io_activity != Env::IOActivity::kMultiGetEntity) {
Status s = Status::InvalidArgument(
"Can only call MultiGetEntity with `ReadOptions::io_activity` is "
const Status s = Status::InvalidArgument(
"Can only call MultiGetEntity with `ReadOptions::io_activity` set to "
"`Env::IOActivity::kUnknown` or `Env::IOActivity::kMultiGetEntity`");
for (size_t i = 0; i < num_keys; ++i) {
if (statuses[i].ok()) {
statuses[i] = s;
}
statuses[i] = s;
}
return;
}

ReadOptions read_options(_read_options);
if (read_options.io_activity == Env::IOActivity::kUnknown) {
read_options.io_activity = Env::IOActivity::kMultiGetEntity;
}

MultiGetCommon(read_options, column_family, num_keys, keys,
/* values */ nullptr, results, /* timestamps */ nullptr,
statuses, sorted_input);
Expand All @@ -3234,18 +3299,34 @@ void DBImpl::MultiGetEntity(const ReadOptions& _read_options,
void DBImpl::MultiGetEntity(const ReadOptions& _read_options, size_t num_keys,
const Slice* keys,
PinnableAttributeGroups* results) {
assert(results);

if (!keys) {
const Status s =
Status::InvalidArgument("Cannot call MultiGetEntity without keys");
for (size_t i = 0; i < num_keys; ++i) {
for (size_t j = 0; j < results[i].size(); ++j) {
results[i][j].SetStatus(s);
}
}

return;
}

if (_read_options.io_activity != Env::IOActivity::kUnknown &&
_read_options.io_activity != Env::IOActivity::kMultiGetEntity) {
Status s = Status::InvalidArgument(
"Can only call MultiGetEntity with ReadOptions::io_activity` is "
const Status s = Status::InvalidArgument(
"Can only call MultiGetEntity with `ReadOptions::io_activity` set to "
"`Env::IOActivity::kUnknown` or `Env::IOActivity::kMultiGetEntity`");
for (size_t i = 0; i < num_keys; ++i) {
for (size_t j = 0; j < results[i].size(); ++j) {
results[i][j].SetStatus(s);
}
}

return;
}

ReadOptions read_options(_read_options);
if (read_options.io_activity == Env::IOActivity::kUnknown) {
read_options.io_activity = Env::IOActivity::kMultiGetEntity;
Expand All @@ -3263,6 +3344,7 @@ void DBImpl::MultiGetEntity(const ReadOptions& _read_options, size_t num_keys,
++total_count;
}
}

std::vector<Status> statuses(total_count);
std::vector<PinnableWideColumns> columns(total_count);
MultiGetCommon(read_options, total_count, column_families.data(),
Expand All @@ -3283,6 +3365,15 @@ void DBImpl::MultiGetEntity(const ReadOptions& _read_options, size_t num_keys,
}
}

void DBImpl::MultiGetEntityWithCallback(
const ReadOptions& read_options, ColumnFamilyHandle* column_family,
ReadCallback* callback,
autovector<KeyContext*, MultiGetContext::MAX_BATCH_SIZE>* sorted_keys) {
assert(read_options.io_activity == Env::IOActivity::kMultiGetEntity);

MultiGetWithCallbackImpl(read_options, column_family, callback, sorted_keys);
}

Status DBImpl::WrapUpCreateColumnFamilies(
const ReadOptions& read_options, const WriteOptions& write_options,
const std::vector<const ColumnFamilyOptions*>& cf_options) {
Expand Down
5 changes: 5 additions & 0 deletions db/db_impl/db_impl.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -291,6 +291,11 @@ class DBImpl : public DB {
const Slice* keys,
PinnableAttributeGroups* results) override;

void MultiGetEntityWithCallback(
const ReadOptions& read_options, ColumnFamilyHandle* column_family,
ReadCallback* callback,
autovector<KeyContext*, MultiGetContext::MAX_BATCH_SIZE>* sorted_keys);

Status CreateColumnFamily(const ColumnFamilyOptions& cf_options,
const std::string& column_family,
ColumnFamilyHandle** handle) override {
Expand Down
138 changes: 138 additions & 0 deletions db/wide/db_wide_basic_test.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1794,6 +1794,144 @@ TEST_F(DBWideBasicTest, PinnableWideColumnsMove) {
test_move(/* fill_cache*/ true);
}

TEST_F(DBWideBasicTest, SanityChecks) {
constexpr char foo[] = "foo";
constexpr char bar[] = "bar";
constexpr size_t num_keys = 2;

{
constexpr ColumnFamilyHandle* column_family = nullptr;
PinnableWideColumns columns;
ASSERT_TRUE(db_->GetEntity(ReadOptions(), column_family, foo, &columns)
.IsInvalidArgument());
}

{
constexpr PinnableWideColumns* columns = nullptr;
ASSERT_TRUE(
db_->GetEntity(ReadOptions(), db_->DefaultColumnFamily(), foo, columns)
.IsInvalidArgument());
}

{
ReadOptions read_options;
read_options.io_activity = Env::IOActivity::kGet;

PinnableWideColumns columns;
ASSERT_TRUE(
db_->GetEntity(read_options, db_->DefaultColumnFamily(), foo, &columns)
.IsInvalidArgument());
}

{
constexpr ColumnFamilyHandle* column_family = nullptr;
std::array<Slice, num_keys> keys{{foo, bar}};
std::array<PinnableWideColumns, num_keys> results;
std::array<Status, num_keys> statuses;

db_->MultiGetEntity(ReadOptions(), column_family, num_keys, keys.data(),
results.data(), statuses.data());

ASSERT_TRUE(statuses[0].IsInvalidArgument());
ASSERT_TRUE(statuses[1].IsInvalidArgument());
}

{
constexpr Slice* keys = nullptr;
std::array<PinnableWideColumns, num_keys> results;
std::array<Status, num_keys> statuses;

db_->MultiGetEntity(ReadOptions(), db_->DefaultColumnFamily(), num_keys,
keys, results.data(), statuses.data());

ASSERT_TRUE(statuses[0].IsInvalidArgument());
ASSERT_TRUE(statuses[1].IsInvalidArgument());
}

{
std::array<Slice, num_keys> keys{{foo, bar}};
constexpr PinnableWideColumns* results = nullptr;
std::array<Status, num_keys> statuses;

db_->MultiGetEntity(ReadOptions(), db_->DefaultColumnFamily(), num_keys,
keys.data(), results, statuses.data());

ASSERT_TRUE(statuses[0].IsInvalidArgument());
ASSERT_TRUE(statuses[1].IsInvalidArgument());
}

{
ReadOptions read_options;
read_options.io_activity = Env::IOActivity::kMultiGet;

std::array<Slice, num_keys> keys{{foo, bar}};
std::array<PinnableWideColumns, num_keys> results;
std::array<Status, num_keys> statuses;

db_->MultiGetEntity(read_options, db_->DefaultColumnFamily(), num_keys,
keys.data(), results.data(), statuses.data());
ASSERT_TRUE(statuses[0].IsInvalidArgument());
ASSERT_TRUE(statuses[1].IsInvalidArgument());
}

{
constexpr ColumnFamilyHandle** column_families = nullptr;
std::array<Slice, num_keys> keys{{foo, bar}};
std::array<PinnableWideColumns, num_keys> results;
std::array<Status, num_keys> statuses;

db_->MultiGetEntity(ReadOptions(), num_keys, column_families, keys.data(),
results.data(), statuses.data());

ASSERT_TRUE(statuses[0].IsInvalidArgument());
ASSERT_TRUE(statuses[1].IsInvalidArgument());
}

{
std::array<ColumnFamilyHandle*, num_keys> column_families{
{db_->DefaultColumnFamily(), db_->DefaultColumnFamily()}};
constexpr Slice* keys = nullptr;
std::array<PinnableWideColumns, num_keys> results;
std::array<Status, num_keys> statuses;

db_->MultiGetEntity(ReadOptions(), num_keys, column_families.data(), keys,
results.data(), statuses.data());

ASSERT_TRUE(statuses[0].IsInvalidArgument());
ASSERT_TRUE(statuses[1].IsInvalidArgument());
}

{
std::array<ColumnFamilyHandle*, num_keys> column_families{
{db_->DefaultColumnFamily(), db_->DefaultColumnFamily()}};
std::array<Slice, num_keys> keys{{foo, bar}};
constexpr PinnableWideColumns* results = nullptr;
std::array<Status, num_keys> statuses;

db_->MultiGetEntity(ReadOptions(), num_keys, column_families.data(),
keys.data(), results, statuses.data());

ASSERT_TRUE(statuses[0].IsInvalidArgument());
ASSERT_TRUE(statuses[1].IsInvalidArgument());
}

{
ReadOptions read_options;
read_options.io_activity = Env::IOActivity::kMultiGet;

std::array<ColumnFamilyHandle*, num_keys> column_families{
{db_->DefaultColumnFamily(), db_->DefaultColumnFamily()}};
std::array<Slice, num_keys> keys{{foo, bar}};
std::array<PinnableWideColumns, num_keys> results;
std::array<Status, num_keys> statuses;

db_->MultiGetEntity(read_options, num_keys, column_families.data(),
keys.data(), results.data(), statuses.data());
ASSERT_TRUE(statuses[0].IsInvalidArgument());
ASSERT_TRUE(statuses[1].IsInvalidArgument());
}
}

} // namespace ROCKSDB_NAMESPACE

int main(int argc, char** argv) {
Expand Down
15 changes: 13 additions & 2 deletions include/rocksdb/utilities/write_batch_with_index.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -284,7 +284,12 @@ class WriteBatchWithIndex : public WriteBatchBase {
Status GetEntityFromBatchAndDB(DB* db, const ReadOptions& read_options,
ColumnFamilyHandle* column_family,
const Slice& key,
PinnableWideColumns* columns);
PinnableWideColumns* columns) {
constexpr ReadCallback* callback = nullptr;

return GetEntityFromBatchAndDB(db, read_options, column_family, key,
columns, callback);
}

void MultiGetFromBatchAndDB(DB* db, const ReadOptions& read_options,
ColumnFamilyHandle* column_family,
Expand All @@ -310,7 +315,13 @@ class WriteBatchWithIndex : public WriteBatchBase {
ColumnFamilyHandle* column_family,
size_t num_keys, const Slice* keys,
PinnableWideColumns* results,
Status* statuses, bool sorted_input);
Status* statuses, bool sorted_input) {
constexpr ReadCallback* callback = nullptr;

MultiGetEntityFromBatchAndDB(db, read_options, column_family, num_keys,
keys, results, statuses, sorted_input,
callback);
}

// Records the state of the batch for future calls to RollbackToSavePoint().
// May be called multiple times to set multiple save points.
Expand Down
Loading

0 comments on commit 97e7090

Please sign in to comment.