feat(server): adding support for bitops #213 (#295)

feat(server): bitset bitget commands - update docs/api_status.md #213

Signed-off-by: Boaz Sade <boaz@dragonflydb.io>

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

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)

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 <bitset>
+
+extern "C" {
+#include "redis/object.h"
+}
+
+#include <new>
+
+#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<std::string> ReadValue(const OpArgs& op_args, std::string_view key);
+OpResult<bool> 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<uint8_t>(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<bool> Set(std::string_view key, uint32_t offset, bool bit_value);
+};
+
+OpResult<bool> 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<PrimeIterator, bool> 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<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
+    }
+  }
+}
+
+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<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;
+    }
+  }
+}
+
+// ------------------------------------------------------------------------- //
+// 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<bool> ReadValueBitsetAt(const OpArgs& op_args, std::string_view key, uint32_t offset) {
+  OpResult<std::string> result = ReadValue(op_args, key);
+  if (result) {
+    return GetBitValueSafe(result.value(), offset);
+  } else {
+    return result.status();
+  }
+}
+
+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_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

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

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

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(&registry_);
   ZSetFamily::Register(&registry_);
   JsonFamily::Register(&registry_);
+  BitOpsFamily::Register(&registry_);
 
   server_family_.Register(&registry_);
 

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.