ColumnFamilyHandle Nullcheck in GetEntity and MultiGetEntity (#12057)

Summary:
- Add missing null check for ColumnFamilyHandle in `GetEntity()`
- `FailIfCfHasTs()` now returns `Status::InvalidArgument()` if `column_family` is null. `MultiGetEntity()` can rely on this for cfh null check.
- Added `DeleteRange` API using Default Column Family to be consistent with other major APIs (This was also causing Java Test failure after the `FailIfCfHasTs()` change)

Pull Request resolved: https://github.com/facebook/rocksdb/pull/12057

Test Plan:
- Updated `DBWideBasicTest::GetEntityAsPinnableAttributeGroups` to include null CF case
- Updated `DBWideBasicTest::MultiCFMultiGetEntityAsPinnableAttributeGroups` to include null CF case

Reviewed By: jowlyzhang

Differential Revision: D51167445

Pulled By: jaykorean

fbshipit-source-id: 1c1e44fd7b7df4d2dc3bb2d7d251da85bad7d664
This commit is contained in:
Jay Huh 2023-11-13 14:30:04 -08:00 committed by Facebook GitHub Bot
parent b3ffca0e29
commit 8b8f6c63ef
5 changed files with 96 additions and 11 deletions

View File

@ -2062,10 +2062,11 @@ Status DBImpl::GetEntity(const ReadOptions& _read_options, const Slice& key,
return Status::InvalidArgument(
"Cannot call GetEntity without PinnableAttributeGroups object");
}
Status s;
const size_t num_column_families = result->size();
if (_read_options.io_activity != Env::IOActivity::kUnknown &&
_read_options.io_activity != Env::IOActivity::kGetEntity) {
Status s = Status::InvalidArgument(
s = Status::InvalidArgument(
"Cannot call GetEntity with `ReadOptions::io_activity` != "
"`Env::IOActivity::kUnknown` or `Env::IOActivity::kGetEntity`");
for (size_t i = 0; i < num_column_families; ++i) {
@ -2075,7 +2076,7 @@ Status DBImpl::GetEntity(const ReadOptions& _read_options, const Slice& key,
}
// return early if no CF was passed in
if (num_column_families == 0) {
return Status::OK();
return s;
}
ReadOptions read_options(_read_options);
if (read_options.io_activity == Env::IOActivity::kUnknown) {
@ -2084,10 +2085,29 @@ Status DBImpl::GetEntity(const ReadOptions& _read_options, const Slice& key,
std::vector<Slice> keys;
std::vector<ColumnFamilyHandle*> column_families;
for (size_t i = 0; i < num_column_families; ++i) {
// If any of the CFH is null, break early since the entire query will fail
if (!(*result)[i].column_family()) {
s = Status::InvalidArgument(
"DB failed to query because one or more group(s) have null column "
"family handle");
(*result)[i].SetStatus(
Status::InvalidArgument("Column family handle cannot be null"));
break;
}
// Adding the same key slice for different CFs
keys.emplace_back(key);
column_families.emplace_back((*result)[i].column_family());
}
if (!s.ok()) {
for (size_t i = 0; i < num_column_families; ++i) {
if ((*result)[i].status().ok()) {
(*result)[i].SetStatus(
Status::Incomplete("DB not queried due to invalid argument(s) in "
"one or more of the attribute groups"));
}
}
return s;
}
std::vector<PinnableWideColumns> columns(num_column_families);
std::vector<Status> statuses(num_column_families);
MultiGetCommon(
@ -2100,7 +2120,7 @@ Status DBImpl::GetEntity(const ReadOptions& _read_options, const Slice& key,
(*result)[i].SetStatus(statuses[i]);
(*result)[i].SetColumns(std::move(columns[i]));
}
return Status::OK();
return s;
}
bool DBImpl::ShouldReferenceSuperVersion(const MergeContext& merge_context) {
@ -2882,7 +2902,6 @@ void DBImpl::MultiGetCommon(const ReadOptions& read_options,
bool should_fail = false;
for (size_t i = 0; i < num_keys; ++i) {
ColumnFamilyHandle* cfh = column_families[i];
assert(cfh);
if (read_options.timestamp) {
statuses[i] = FailIfTsMismatchCf(cfh, *(read_options.timestamp));
if (!statuses[i].ok()) {

View File

@ -2870,7 +2870,9 @@ static void ClipToRange(T* ptr, V minvalue, V maxvalue) {
inline Status DBImpl::FailIfCfHasTs(
const ColumnFamilyHandle* column_family) const {
column_family = column_family ? column_family : DefaultColumnFamily();
if (!column_family) {
return Status::InvalidArgument("column family handle cannot be null");
}
assert(column_family);
const Comparator* const ucmp = column_family->GetComparator();
assert(ucmp);

View File

@ -277,6 +277,9 @@ TEST_F(DBWideBasicTest, GetEntityAsPinnableAttributeGroups) {
{handles_[kDefaultCfHandleIndex], handles_[kHotCfHandleIndex]}};
std::vector<ColumnFamilyHandle*> hot_and_cold_cfs{
{handles_[kHotCfHandleIndex], handles_[kColdCfHandleIndex]}};
std::vector<ColumnFamilyHandle*> default_null_and_hot_cfs{
handles_[kDefaultCfHandleIndex], nullptr, handles_[kHotCfHandleIndex],
nullptr};
auto create_result =
[](const std::vector<ColumnFamilyHandle*>& column_families)
-> PinnableAttributeGroups {
@ -286,9 +289,29 @@ TEST_F(DBWideBasicTest, GetEntityAsPinnableAttributeGroups) {
}
return result;
};
{
// Case 1. Get first key from default cf and hot_cf and second key from
// Case 1. Invalid Argument (passing in null CF)
AttributeGroups ag{
AttributeGroup(nullptr, first_default_columns),
AttributeGroup(handles_[kHotCfHandleIndex], first_hot_columns)};
ASSERT_NOK(db_->PutEntity(WriteOptions(), first_key, ag));
PinnableAttributeGroups result = create_result(default_null_and_hot_cfs);
Status s = db_->GetEntity(ReadOptions(), first_key, &result);
ASSERT_NOK(s);
ASSERT_TRUE(s.IsInvalidArgument());
// Valid CF, but failed with Incomplete status due to other attribute groups
ASSERT_TRUE(result[0].status().IsIncomplete());
// Null CF
ASSERT_TRUE(result[1].status().IsInvalidArgument());
// Valid CF, but failed with Incomplete status due to other attribute groups
ASSERT_TRUE(result[2].status().IsIncomplete());
// Null CF, but failed with Incomplete status because the nullcheck break
// out early in the loop
ASSERT_TRUE(result[3].status().IsIncomplete());
}
{
// Case 2. Get first key from default cf and hot_cf and second key from
// hot_cf and cold_cf
constexpr size_t num_column_families = 2;
PinnableAttributeGroups first_key_result =
@ -318,7 +341,7 @@ TEST_F(DBWideBasicTest, GetEntityAsPinnableAttributeGroups) {
ASSERT_EQ(second_cold_columns, second_key_result[1].columns());
}
{
// Case 2. Get first key and second key from all cfs. For the second key, we
// Case 3. Get first key and second key from all cfs. For the second key, we
// don't expect to get columns from default cf.
constexpr size_t num_column_families = 3;
PinnableAttributeGroups first_key_result = create_result(all_cfs);
@ -428,6 +451,8 @@ TEST_F(DBWideBasicTest, MultiCFMultiGetEntityAsPinnableAttributeGroups) {
{handles_[kDefaultCfHandleIndex], handles_[kHotCfHandleIndex]}};
std::vector<ColumnFamilyHandle*> hot_and_cold_cfs{
{handles_[kHotCfHandleIndex], handles_[kColdCfHandleIndex]}};
std::vector<ColumnFamilyHandle*> null_and_hot_cfs{
nullptr, handles_[kHotCfHandleIndex], nullptr};
auto create_result =
[](const std::vector<ColumnFamilyHandle*>& column_families)
-> PinnableAttributeGroups {
@ -438,7 +463,7 @@ TEST_F(DBWideBasicTest, MultiCFMultiGetEntityAsPinnableAttributeGroups) {
return result;
};
{
// Check for invalid argument
// Check for invalid read option argument
ReadOptions read_options;
read_options.io_activity = Env::IOActivity::kGetEntity;
std::vector<PinnableAttributeGroups> results;
@ -452,6 +477,31 @@ TEST_F(DBWideBasicTest, MultiCFMultiGetEntityAsPinnableAttributeGroups) {
ASSERT_TRUE(results[i][j].status().IsInvalidArgument());
}
}
// Check for invalid column family in Attribute Group result
results.clear();
results.emplace_back(create_result(null_and_hot_cfs));
results.emplace_back(create_result(all_cfs));
db_->MultiGetEntity(ReadOptions(), num_keys, keys.data(), results.data());
// First one failed due to null CFs in the AttributeGroup
// Null CF
ASSERT_NOK(results[0][0].status());
ASSERT_TRUE(results[0][0].status().IsInvalidArgument());
// Valid CF, but failed with incomplete status because of other attribute
// groups
ASSERT_NOK(results[0][1].status());
ASSERT_TRUE(results[0][1].status().IsIncomplete());
// Null CF
ASSERT_NOK(results[0][2].status());
ASSERT_TRUE(results[0][2].status().IsInvalidArgument());
// Second one failed with Incomplete because first one failed
ASSERT_NOK(results[1][0].status());
ASSERT_TRUE(results[1][0].status().IsIncomplete());
ASSERT_NOK(results[1][1].status());
ASSERT_TRUE(results[1][1].status().IsIncomplete());
ASSERT_NOK(results[1][2].status());
ASSERT_TRUE(results[1][2].status().IsIncomplete());
}
{
// Case 1. Get first key from default cf and hot_cf and second key from

View File

@ -509,6 +509,15 @@ class DB {
ColumnFamilyHandle* column_family,
const Slice& begin_key, const Slice& end_key,
const Slice& ts);
virtual Status DeleteRange(const WriteOptions& options,
const Slice& begin_key, const Slice& end_key) {
return DeleteRange(options, DefaultColumnFamily(), begin_key, end_key);
}
virtual Status DeleteRange(const WriteOptions& options,
const Slice& begin_key, const Slice& end_key,
const Slice& ts) {
return DeleteRange(options, DefaultColumnFamily(), begin_key, end_key, ts);
}
// Merge the database entry for "key" with "value". Returns OK on success,
// and a non-OK status on error. The semantics of this operation is

View File

@ -1036,8 +1036,13 @@ bool rocksdb_delete_range_helper(
ROCKSDB_NAMESPACE::Slice end_key_slice(reinterpret_cast<char*>(end_key),
jend_key_len);
ROCKSDB_NAMESPACE::Status s =
db->DeleteRange(write_options, cf_handle, begin_key_slice, end_key_slice);
ROCKSDB_NAMESPACE::Status s;
if (cf_handle != nullptr) {
s = db->DeleteRange(write_options, cf_handle, begin_key_slice,
end_key_slice);
} else {
s = db->DeleteRange(write_options, begin_key_slice, end_key_slice);
}
// cleanup
delete[] begin_key;