From 6ec453a59960dd5e0b0e89831655ea19a74b0445 Mon Sep 17 00:00:00 2001 From: Boaz Sade Date: Sun, 18 Sep 2022 16:20:51 +0300 Subject: [PATCH] feat(server): adding support for bitops #213 (#295) feat(server): bitset bitget commands - update docs/api_status.md #213 Signed-off-by: Boaz Sade --- docs/api_status.md | 4 +- src/server/CMakeLists.txt | 5 +- src/server/bitops_family.cc | 302 +++++++++++++++++++++++++++++++ src/server/bitops_family.h | 30 +++ src/server/bitops_family_test.cc | 79 ++++++++ src/server/main_service.cc | 5 +- src/server/string_family.h | 2 +- 7 files changed, 421 insertions(+), 6 deletions(-) create mode 100644 src/server/bitops_family.cc create mode 100644 src/server/bitops_family.h create mode 100644 src/server/bitops_family_test.cc diff --git a/docs/api_status.md b/docs/api_status.md index 0f9ee3a1d..6fd2af08a 100644 --- a/docs/api_status.md +++ b/docs/api_status.md @@ -123,11 +123,11 @@ with respect to Memcached and Redis APIs. - [ ] BITFIELD - [ ] BITOP - [ ] BITPOS - - [ ] GETBIT + - [x] GETBIT - [X] GETRANGE - [X] INCRBYFLOAT - [X] PSETEX - - [ ] SETBIT + - [x] SETBIT - [X] SETRANGE - [X] STRLEN - [X] HashSet Family diff --git a/src/server/CMakeLists.txt b/src/server/CMakeLists.txt index ad5a90aba..6a00fdd0f 100644 --- a/src/server/CMakeLists.txt +++ b/src/server/CMakeLists.txt @@ -12,7 +12,7 @@ add_library(dragonfly_lib channel_slice.cc command_registry.cc list_family.cc main_service.cc rdb_load.cc rdb_save.cc replica.cc snapshot.cc script_mgr.cc server_family.cc malloc_stats.cc set_family.cc stream_family.cc string_family.cc - zset_family.cc version.cc) + zset_family.cc version.cc bitops_family.cc) cxx_link(dragonfly_lib dfly_transaction dfly_facade redis_lib strings_lib html_lib absl::random_random TRDP::jsoncons) @@ -27,6 +27,7 @@ cxx_test(list_family_test dfly_test_lib LABELS DFLY) cxx_test(set_family_test dfly_test_lib LABELS DFLY) cxx_test(stream_family_test dfly_test_lib LABELS DFLY) cxx_test(string_family_test dfly_test_lib LABELS DFLY) +cxx_test(bitops_family_test dfly_test_lib LABELS DFLY) cxx_test(rdb_test dfly_test_lib DATA testdata/empty.rdb testdata/redis6_small.rdb testdata/redis6_stream.rdb LABELS DFLY) cxx_test(zset_family_test dfly_test_lib LABELS DFLY) @@ -38,4 +39,4 @@ cxx_test(json_family_test dfly_test_lib LABELS DFLY) add_custom_target(check_dfly WORKING_DIRECTORY .. COMMAND ctest -L DFLY) add_dependencies(check_dfly dragonfly_test json_family_test list_family_test generic_family_test memcache_parser_test rdb_test - redis_parser_test snapshot_test stream_family_test string_family_test) + redis_parser_test snapshot_test stream_family_test string_family_test bitops_family_test) diff --git a/src/server/bitops_family.cc b/src/server/bitops_family.cc new file mode 100644 index 000000000..3a8076b3b --- /dev/null +++ b/src/server/bitops_family.cc @@ -0,0 +1,302 @@ +// Copyright 2022, Roman Gershman. All rights reserved. +// See LICENSE for licensing terms. +// + +#include "server/bitops_family.h" + +#include + +extern "C" { +#include "redis/object.h" +} + +#include + +#include "base/logging.h" +#include "server/command_registry.h" +#include "server/common.h" +#include "server/conn_context.h" +#include "server/engine_shard_set.h" +#include "server/error.h" +#include "server/tiered_storage.h" +#include "server/transaction.h" +#include "util/varz.h" + +namespace dfly { +using namespace facade; + +namespace { +static const int32_t OFFSET_FACTOR = 8; // number of bits in byte + +// The following is the list of the functions that would handle the +// commands that handle the bit operations +void BitPos(CmdArgList args, ConnectionContext* cntx); +void BitCount(CmdArgList args, ConnectionContext* cntx); +void BitField(CmdArgList args, ConnectionContext* cntx); +void BitFieldRo(CmdArgList args, ConnectionContext* cntx); +void BitOp(CmdArgList args, ConnectionContext* cntx); +void GetBit(CmdArgList args, ConnectionContext* cntx); +void SetBit(CmdArgList args, ConnectionContext* cntx); + +OpResult ReadValue(const OpArgs& op_args, std::string_view key); +OpResult ReadValueBitsetAt(const OpArgs& op_args, std::string_view key, uint32_t offset); +std::string GetString(EngineShard* shard, const PrimeValue& pv); +bool SetBitValue(uint32_t offset, bool bit_value, std::string* entry); + +// ------------------------------------------------------------------------- // +// Bits manipulation functions +constexpr int32_t GetBitIndex(uint32_t offset) noexcept { + return offset % OFFSET_FACTOR; +} + +constexpr int32_t GetNormalizedBitIndex(uint32_t offset) noexcept { + return (OFFSET_FACTOR - 1) - GetBitIndex(offset); +} + +constexpr int32_t GetByteIndex(uint32_t offset) noexcept { + return offset / OFFSET_FACTOR; +} + +uint8_t GetByteValue(const std::string& str, uint32_t offset) { + return static_cast(str[GetByteIndex(offset)]); +} + +constexpr bool CheckBitStatus(uint8_t byte, uint32_t offset) { + return byte & (0x1 << offset); +} + +// return true if bit is on +bool GetBitValue(const std::string& entry, uint32_t offset) { + const auto byte_val{GetByteValue(entry, offset)}; + const auto index{GetNormalizedBitIndex(offset)}; + return CheckBitStatus(byte_val, index); +} + +bool GetBitValueSafe(const std::string& entry, uint32_t offset) { + return ((entry.size() * OFFSET_FACTOR) > offset) ? GetBitValue(entry, offset) : false; +} + +constexpr uint8_t TurnBitOn(uint8_t on, uint32_t offset) { + return on |= 1 << offset; +} + +constexpr uint8_t TunBitOff(uint8_t on, uint32_t offset) { + return on &= ~(1 << offset); +} + +bool SetBitValue(uint32_t offset, bool bit_value, std::string* entry) { + // we need to return the old value after setting the value for offset + const auto old_value{GetBitValue(*entry, offset)}; // save this as the return value + auto byte{GetByteValue(*entry, offset)}; + std::bitset<8> bits{byte}; + const auto bit_index{GetNormalizedBitIndex(offset)}; + byte = bit_value ? TurnBitOn(byte, bit_index) : TunBitOff(byte, bit_index); + (*entry)[GetByteIndex(offset)] = byte; + return old_value; +} + +// ------------------------------------------------------------------------- // +// Helper functions to access the data or change it + +class OverrideValue { + EngineShard* shard_ = nullptr; + DbIndex index_ = 0; + + public: + explicit OverrideValue(const OpArgs& args) : shard_{args.shard}, index_{args.db_ind} { + } + + OpResult Set(std::string_view key, uint32_t offset, bool bit_value); +}; + +OpResult OverrideValue::Set(std::string_view key, uint32_t offset, bool bit_value) { + auto& db_slice = shard_->db_slice(); + DbIndex index = index_; + + DCHECK(db_slice.IsDbValid(index_)); + + std::pair add_res; + try { + add_res = db_slice.AddOrFind(index_, key); + } 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(index, it); + it->second.SetString(value); + db_slice.PostUpdate(index, it, key, !added); + }; + + if (added) { // this is a new entry in the "table" + std::string new_entry(GetByteIndex(offset) + 1, 0); + old_value = SetBitValue(offset, bit_value, &new_entry); + UpdateBitMapValue(new_entry); + } else { + if (it->second.ObjType() != OBJ_STRING) { + return OpStatus::WRONG_TYPE; + } + bool reset = false; + std::string existing_entry{GetString(shard_, it->second)}; + if ((existing_entry.size() * OFFSET_FACTOR) <= offset) { // need to resize first + 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); + } + } + return old_value; +} + +// ------------------------------------------------------------------------- // +// Impl for the command functions +void BitPos(CmdArgList args, ConnectionContext* cntx) { + (*cntx)->SendLong(0); +} + +void BitCount(CmdArgList args, ConnectionContext* cntx) { + (*cntx)->SendLong(0); +} + +void BitField(CmdArgList args, ConnectionContext* cntx) { + (*cntx)->SendLong(0); +} + +void BitFieldRo(CmdArgList args, ConnectionContext* cntx) { + (*cntx)->SendLong(0); +} + +void BitOp(CmdArgList args, ConnectionContext* cntx) { + (*cntx)->SendOk(); +} + +void GetBit(CmdArgList args, ConnectionContext* cntx) { + // Support for the command "GETBIT key offset" + // see https://redis.io/commands/getbit/ + + uint32_t offset{0}; + std::string_view key = ArgS(args, 1); + + if (!absl::SimpleAtoi(ArgS(args, 2), &offset)) { + return (*cntx)->SendError(kInvalidIntErr); + } + auto cb = [&](Transaction* t, EngineShard* shard) { + return ReadValueBitsetAt(t->GetOpArgs(shard), key, offset); + }; + Transaction* trans = cntx->transaction; + OpResult 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 + } + } +} + +void SetBit(CmdArgList args, ConnectionContext* cntx) { + // Support for the command "SETBIT key offset new_value" + // see https://redis.io/commands/setbit/ + + uint32_t offset{0}; + int32_t value{0}; + std::string_view key = ArgS(args, 1); + + if (!absl::SimpleAtoi(ArgS(args, 2), &offset) || !absl::SimpleAtoi(ArgS(args, 3), &value)) { + return (*cntx)->SendError(kInvalidIntErr); + } + + auto cb = [&](Transaction* t, EngineShard* shard) { + OverrideValue set_operation{t->GetOpArgs(shard)}; + + return set_operation.Set(key, offset, value != 0); + }; + + Transaction* trans = cntx->transaction; + OpResult 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; + } + } +} + +// ------------------------------------------------------------------------- // +// This are the "callbacks" that we're using from above +std::string GetString(EngineShard* shard, const PrimeValue& pv) { + std::string res; + if (pv.IsExternal()) { + auto* tiered = shard->tiered_storage(); + auto [offset, size] = pv.GetExternalPtr(); + res.resize(size); + + std::error_code ec = tiered->Read(offset, size, res.data()); + CHECK(!ec) << "TBD: " << ec; + } else { + pv.GetString(&res); + } + + return res; +} + +OpResult ReadValueBitsetAt(const OpArgs& op_args, std::string_view key, uint32_t offset) { + OpResult result = ReadValue(op_args, key); + if (result) { + return GetBitValueSafe(result.value(), offset); + } else { + return result.status(); + } +} + +OpResult ReadValue(const OpArgs& op_args, std::string_view key) { + OpResult it_res = op_args.shard->db_slice().Find(op_args.db_ind, key, OBJ_STRING); + if (!it_res.ok()) { + return it_res.status(); + } + + const PrimeValue& pv = it_res.value()->second; + + return GetString(op_args.shard, pv); +} + +} // namespace + +void BitOpsFamily::Register(CommandRegistry* registry) { + using CI = CommandId; + + *registry << CI{"BITPOS", CO::CommandOpt::READONLY, -3, 1, 1, 1}.SetHandler(&BitPos) + << 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{"GETBIT", CO::READONLY | CO::FAST | CO::FAST, 3, 1, 1, 1}.SetHandler(&GetBit) + << CI{"SETBIT", CO::WRITE, 4, 1, 1, 1}.SetHandler(&SetBit); +} + +} // namespace dfly diff --git a/src/server/bitops_family.h b/src/server/bitops_family.h new file mode 100644 index 000000000..fe368eae1 --- /dev/null +++ b/src/server/bitops_family.h @@ -0,0 +1,30 @@ +// Copyright 2021, Roman Gershman. All rights reserved. +// See LICENSE for licensing terms. +// + +#pragma once + +/// @brief This would implement bits related string commands: GETBIT, SETBIT, BITCOUNT, BITOP. +/// Note: The name of this derive from the same file name in Redis source code. +/// For more details about these command see: +/// BITPOS: https://redis.io/commands/bitpos/ +/// BITCOUNT: https://redis.io/commands/bitcount/ +/// BITFIELD: https://redis.io/commands/bitfield/ +/// BITFIELD_RO: https://redis.io/commands/bitfield_ro/ +/// BITOP: https://redis.io/commands/bitop/ +/// GETBIT: https://redis.io/commands/getbit/ +/// SETBIT: https://redis.io/commands/setbit/ +namespace dfly { +class CommandRegistry; + +class BitOpsFamily { + public: + /// @brief Register the function that would be called to operate on user commands. + /// @param registry The location to which the handling functions would be registered. + /// + /// We are assuming that this would have a valid registry to work on (i.e this do not point to + /// null!). + static void Register(CommandRegistry* registry); +}; + +} // end of namespace dfly diff --git a/src/server/bitops_family_test.cc b/src/server/bitops_family_test.cc new file mode 100644 index 000000000..9235fc193 --- /dev/null +++ b/src/server/bitops_family_test.cc @@ -0,0 +1,79 @@ +// Copyright 2022, Roman Gershman. All rights reserved. +// See LICENSE for licensing terms. +// + +#include "server/bitops_family.h" + +#include "base/gtest.h" +#include "base/logging.h" +#include "facade/facade_test.h" +#include "server/command_registry.h" +#include "server/conn_context.h" +#include "server/engine_shard_set.h" +#include "server/error.h" +#include "server/test_utils.h" +#include "server/transaction.h" +#include "util/uring/uring_pool.h" + +using namespace testing; +using namespace std; +using namespace util; +using absl::StrCat; + +namespace dfly { + +class BitOpsFamilyTest : public BaseFamilyTest { + protected: +}; + +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]); + +TEST_F(BitOpsFamilyTest, GetBit) { + auto resp = Run({"set", "foo", "abc"}); + + EXPECT_EQ(resp, "OK"); + + for (int32_t i = 0; i < ITERATIONS; i++) { + EXPECT_EQ(EXPECTED_VALUE_SETBIT[i], CheckedInt({"getbit", "foo", std::to_string(i)})); + } + + // make sure that when accessing bit that is not in the range its working and we are + // getting 0 + EXPECT_EQ(0, CheckedInt({"getbit", "foo", std::to_string(strlen("abc") + 5)})); +} + +TEST_F(BitOpsFamilyTest, SetBitExistingKey) { + // this test would test when we have the value in place and + // we are overriding and existing key + // so there are no allocations of keys + auto resp = Run({"set", "foo", "abc"}); + + EXPECT_EQ(resp, "OK"); + + // we are setting all to 1s first, we are expecting to get the old values + for (int32_t i = 0; i < ITERATIONS; i++) { + EXPECT_EQ(EXPECTED_VALUE_SETBIT[i], CheckedInt({"setbit", "foo", std::to_string(i), "1"})); + } + + for (int32_t i = 0; i < ITERATIONS; i++) { + EXPECT_EQ(1, CheckedInt({"getbit", "foo", std::to_string(i)})); + } +} + +TEST_F(BitOpsFamilyTest, SetBitMissingKey) { + // This test would run without pre-allocated existing key + // so we need to allocate the key as part of setting the values + for (int32_t i = 0; i < ITERATIONS; i++) { // we are setting all to 1s first, we are expecting + // 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)})); + } +} + +} // end of namespace dfly diff --git a/src/server/main_service.cc b/src/server/main_service.cc index ef460a5f9..63dbb30de 100644 --- a/src/server/main_service.cc +++ b/src/server/main_service.cc @@ -20,6 +20,8 @@ extern "C" { #include "base/logging.h" #include "facade/dragonfly_connection.h" #include "facade/error.h" +#include "server/bitops_family.h" +#include "server/conn_context.h" #include "server/error.h" #include "server/generic_family.h" #include "server/hset_family.h" @@ -548,7 +550,7 @@ void Service::DispatchCommand(CmdArgList args, facade::ConnectionContext* cntx) try { cid->Invoke(args, dfly_cntx); - } catch(std::exception& e) { + } catch (std::exception& e) { LOG(ERROR) << "Internal error, system probably unstable " << e.what(); dfly_cntx->reply_builder()->SendError("Internal Error"); dfly_cntx->reply_builder()->CloseConnection(); @@ -1285,6 +1287,7 @@ void Service::RegisterCommands() { HSetFamily::Register(®istry_); ZSetFamily::Register(®istry_); JsonFamily::Register(®istry_); + BitOpsFamily::Register(®istry_); server_family_.Register(®istry_); diff --git a/src/server/string_family.h b/src/server/string_family.h index 6fc56629c..56c6f5141 100644 --- a/src/server/string_family.h +++ b/src/server/string_family.h @@ -26,7 +26,7 @@ class SetCmd { struct SetParams { SetHow how = SET_ALWAYS; - DbIndex db_index; + DbIndex db_index = 0; uint32_t memcache_flags = 0; // Relative value based on now. 0 means no expiration.