feat(server): bitop command - more code refactor from code review #213 Signed-off-by: Boaz Sade <boaz@dragonflydb.io> Co-authored-by: Boaz Sade <boaz@dragonflydb.io>
This commit is contained in:
Boaz Sade
GitHub
parent
ad46e5e854
commit
ac90aecde1
|
|
@ -121,7 +121,7 @@ with respect to Memcached and Redis APIs.
|
|||
- [X] PREPEND (dragonfly specific)
|
||||
- [x] BITCOUNT
|
||||
- [ ] BITFIELD
|
||||
- [ ] BITOP
|
||||
- [x] BITOP
|
||||
- [ ] BITPOS
|
||||
- [x] GETBIT
|
||||
- [X] GETRANGE
|
||||
|
|
|
|||
|
|
@ -10,8 +10,6 @@ extern "C" {
|
|||
#include "redis/object.h"
|
||||
}
|
||||
|
||||
#include <new>
|
||||
|
||||
#include "base/logging.h"
|
||||
#include "server/command_registry.h"
|
||||
#include "server/common.h"
|
||||
|
|
@ -26,7 +24,15 @@ namespace dfly {
|
|||
using namespace facade;
|
||||
|
||||
namespace {
|
||||
|
||||
using ShardStringResults = std::vector<OpResult<std::string>>;
|
||||
const int32_t OFFSET_FACTOR = 8; // number of bits in byte
|
||||
const char* OR_OP_NAME = "OR";
|
||||
const char* XOR_OP_NAME = "XOR";
|
||||
const char* AND_OP_NAME = "AND";
|
||||
const char* NOT_OP_NAME = "NOT";
|
||||
|
||||
using BitsStrVec = std::vector<std::string>;
|
||||
|
||||
// The following is the list of the functions that would handle the
|
||||
// commands that handle the bit operations
|
||||
|
|
@ -38,17 +44,83 @@ void BitOp(CmdArgList args, ConnectionContext* cntx);
|
|||
void GetBit(CmdArgList args, ConnectionContext* cntx);
|
||||
void SetBit(CmdArgList args, ConnectionContext* cntx);
|
||||
|
||||
OpResult<std::string> ReadValue(const OpArgs& op_args, std::string_view key);
|
||||
OpResult<std::string> ReadValue(const DbContext& context, std::string_view key, EngineShard* shard);
|
||||
OpResult<bool> ReadValueBitsetAt(const OpArgs& op_args, std::string_view key, uint32_t offset);
|
||||
OpResult<std::size_t> CountBitsForValue(const OpArgs& op_args, std::string_view key, int64_t start,
|
||||
int64_t end, bool bit_value);
|
||||
std::string GetString(EngineShard* shard, const PrimeValue& pv);
|
||||
std::string GetString(const PrimeValue& pv, EngineShard* shard);
|
||||
bool SetBitValue(uint32_t offset, bool bit_value, std::string* entry);
|
||||
std::size_t CountBitSetByByteIndices(std::string_view at, std::size_t start, std::size_t end);
|
||||
std::size_t CountBitSet(std::string_view str, int64_t start, int64_t end, bool bits);
|
||||
std::size_t CountBitSetByBitIndices(std::string_view at, std::size_t start, std::size_t end);
|
||||
OpResult<std::string> RunBitOpOnShard(std::string_view op, const OpArgs& op_args, ArgSlice keys);
|
||||
std::string RunBitOperationOnValues(std::string_view op, const BitsStrVec& values);
|
||||
|
||||
// ------------------------------------------------------------------------- //
|
||||
|
||||
// This function can be used for any case where we allowing out of bound
|
||||
// access where the default in this case would be 0 -such as bitop
|
||||
uint8_t GetByteAt(std::string_view s, std::size_t at) {
|
||||
return at >= s.size() ? 0 : s[at];
|
||||
}
|
||||
|
||||
// For XOR, OR, AND operations on a collection of bytes
|
||||
template <typename BitOp, typename SkipOp>
|
||||
std::string BitOpString(BitOp operation_f, SkipOp skip_f, const BitsStrVec& values,
|
||||
std::string&& new_value) {
|
||||
// at this point, values are not empty
|
||||
std::size_t max_size = new_value.size();
|
||||
|
||||
if (values.size() > 1) {
|
||||
for (std::size_t i = 0; i < max_size; i++) {
|
||||
std::uint8_t new_entry = operation_f(GetByteAt(values[0], i), GetByteAt(values[1], i));
|
||||
for (std::size_t j = 2; j < values.size(); ++j) {
|
||||
new_entry = operation_f(new_entry, GetByteAt(values[j], i));
|
||||
if (skip_f(new_entry)) {
|
||||
break;
|
||||
}
|
||||
}
|
||||
new_value[i] = new_entry;
|
||||
}
|
||||
return new_value;
|
||||
} else {
|
||||
return values[0];
|
||||
}
|
||||
}
|
||||
|
||||
// Helper functions to support operations
|
||||
// so we would not need to check which
|
||||
// operations to run in the look (unlike
|
||||
// https://github.com/redis/redis/blob/c2b0c13d5c0fab49131f6f5e844f80bfa43f6219/src/bitops.c#L607)
|
||||
constexpr bool SkipAnd(uint8_t byte) {
|
||||
return byte == 0x0;
|
||||
}
|
||||
|
||||
constexpr bool SkipOr(uint8_t byte) {
|
||||
return byte == 0xff;
|
||||
}
|
||||
|
||||
constexpr bool SkipXor(uint8_t) {
|
||||
return false;
|
||||
}
|
||||
|
||||
constexpr uint8_t AndOp(uint8_t left, uint8_t right) {
|
||||
return left & right;
|
||||
}
|
||||
|
||||
constexpr uint8_t OrOp(uint8_t left, uint8_t right) {
|
||||
return left | right;
|
||||
}
|
||||
|
||||
constexpr uint8_t XorOp(uint8_t left, uint8_t right) {
|
||||
return left ^ right;
|
||||
}
|
||||
|
||||
std::string BitOpNotString(std::string from) {
|
||||
std::transform(from.begin(), from.end(), from.begin(), [](auto c) { return ~c; });
|
||||
return from;
|
||||
}
|
||||
|
||||
// Bits manipulation functions
|
||||
constexpr int32_t GetBitIndex(uint32_t offset) noexcept {
|
||||
return offset % OFFSET_FACTOR;
|
||||
|
|
@ -181,60 +253,233 @@ bool SetBitValue(uint32_t offset, bool bit_value, std::string* entry) {
|
|||
}
|
||||
|
||||
// ------------------------------------------------------------------------- //
|
||||
// Helper functions to access the data or change it
|
||||
|
||||
class OverrideValue {
|
||||
const OpArgs& args_;
|
||||
class ElementAccess {
|
||||
bool added_ = false;
|
||||
PrimeIterator element_iter_;
|
||||
std::string_view key_;
|
||||
DbContext context_;
|
||||
EngineShard* shard_ = nullptr;
|
||||
|
||||
public:
|
||||
explicit OverrideValue(const OpArgs& args) : args_{args} {
|
||||
ElementAccess(std::string_view key, const OpArgs& args) : key_{key}, context_{args.db_cntx} {
|
||||
}
|
||||
|
||||
OpResult<bool> Set(std::string_view key, uint32_t offset, bool bit_value);
|
||||
OpStatus Find(EngineShard* shard);
|
||||
|
||||
bool IsNewEntry() const {
|
||||
CHECK_NOTNULL(shard_);
|
||||
return added_;
|
||||
}
|
||||
|
||||
constexpr DbIndex Index() const {
|
||||
return context_.db_index;
|
||||
}
|
||||
|
||||
std::string Value() const;
|
||||
|
||||
void Commit(std::string_view new_value) const;
|
||||
};
|
||||
|
||||
OpResult<bool> OverrideValue::Set(std::string_view key, uint32_t offset, bool bit_value) {
|
||||
auto& db_slice = args_.shard->db_slice();
|
||||
|
||||
DCHECK(db_slice.IsDbValid(args_.db_cntx.db_index));
|
||||
|
||||
std::pair<PrimeIterator, bool> add_res;
|
||||
OpStatus ElementAccess::Find(EngineShard* shard) {
|
||||
try {
|
||||
add_res = db_slice.AddOrFind(args_.db_cntx, key);
|
||||
std::pair<PrimeIterator, bool> add_res = shard->db_slice().AddOrFind(context_, key_);
|
||||
if (!add_res.second) {
|
||||
if (add_res.first->second.ObjType() != OBJ_STRING) {
|
||||
return OpStatus::WRONG_TYPE;
|
||||
}
|
||||
}
|
||||
element_iter_ = add_res.first;
|
||||
added_ = add_res.second;
|
||||
shard_ = shard;
|
||||
return OpStatus::OK;
|
||||
} catch (const std::bad_alloc&) {
|
||||
return OpStatus::OUT_OF_MEMORY;
|
||||
}
|
||||
bool old_value = false;
|
||||
PrimeIterator& it = add_res.first;
|
||||
bool added = add_res.second;
|
||||
auto UpdateBitMapValue = [&](std::string_view value) {
|
||||
db_slice.PreUpdate(args_.db_cntx.db_index, it);
|
||||
it->second.SetString(value);
|
||||
db_slice.PostUpdate(args_.db_cntx.db_index, it, key, !added);
|
||||
};
|
||||
}
|
||||
|
||||
if (added) { // this is a new entry in the "table"
|
||||
std::string ElementAccess::Value() const {
|
||||
CHECK_NOTNULL(shard_);
|
||||
if (!added_) { // Exist entry - return it
|
||||
return GetString(element_iter_->second, shard_);
|
||||
} else { // we only have reference to the new entry but no value
|
||||
return std::string{};
|
||||
}
|
||||
}
|
||||
|
||||
void ElementAccess::Commit(std::string_view new_value) const {
|
||||
if (shard_) {
|
||||
auto& db_slice = shard_->db_slice();
|
||||
db_slice.PreUpdate(Index(), element_iter_);
|
||||
element_iter_->second.SetString(new_value);
|
||||
db_slice.PostUpdate(Index(), element_iter_, key_, !added_);
|
||||
}
|
||||
}
|
||||
|
||||
// =============================================
|
||||
// Set a new value to a given bit
|
||||
|
||||
OpResult<bool> BitNewValue(const OpArgs& args, std::string_view key, uint32_t offset,
|
||||
bool bit_value) {
|
||||
EngineShard* shard = args.shard;
|
||||
ElementAccess element_access{key, args};
|
||||
auto& db_slice = shard->db_slice();
|
||||
DCHECK(db_slice.IsDbValid(element_access.Index()));
|
||||
bool old_value = false;
|
||||
|
||||
auto find_res = element_access.Find(shard);
|
||||
|
||||
if (find_res != OpStatus::OK) {
|
||||
return find_res;
|
||||
}
|
||||
|
||||
if (element_access.IsNewEntry()) {
|
||||
std::string new_entry(GetByteIndex(offset) + 1, 0);
|
||||
old_value = SetBitValue(offset, bit_value, &new_entry);
|
||||
UpdateBitMapValue(new_entry);
|
||||
element_access.Commit(new_entry);
|
||||
} else {
|
||||
if (it->second.ObjType() != OBJ_STRING) {
|
||||
return OpStatus::WRONG_TYPE;
|
||||
}
|
||||
bool reset = false;
|
||||
std::string existing_entry{GetString(args_.shard, it->second)};
|
||||
if ((existing_entry.size() * OFFSET_FACTOR) <= offset) { // need to resize first
|
||||
std::string existing_entry{element_access.Value()};
|
||||
if ((existing_entry.size() * OFFSET_FACTOR) <= offset) {
|
||||
existing_entry.resize(GetByteIndex(offset) + 1, 0);
|
||||
reset = true;
|
||||
}
|
||||
old_value = SetBitValue(offset, bit_value, &existing_entry);
|
||||
if (reset || old_value != bit_value) { // we made a "real" change to the entry, save it
|
||||
UpdateBitMapValue(existing_entry);
|
||||
element_access.Commit(existing_entry);
|
||||
}
|
||||
}
|
||||
return old_value;
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------
|
||||
|
||||
std::string RunBitOperationOnValues(std::string_view op, const BitsStrVec& values) {
|
||||
// This function accept an operation (either OR, XOR, NOT or OR), and run bit operation
|
||||
// on all the values we got from the database. Note that in case that one of the values
|
||||
// is shorter than the other it would return a 0 and the operation would continue
|
||||
// until we ran the longest value. The function will return the resulting new value
|
||||
std::size_t max_len = 0;
|
||||
std::size_t max_len_index = 0;
|
||||
|
||||
const auto BitOperation = [&]() {
|
||||
if (op == OR_OP_NAME) {
|
||||
std::string default_str{values[max_len_index]};
|
||||
return BitOpString(OrOp, SkipOr, std::move(values), std::move(default_str));
|
||||
} else if (op == XOR_OP_NAME) {
|
||||
return BitOpString(XorOp, SkipXor, std::move(values), std::string(max_len, 0));
|
||||
} else if (op == AND_OP_NAME) {
|
||||
return BitOpString(AndOp, SkipAnd, std::move(values), std::string(max_len, 0));
|
||||
} else if (op == NOT_OP_NAME) {
|
||||
return BitOpNotString(values[0]);
|
||||
} else {
|
||||
LOG(FATAL) << "Operation not supported '" << op << "'";
|
||||
return std::string{}; // otherwise we will have warning of not returning value
|
||||
}
|
||||
};
|
||||
|
||||
if (values.empty()) { // this is ok in case we don't have the src keys
|
||||
return std::string{};
|
||||
}
|
||||
// The new result is the max length input
|
||||
max_len = values[0].size();
|
||||
for (std::size_t i = 1; i < values.size(); ++i) {
|
||||
if (values[i].size() > max_len) {
|
||||
max_len = values[i].size();
|
||||
max_len_index = i;
|
||||
}
|
||||
}
|
||||
return BitOperation();
|
||||
}
|
||||
|
||||
OpResult<std::string> CombineResultOp(ShardStringResults result, std::string_view op) {
|
||||
// take valid result for each shard
|
||||
BitsStrVec values;
|
||||
for (auto&& res : result) {
|
||||
if (res) {
|
||||
auto v = res.value();
|
||||
values.emplace_back(std::move(v));
|
||||
} else {
|
||||
if (res.status() != OpStatus::KEY_NOTFOUND) {
|
||||
// something went wrong, just bale out
|
||||
return res;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// and combine them to single result
|
||||
return RunBitOperationOnValues(op, values);
|
||||
}
|
||||
|
||||
// For bitop not - we cannot accumulate
|
||||
OpResult<std::string> RunBitOpNot(const OpArgs& op_args, ArgSlice keys) {
|
||||
DCHECK(keys.size() == 1);
|
||||
|
||||
EngineShard* es = op_args.shard;
|
||||
// if we found the value, just return, if not found then skip, otherwise report an error
|
||||
auto key = keys.front();
|
||||
OpResult<PrimeIterator> find_res = es->db_slice().Find(op_args.db_cntx, key, OBJ_STRING);
|
||||
if (find_res) {
|
||||
return GetString(find_res.value()->second, es);
|
||||
} else {
|
||||
return find_res.status();
|
||||
}
|
||||
}
|
||||
|
||||
// Read only operation where we are running the bit operation on all the
|
||||
// values that belong to same shard.
|
||||
OpResult<std::string> RunBitOpOnShard(std::string_view op, const OpArgs& op_args, ArgSlice keys) {
|
||||
DCHECK(!keys.empty());
|
||||
if (op == NOT_OP_NAME) {
|
||||
return RunBitOpNot(op_args, keys);
|
||||
}
|
||||
EngineShard* es = op_args.shard;
|
||||
BitsStrVec values;
|
||||
values.reserve(keys.size());
|
||||
|
||||
// collect all the value for this shard
|
||||
for (auto& key : keys) {
|
||||
OpResult<PrimeIterator> find_res = es->db_slice().Find(op_args.db_cntx, key, OBJ_STRING);
|
||||
if (find_res) {
|
||||
values.emplace_back(std::move(GetString(find_res.value()->second, es)));
|
||||
} else {
|
||||
if (find_res.status() == OpStatus::KEY_NOTFOUND) {
|
||||
continue; // this is allowed, just return empty string per Redis
|
||||
} else {
|
||||
return find_res.status();
|
||||
}
|
||||
}
|
||||
}
|
||||
// Run the operation on all the values that we found
|
||||
std::string op_result = RunBitOperationOnValues(op, values);
|
||||
return op_result;
|
||||
}
|
||||
|
||||
template <typename T> void HandleOpValueResult(const OpResult<T>& result, ConnectionContext* cntx) {
|
||||
static_assert(std::is_integral<T>::value,
|
||||
"we are only handling types that are integral types in the return types from "
|
||||
"here");
|
||||
if (result) {
|
||||
(*cntx)->SendLong(result.value());
|
||||
} else {
|
||||
switch (result.status()) {
|
||||
case OpStatus::WRONG_TYPE:
|
||||
(*cntx)->SendError(kWrongTypeErr);
|
||||
break;
|
||||
case OpStatus::OUT_OF_MEMORY:
|
||||
(*cntx)->SendError(kOutOfMemory);
|
||||
break;
|
||||
default:
|
||||
(*cntx)->SendLong(0); // in case we don't have the value we should just send 0
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
OpStatus NoOpCb(Transaction* t, EngineShard* shard) {
|
||||
return OpStatus::OK;
|
||||
}
|
||||
|
||||
// ------------------------------------------------------------------------- //
|
||||
// Impl for the command functions
|
||||
void BitPos(CmdArgList args, ConnectionContext* cntx) {
|
||||
|
|
@ -268,19 +513,8 @@ void BitCount(CmdArgList args, ConnectionContext* cntx) {
|
|||
return CountBitsForValue(t->GetOpArgs(shard), key, start, end, as_bit);
|
||||
};
|
||||
Transaction* trans = cntx->transaction;
|
||||
OpResult<std::size_t> result = trans->ScheduleSingleHopT(std::move(cb));
|
||||
if (result) {
|
||||
(*cntx)->SendLong(result.value());
|
||||
} else {
|
||||
switch (result.status()) {
|
||||
case OpStatus::WRONG_TYPE:
|
||||
(*cntx)->SendError(kWrongTypeErr);
|
||||
break;
|
||||
default:
|
||||
(*cntx)->SendLong(0);
|
||||
break;
|
||||
}
|
||||
}
|
||||
OpResult<std::size_t> res = trans->ScheduleSingleHopT(std::move(cb));
|
||||
HandleOpValueResult(res, cntx);
|
||||
}
|
||||
|
||||
void BitField(CmdArgList args, ConnectionContext* cntx) {
|
||||
|
|
@ -292,7 +526,64 @@ void BitFieldRo(CmdArgList args, ConnectionContext* cntx) {
|
|||
}
|
||||
|
||||
void BitOp(CmdArgList args, ConnectionContext* cntx) {
|
||||
(*cntx)->SendOk();
|
||||
static const std::array<std::string_view, 4> BITOP_OP_NAMES{OR_OP_NAME, XOR_OP_NAME, AND_OP_NAME,
|
||||
NOT_OP_NAME};
|
||||
ToUpper(&args[1]);
|
||||
std::string_view op = ArgS(args, 1);
|
||||
std::string_view dest_key = ArgS(args, 2);
|
||||
bool illegal = std::none_of(BITOP_OP_NAMES.begin(), BITOP_OP_NAMES.end(),
|
||||
[&op](auto val) { return op == val; });
|
||||
|
||||
if (illegal || (op == NOT_OP_NAME && args.size() > 4)) {
|
||||
return (*cntx)->SendError(kSyntaxErr); // too many arguments
|
||||
}
|
||||
|
||||
// Multi shard access - read only
|
||||
ShardStringResults result_set(shard_set->size(), OpStatus::KEY_NOTFOUND);
|
||||
ShardId dest_shard = Shard(dest_key, result_set.size());
|
||||
|
||||
auto shard_bitop = [&](Transaction* t, EngineShard* shard) {
|
||||
ArgSlice largs = t->ShardArgsInShard(shard->shard_id());
|
||||
DCHECK(!largs.empty());
|
||||
|
||||
if (shard->shard_id() == dest_shard) {
|
||||
CHECK_EQ(largs.front(), dest_key);
|
||||
largs.remove_prefix(1);
|
||||
if (largs.empty()) { // no more keys to check
|
||||
return OpStatus::OK;
|
||||
}
|
||||
}
|
||||
OpArgs op_args = t->GetOpArgs(shard);
|
||||
result_set[shard->shard_id()] = RunBitOpOnShard(op, op_args, largs);
|
||||
return OpStatus::OK;
|
||||
};
|
||||
|
||||
cntx->transaction->Schedule();
|
||||
cntx->transaction->Execute(std::move(shard_bitop), false); // we still have more work to do
|
||||
// All result from each shard
|
||||
const auto joined_results = CombineResultOp(result_set, op);
|
||||
// Second phase - save to targe key if successful
|
||||
if (!joined_results) {
|
||||
cntx->transaction->Execute(NoOpCb, true);
|
||||
(*cntx)->SendError(joined_results.status());
|
||||
return;
|
||||
} else {
|
||||
auto op_result = joined_results.value();
|
||||
auto store_cb = [&](Transaction* t, EngineShard* shard) {
|
||||
if (shard->shard_id() == dest_shard) {
|
||||
ElementAccess operation{dest_key, t->GetOpArgs(shard)};
|
||||
auto find_res = operation.Find(shard);
|
||||
|
||||
if (find_res == OpStatus::OK) {
|
||||
operation.Commit(op_result);
|
||||
}
|
||||
}
|
||||
return OpStatus::OK;
|
||||
};
|
||||
|
||||
cntx->transaction->Execute(std::move(store_cb), true);
|
||||
(*cntx)->SendLong(op_result.size());
|
||||
}
|
||||
}
|
||||
|
||||
void GetBit(CmdArgList args, ConnectionContext* cntx) {
|
||||
|
|
@ -309,24 +600,8 @@ void GetBit(CmdArgList args, ConnectionContext* cntx) {
|
|||
return ReadValueBitsetAt(t->GetOpArgs(shard), key, offset);
|
||||
};
|
||||
Transaction* trans = cntx->transaction;
|
||||
OpResult<bool> result = trans->ScheduleSingleHopT(std::move(cb));
|
||||
|
||||
if (result) {
|
||||
DVLOG(2) << "GET" << trans->DebugId() << "': key: '" << key << ", value '" << result.value()
|
||||
<< "'\n";
|
||||
// we have the value, now we need to get the bit at the location
|
||||
long val = result.value() ? 1 : 0;
|
||||
(*cntx)->SendLong(val);
|
||||
} else {
|
||||
switch (result.status()) {
|
||||
case OpStatus::WRONG_TYPE:
|
||||
(*cntx)->SendError(kWrongTypeErr);
|
||||
break;
|
||||
default:
|
||||
DVLOG(2) << "GET " << key << " nil";
|
||||
(*cntx)->SendLong(0); // in case we don't have the value we should just send 0
|
||||
}
|
||||
}
|
||||
OpResult<bool> res = trans->ScheduleSingleHopT(std::move(cb));
|
||||
HandleOpValueResult(res, cntx);
|
||||
}
|
||||
|
||||
void SetBit(CmdArgList args, ConnectionContext* cntx) {
|
||||
|
|
@ -342,35 +617,17 @@ void SetBit(CmdArgList args, ConnectionContext* cntx) {
|
|||
}
|
||||
|
||||
auto cb = [&](Transaction* t, EngineShard* shard) {
|
||||
OverrideValue set_operation{t->GetOpArgs(shard)};
|
||||
|
||||
return set_operation.Set(key, offset, value != 0);
|
||||
return BitNewValue(t->GetOpArgs(shard), key, offset, value != 0);
|
||||
};
|
||||
|
||||
Transaction* trans = cntx->transaction;
|
||||
OpResult<bool> result = trans->ScheduleSingleHopT(std::move(cb));
|
||||
if (result) {
|
||||
long res = result.value() ? 1 : 0;
|
||||
(*cntx)->SendLong(res);
|
||||
} else {
|
||||
switch (result.status()) {
|
||||
case OpStatus::WRONG_TYPE:
|
||||
(*cntx)->SendError(kWrongTypeErr);
|
||||
break;
|
||||
case OpStatus::OUT_OF_MEMORY:
|
||||
(*cntx)->SendError(kOutOfMemory);
|
||||
break;
|
||||
default:
|
||||
DVLOG(2) << "SETBIT " << key << " nil" << result.status();
|
||||
(*cntx)->SendLong(0); // in case we don't have the value we should just send 0
|
||||
break;
|
||||
}
|
||||
}
|
||||
OpResult<bool> res = trans->ScheduleSingleHopT(std::move(cb));
|
||||
HandleOpValueResult(res, cntx);
|
||||
}
|
||||
|
||||
// ------------------------------------------------------------------------- //
|
||||
// This are the "callbacks" that we're using from above
|
||||
std::string GetString(EngineShard* shard, const PrimeValue& pv) {
|
||||
std::string GetString(const PrimeValue& pv, EngineShard* shard) {
|
||||
std::string res;
|
||||
if (pv.IsExternal()) {
|
||||
auto* tiered = shard->tiered_storage();
|
||||
|
|
@ -387,7 +644,7 @@ std::string GetString(EngineShard* shard, const PrimeValue& pv) {
|
|||
}
|
||||
|
||||
OpResult<bool> ReadValueBitsetAt(const OpArgs& op_args, std::string_view key, uint32_t offset) {
|
||||
OpResult<std::string> result = ReadValue(op_args, key);
|
||||
OpResult<std::string> result = ReadValue(op_args.db_cntx, key, op_args.shard);
|
||||
if (result) {
|
||||
return GetBitValueSafe(result.value(), offset);
|
||||
} else {
|
||||
|
|
@ -395,22 +652,23 @@ OpResult<bool> ReadValueBitsetAt(const OpArgs& op_args, std::string_view key, ui
|
|||
}
|
||||
}
|
||||
|
||||
OpResult<std::string> ReadValue(const OpArgs& op_args, std::string_view key) {
|
||||
OpResult<PrimeIterator> it_res = op_args.shard->db_slice().Find(op_args.db_cntx, key, OBJ_STRING);
|
||||
OpResult<std::string> ReadValue(const DbContext& context, std::string_view key,
|
||||
EngineShard* shard) {
|
||||
OpResult<PrimeIterator> it_res = shard->db_slice().Find(context, key, OBJ_STRING);
|
||||
if (!it_res.ok()) {
|
||||
return it_res.status();
|
||||
}
|
||||
|
||||
const PrimeValue& pv = it_res.value()->second;
|
||||
|
||||
return GetString(op_args.shard, pv);
|
||||
return GetString(pv, shard);
|
||||
}
|
||||
|
||||
OpResult<std::size_t> CountBitsForValue(const OpArgs& op_args, std::string_view key, int64_t start,
|
||||
int64_t end, bool bit_value) {
|
||||
OpResult<std::string> result = ReadValue(op_args, key);
|
||||
OpResult<std::string> result = ReadValue(op_args.db_cntx, key, op_args.shard);
|
||||
|
||||
if (result) {
|
||||
if (result) { // if this is not found, just return 0 - per Redis
|
||||
if (result.value().empty()) {
|
||||
return 0;
|
||||
}
|
||||
|
|
@ -432,7 +690,7 @@ void BitOpsFamily::Register(CommandRegistry* registry) {
|
|||
<< CI{"BITCOUNT", CO::READONLY, -2, 1, 1, 1}.SetHandler(&BitCount)
|
||||
<< CI{"BITFIELD", CO::WRITE, -3, 1, 1, 1}.SetHandler(&BitField)
|
||||
<< CI{"BITFIELD_RO", CO::READONLY, -5, 1, 1, 1}.SetHandler(&BitFieldRo)
|
||||
<< CI{"BITOP", CO::WRITE, -4, 1, 1, 1}.SetHandler(&BitOp)
|
||||
<< CI{"BITOP", CO::WRITE, -4, 2, -1, 1}.SetHandler(&BitOp)
|
||||
<< CI{"GETBIT", CO::READONLY | CO::FAST | CO::FAST, 3, 1, 1, 1}.SetHandler(&GetBit)
|
||||
<< CI{"SETBIT", CO::WRITE, 4, 1, 1, 1}.SetHandler(&SetBit);
|
||||
}
|
||||
|
|
|
|||
|
|
@ -4,6 +4,12 @@
|
|||
|
||||
#include "server/bitops_family.h"
|
||||
|
||||
#include <bitset>
|
||||
#include <iomanip>
|
||||
#include <iostream>
|
||||
#include <string>
|
||||
#include <string_view>
|
||||
|
||||
#include "base/gtest.h"
|
||||
#include "base/logging.h"
|
||||
#include "facade/facade_test.h"
|
||||
|
|
@ -22,10 +28,163 @@ using absl::StrCat;
|
|||
|
||||
namespace dfly {
|
||||
|
||||
class Bytes {
|
||||
using char_t = std::uint8_t;
|
||||
using string_type = std::basic_string<char_t>;
|
||||
|
||||
public:
|
||||
enum State { GOOD, ERROR, NIL };
|
||||
|
||||
Bytes(std::initializer_list<std::uint8_t> bytes) : data_(bytes.size(), 0) {
|
||||
// note - we want this to be like its would be used in redis where most significate bit is to
|
||||
// the "left"
|
||||
std::copy(rbegin(bytes), rend(bytes), data_.begin());
|
||||
}
|
||||
|
||||
explicit Bytes(unsigned long long n) : data_(sizeof(n), 0) {
|
||||
FromNumber(n);
|
||||
}
|
||||
|
||||
static Bytes From(unsigned long long x) {
|
||||
return Bytes(x);
|
||||
}
|
||||
|
||||
explicit Bytes(State state) : state_{state} {
|
||||
}
|
||||
|
||||
Bytes(const char_t* ch, std::size_t len) : data_(ch, len) {
|
||||
}
|
||||
|
||||
Bytes(const char* ch, std::size_t len) : Bytes(reinterpret_cast<const char_t*>(ch), len) {
|
||||
}
|
||||
|
||||
explicit Bytes(std::string_view from) : Bytes(from.data(), from.size()) {
|
||||
}
|
||||
|
||||
static Bytes From(RespExpr&& r);
|
||||
|
||||
std::size_t Size() const {
|
||||
return data_.size();
|
||||
}
|
||||
|
||||
operator std::string_view() const {
|
||||
return std::string_view(reinterpret_cast<const char*>(data_.data()), Size());
|
||||
}
|
||||
|
||||
std::ostream& Print(std::ostream& os) const;
|
||||
|
||||
std::ostream& PrintHex(std::ostream& os) const;
|
||||
|
||||
private:
|
||||
template <typename T> void FromNumber(T num) {
|
||||
// note - we want this to be like its would be used in redis where most significate bit is to
|
||||
// the "left"
|
||||
std::size_t i = 0;
|
||||
for (const char_t* s = reinterpret_cast<const char_t*>(&num); i < sizeof(T); s++, i++) {
|
||||
data_[i] = *s;
|
||||
}
|
||||
}
|
||||
|
||||
string_type data_;
|
||||
State state_ = GOOD;
|
||||
};
|
||||
|
||||
Bytes Bytes::From(RespExpr&& r) {
|
||||
if (r.type == RespExpr::STRING) {
|
||||
return Bytes(ToSV(r.GetBuf()));
|
||||
} else {
|
||||
if (r.type == RespExpr::NIL || r.type == RespExpr::NIL_ARRAY) {
|
||||
return Bytes{Bytes::NIL};
|
||||
} else {
|
||||
return Bytes(Bytes::ERROR);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
std::ostream& Bytes::Print(std::ostream& os) const {
|
||||
if (state_ == GOOD) {
|
||||
for (auto c : data_) {
|
||||
std::bitset<8> b{c};
|
||||
os << b << ":";
|
||||
}
|
||||
} else {
|
||||
if (state_ == NIL) {
|
||||
os << "nil";
|
||||
} else {
|
||||
os << "error";
|
||||
}
|
||||
}
|
||||
return os;
|
||||
}
|
||||
|
||||
std::ostream& Bytes::PrintHex(std::ostream& os) const {
|
||||
if (state_ == GOOD) {
|
||||
for (auto c : data_) {
|
||||
os << std::hex << std::setfill('0') << std::setw(2) << (std::uint16_t)c << ":";
|
||||
}
|
||||
} else {
|
||||
if (state_ == NIL) {
|
||||
os << "nil";
|
||||
} else {
|
||||
os << "error";
|
||||
}
|
||||
}
|
||||
return os;
|
||||
}
|
||||
|
||||
inline bool operator==(const Bytes& left, const Bytes& right) {
|
||||
return static_cast<const std::string_view&>(left) == static_cast<const std::string_view&>(right);
|
||||
}
|
||||
|
||||
inline bool operator!=(const Bytes& left, const Bytes& right) {
|
||||
return !(left == right);
|
||||
}
|
||||
|
||||
inline Bytes operator"" _b(unsigned long long x) {
|
||||
return Bytes::From(x);
|
||||
}
|
||||
|
||||
inline Bytes operator"" _b(const char* x, std::size_t s) {
|
||||
return Bytes{x, s};
|
||||
}
|
||||
|
||||
inline Bytes operator"" _b(const char* x) {
|
||||
return Bytes{x, std::strlen(x)};
|
||||
}
|
||||
|
||||
inline std::ostream& operator<<(std::ostream& os, const Bytes& bs) {
|
||||
return bs.PrintHex(os);
|
||||
}
|
||||
|
||||
class BitOpsFamilyTest : public BaseFamilyTest {
|
||||
protected:
|
||||
// only for bitop XOR, OR, AND tests
|
||||
void BitOpSetKeys();
|
||||
};
|
||||
|
||||
// for the bitop tests we need to test with multiple keys as the issue
|
||||
// is that we need to make sure that accessing multiple shards creates
|
||||
// the correct result
|
||||
// Since this is bit operations, we are using the bytes data type
|
||||
// that makes the verification more ergonomics.
|
||||
const std::pair<std::string_view, Bytes> KEY_VALUES_BIT_OP[] = {
|
||||
{"first_key", 0xFFAACC01_b},
|
||||
{"key_second", {0x1, 0xBB}},
|
||||
{"_this_is_the_third_key", {0x01, 0x05, 0x15, 0x20, 0xAA, 0xCC}},
|
||||
{"the_last_key_we_have", 0xAACC_b}};
|
||||
|
||||
// For the bitop XOR OR and AND we are setting these keys/value pairs
|
||||
void BitOpsFamilyTest::BitOpSetKeys() {
|
||||
auto resp = Run({"set", KEY_VALUES_BIT_OP[0].first, KEY_VALUES_BIT_OP[0].second});
|
||||
EXPECT_EQ(resp, "OK");
|
||||
resp = Run({"set", KEY_VALUES_BIT_OP[1].first, KEY_VALUES_BIT_OP[1].second});
|
||||
EXPECT_EQ(resp, "OK");
|
||||
resp = Run({"set", KEY_VALUES_BIT_OP[2].first, KEY_VALUES_BIT_OP[2].second});
|
||||
EXPECT_EQ(resp, "OK");
|
||||
resp = Run({"set", KEY_VALUES_BIT_OP[3].first, KEY_VALUES_BIT_OP[3].second});
|
||||
EXPECT_EQ(resp, "OK");
|
||||
}
|
||||
|
||||
const long EXPECTED_VALUE_SETBIT[] = {0, 1, 1, 0, 0, 0,
|
||||
0, 1, 0, 1, 1, 0}; // taken from running this on redis
|
||||
const int32_t ITERATIONS = sizeof(EXPECTED_VALUE_SETBIT) / sizeof(EXPECTED_VALUE_SETBIT[0]);
|
||||
|
|
@ -69,7 +228,6 @@ TEST_F(BitOpsFamilyTest, SetBitMissingKey) {
|
|||
// get 0s since we didn't have this key before
|
||||
EXPECT_EQ(0, CheckedInt({"setbit", "foo", std::to_string(i), "1"}));
|
||||
}
|
||||
|
||||
// now all that we set are at 1s
|
||||
for (int32_t i = 0; i < ITERATIONS; i++) {
|
||||
EXPECT_EQ(1, CheckedInt({"getbit", "foo", std::to_string(i)}));
|
||||
|
|
@ -126,4 +284,141 @@ TEST_F(BitOpsFamilyTest, BitCountByteBitSubRange) {
|
|||
EXPECT_EQ(0, CheckedInt({"bitcount", "foo", "-1", "-2", "bit"})); // illegal range
|
||||
}
|
||||
|
||||
// ------------------------- BITOP tests
|
||||
|
||||
const auto EXPECTED_LEN_BITOP =
|
||||
std::max(KEY_VALUES_BIT_OP[0].second.Size(), KEY_VALUES_BIT_OP[1].second.Size());
|
||||
const auto EXPECTED_LEN_BITOP2 = std::max(EXPECTED_LEN_BITOP, KEY_VALUES_BIT_OP[2].second.Size());
|
||||
const auto EXPECTED_LEN_BITOP3 = std::max(EXPECTED_LEN_BITOP2, KEY_VALUES_BIT_OP[3].second.Size());
|
||||
|
||||
TEST_F(BitOpsFamilyTest, BitOpsAnd) {
|
||||
BitOpSetKeys();
|
||||
auto resp = Run({"bitop", "foo", "bar", "abc"}); // should failed this is illegal operation
|
||||
ASSERT_THAT(resp, ErrArg("syntax error"));
|
||||
// run with none existing keys, should return 0
|
||||
EXPECT_EQ(0, CheckedInt({"bitop", "and", "dest_key", "1", "2", "3"}));
|
||||
|
||||
// bitop AND single key
|
||||
EXPECT_EQ(KEY_VALUES_BIT_OP[0].second.Size(),
|
||||
CheckedInt({"bitop", "and", "foo_out", KEY_VALUES_BIT_OP[0].first}));
|
||||
|
||||
auto res = Bytes::From(Run({"get", "foo_out"}));
|
||||
EXPECT_EQ(res, KEY_VALUES_BIT_OP[0].second);
|
||||
|
||||
// this will 0 all values other than one bit it would end with result with length ==
|
||||
// FOO_KEY_VALUE && value == BAR_KEY_VALUE
|
||||
EXPECT_EQ(EXPECTED_LEN_BITOP, CheckedInt({"bitop", "and", "foo-out", KEY_VALUES_BIT_OP[0].first,
|
||||
KEY_VALUES_BIT_OP[1].first}));
|
||||
const auto EXPECTED_RESULT = Bytes((0xffaacc01 & 0x1BB)); // first and second values
|
||||
res = Bytes::From(Run({"get", "foo-out"}));
|
||||
EXPECT_EQ(res, EXPECTED_RESULT);
|
||||
|
||||
// test bitop AND with 3 keys
|
||||
EXPECT_EQ(EXPECTED_LEN_BITOP2,
|
||||
CheckedInt({"bitop", "and", "foo-out", KEY_VALUES_BIT_OP[0].first,
|
||||
KEY_VALUES_BIT_OP[1].first, KEY_VALUES_BIT_OP[2].first}));
|
||||
const auto EXPECTED_RES2 = Bytes((0xffaacc01 & 0x1BB & 0x01051520AACC));
|
||||
res = Bytes::From(Run({"get", "foo-out"}));
|
||||
EXPECT_EQ(EXPECTED_RES2, res);
|
||||
|
||||
// test bitop AND with 4 parameters
|
||||
const auto EXPECTED_RES3 = Bytes((0xffaacc01 & 0x1BB & 0x01051520AACC & 0xAACC));
|
||||
EXPECT_EQ(EXPECTED_LEN_BITOP3, CheckedInt({"bitop", "and", "foo-out", KEY_VALUES_BIT_OP[0].first,
|
||||
KEY_VALUES_BIT_OP[1].first, KEY_VALUES_BIT_OP[2].first,
|
||||
KEY_VALUES_BIT_OP[3].first}));
|
||||
res = Bytes::From(Run({"get", "foo-out"}));
|
||||
EXPECT_EQ(EXPECTED_RES3, res);
|
||||
}
|
||||
|
||||
TEST_F(BitOpsFamilyTest, BitOpsOr) {
|
||||
BitOpSetKeys();
|
||||
|
||||
EXPECT_EQ(0, CheckedInt({"bitop", "or", "dest_key", "1", "2", "3"}));
|
||||
|
||||
// bitop or single key
|
||||
EXPECT_EQ(KEY_VALUES_BIT_OP[0].second.Size(),
|
||||
CheckedInt({"bitop", "or", "foo_out", KEY_VALUES_BIT_OP[0].first}));
|
||||
|
||||
auto res = Bytes::From(Run({"get", "foo_out"}));
|
||||
EXPECT_EQ(res, KEY_VALUES_BIT_OP[0].second);
|
||||
|
||||
// bitop OR 2 keys
|
||||
EXPECT_EQ(EXPECTED_LEN_BITOP, CheckedInt({"bitop", "or", "foo-out", KEY_VALUES_BIT_OP[0].first,
|
||||
KEY_VALUES_BIT_OP[1].first}));
|
||||
const auto EXPECTED_RESULT = Bytes((0xffaacc01 | 0x1BB)); // first or second values
|
||||
res = Bytes::From(Run({"get", "foo-out"}));
|
||||
EXPECT_EQ(res, EXPECTED_RESULT);
|
||||
|
||||
// bitop OR with 3 keys
|
||||
EXPECT_EQ(EXPECTED_LEN_BITOP2,
|
||||
CheckedInt({"bitop", "or", "foo-out", KEY_VALUES_BIT_OP[0].first,
|
||||
KEY_VALUES_BIT_OP[1].first, KEY_VALUES_BIT_OP[2].first}));
|
||||
const auto EXPECTED_RES2 = Bytes((0xffaacc01 | 0x1BB | 0x01051520AACC));
|
||||
res = Bytes::From(Run({"get", "foo-out"}));
|
||||
EXPECT_EQ(EXPECTED_RES2, res);
|
||||
|
||||
// bitop OR with 4 keys
|
||||
const auto EXPECTED_RES3 = Bytes((0xffaacc01 | 0x1BB | 0x01051520AACC | 0xAACC));
|
||||
EXPECT_EQ(EXPECTED_LEN_BITOP3, CheckedInt({"bitop", "or", "foo-out", KEY_VALUES_BIT_OP[0].first,
|
||||
KEY_VALUES_BIT_OP[1].first, KEY_VALUES_BIT_OP[2].first,
|
||||
KEY_VALUES_BIT_OP[3].first}));
|
||||
res = Bytes::From(Run({"get", "foo-out"}));
|
||||
EXPECT_EQ(EXPECTED_RES3, res);
|
||||
}
|
||||
|
||||
TEST_F(BitOpsFamilyTest, BitOpsXor) {
|
||||
BitOpSetKeys();
|
||||
|
||||
EXPECT_EQ(0, CheckedInt({"bitop", "or", "dest_key", "1", "2", "3"}));
|
||||
|
||||
// bitop XOR on single key
|
||||
EXPECT_EQ(KEY_VALUES_BIT_OP[0].second.Size(),
|
||||
CheckedInt({"bitop", "xor", "foo_out", KEY_VALUES_BIT_OP[0].first}));
|
||||
auto res = Bytes::From(Run({"get", "foo_out"}));
|
||||
EXPECT_EQ(res, KEY_VALUES_BIT_OP[0].second);
|
||||
|
||||
// bitop on XOR with two keys
|
||||
EXPECT_EQ(EXPECTED_LEN_BITOP, CheckedInt({"bitop", "xor", "foo-out", KEY_VALUES_BIT_OP[0].first,
|
||||
KEY_VALUES_BIT_OP[1].first}));
|
||||
const auto EXPECTED_RESULT = Bytes((0xffaacc01 ^ 0x1BB)); // first xor second values
|
||||
res = Bytes::From(Run({"get", "foo-out"}));
|
||||
EXPECT_EQ(res, EXPECTED_RESULT);
|
||||
|
||||
// bitop XOR with 3 keys
|
||||
EXPECT_EQ(EXPECTED_LEN_BITOP2,
|
||||
CheckedInt({"bitop", "xor", "foo-out", KEY_VALUES_BIT_OP[0].first,
|
||||
KEY_VALUES_BIT_OP[1].first, KEY_VALUES_BIT_OP[2].first}));
|
||||
const auto EXPECTED_RES2 = Bytes((0xffaacc01 ^ 0x1BB ^ 0x01051520AACC));
|
||||
res = Bytes::From(Run({"get", "foo-out"}));
|
||||
EXPECT_EQ(EXPECTED_RES2, res);
|
||||
|
||||
// bitop XOR with 4 keys
|
||||
const auto EXPECTED_RES3 = Bytes((0xffaacc01 ^ 0x1BB ^ 0x01051520AACC ^ 0xAACC));
|
||||
EXPECT_EQ(EXPECTED_LEN_BITOP3, CheckedInt({"bitop", "xor", "foo-out", KEY_VALUES_BIT_OP[0].first,
|
||||
KEY_VALUES_BIT_OP[1].first, KEY_VALUES_BIT_OP[2].first,
|
||||
KEY_VALUES_BIT_OP[3].first}));
|
||||
res = Bytes::From(Run({"get", "foo-out"}));
|
||||
EXPECT_EQ(EXPECTED_RES3, res);
|
||||
}
|
||||
|
||||
TEST_F(BitOpsFamilyTest, BitOpsNot) {
|
||||
// should failed this is illegal number of args
|
||||
auto resp = Run({"bitop", "not", "bar", "abc", "efg"});
|
||||
ASSERT_THAT(resp, ErrArg("syntax error"));
|
||||
|
||||
// Make sure that this works with none existing key as well
|
||||
EXPECT_EQ(0, CheckedInt({"bitop", "NOT", "bit-op-not-none-existing-key-results",
|
||||
"this-key-do-not-exists"}));
|
||||
EXPECT_EQ(Run({"get", "bit-op-not-none-existing-key-results"}), "");
|
||||
|
||||
// test bitop not
|
||||
resp = Run({"set", KEY_VALUES_BIT_OP[0].first, KEY_VALUES_BIT_OP[0].second});
|
||||
EXPECT_EQ(KEY_VALUES_BIT_OP[0].second.Size(),
|
||||
CheckedInt({"bitop", "not", "foo_out", KEY_VALUES_BIT_OP[0].first}));
|
||||
auto res = Bytes::From(Run({"get", "foo_out"}));
|
||||
|
||||
const auto NOT_RESULTS = Bytes(~0xFFAACC01ull);
|
||||
EXPECT_EQ(res, NOT_RESULTS);
|
||||
}
|
||||
|
||||
} // end of namespace dfly
|
||||
|
|
|
|||
Loading…
Reference in New Issue