Add extent tree for later usage in the external allocator

README.md

@@ -109,7 +109,7 @@ API 1.0
   - [X] LSET
   - [X] LTRIM
   - [X] RPOP
-  - [ ] RPOPLPUSH
+  - [X] RPOPLPUSH
   - [X] RPUSH
 - [X] SortedSet Family
   - [X] ZADD
@@ -214,7 +214,7 @@ API 2.0
   - [X] ZREVRANGEBYSCORE
   - [X] ZREVRANK
   - [ ] ZUNIONSTORE
-  - [ ] ZSCAN
+  - [X] ZSCAN
 - [ ] HYPERLOGLOG Family
   - [ ] PFADD
   - [ ] PFCOUNT

src/core/CMakeLists.txt

@@ -1,10 +1,12 @@
-add_library(dfly_core compact_object.cc dragonfly_core.cc external_alloc.cc interpreter.cc 
+add_library(dfly_core compact_object.cc dragonfly_core.cc extent_tree.cc 
+            external_alloc.cc interpreter.cc 
             segment_allocator.cc small_string.cc tx_queue.cc)
 cxx_link(dfly_core base absl::flat_hash_map absl::str_format redis_lib TRDP::lua 
          Boost::fiber crypto)
 
 cxx_test(dfly_core_test dfly_core LABELS DFLY)
 cxx_test(compact_object_test dfly_core LABELS DFLY)
+cxx_test(extent_tree_test dfly_core LABELS DFLY)
 cxx_test(external_alloc_test dfly_core LABELS DFLY)
 cxx_test(dash_test dfly_core LABELS DFLY)
 cxx_test(interpreter_test dfly_core LABELS DFLY)

src/core/extent_tree.cc

@@ -0,0 +1,112 @@
+// Copyright 2022, Roman Gershman.  All rights reserved.
+// See LICENSE for licensing terms.
+//
+
+#include "core/extent_tree.h"
+
+#include "base/logging.h"
+
+namespace dfly {
+
+using namespace std;
+
+// offset, len must be multiplies of 256MB.
+void ExtentTree::Add(size_t start, size_t len) {
+  DCHECK_GT(len, 0u);
+  DCHECK_EQ(len_extents_.size(), extents_.size());
+
+  size_t end = start + len;
+
+  if (extents_.empty()) {
+    extents_.emplace(start, end);
+    len_extents_.emplace(len, start);
+
+    return;
+  }
+
+  auto it = extents_.lower_bound(start);
+  bool merged = false;
+
+  if (it != extents_.begin()) {
+    auto prev = it;
+    --prev;
+
+    DCHECK_LE(prev->second, start);
+    if (prev->second == start) {  // [first, second = start, end)
+      merged = true;
+      len_extents_.erase(pair{prev->second - prev->first, prev->first});
+
+      if (end == it->first) {  // [first, end = it->first, it->second)
+        prev->second = it->second;
+        len_extents_.erase(pair{it->second - it->first, it->first});
+        extents_.erase(it);
+      } else {
+        prev->second = end;
+      }
+      len_extents_.emplace(prev->second - prev->first, prev->first);
+    }
+  }
+
+  if (!merged) {
+    if (end == it->first) {  // [start, end), [it->first, it->second]
+      len_extents_.erase(pair{it->second - it->first, it->first});
+      end = it->second;
+      extents_.erase(it);
+    }
+    extents_.emplace(start, end);
+    len_extents_.emplace(end - start, start);
+  }
+}
+
+optional<pair<size_t, size_t>> ExtentTree::GetRange(size_t len, size_t align) {
+  DCHECK_GT(align, 0u);
+  DCHECK_EQ(0u, align & (align - 1));
+  DCHECK_EQ(0u, len & (align - 1));
+
+  auto it = len_extents_.lower_bound(pair{len, 0});
+  if (it == len_extents_.end())
+    return nullopt;
+  size_t amask = align - 1;
+  size_t aligned_start = 0;
+
+  while (true) {
+    aligned_start = it->second;
+    if ((aligned_start & amask) == 0)
+      break;
+    size_t end = it->first + aligned_start;
+    // round up to the next aligned address
+    aligned_start = align + (aligned_start & (~amask));
+    if (aligned_start + len <= end)
+      break;
+    ++it;
+
+    if (it == len_extents_.end())
+      return nullopt;
+  }
+
+  DCHECK_GE(aligned_start, it->second);
+  // if we are here - we found the range starting at aligned_start.
+  auto eit = extents_.find(it->second);
+  DCHECK(eit != extents_.end());
+  size_t end = eit->second;
+  size_t range_end = aligned_start + len;
+
+  len_extents_.erase(it);
+
+  // we break the interval [eit->first, eit->second] to either 0, 1 or 2 intervals.
+  if (aligned_start > eit->first) {
+    eit->second = aligned_start;
+    len_extents_.emplace(eit->second - eit->first, eit->first);
+  } else {
+    extents_.erase(eit);
+  }
+
+  if (aligned_start + len < end) {
+    extents_.emplace(range_end, end);
+    len_extents_.emplace(end - range_end, range_end);
+  }
+
+  return pair{aligned_start, range_end};
+}
+
+}  // namespace dfly

src/core/extent_tree.h

@@ -0,0 +1,30 @@
+// Copyright 2022, Roman Gershman.  All rights reserved.
+// See LICENSE for licensing terms.
+//
+
+#pragma once
+
+#include <absl/container/btree_map.h>
+#include <absl/container/btree_set.h>
+
+#include <optional>
+
+namespace dfly {
+
+// represents a tree of disjoint extents.
+// check-fails if overlapping ranges are added.
+// automatically handles union of the consequent ranges that are added to the tree.
+class ExtentTree {
+ public:
+  void Add(size_t start, size_t len);
+
+  // in case of success, returns (start, end) pair  where end-start >= len and
+  // start is aligned by align.
+  std::optional<std::pair<size_t, size_t>> GetRange(size_t len, size_t align);
+
+ private:
+  absl::btree_map<size_t, size_t> extents_;                 // start -> end.
+  absl::btree_set<std::pair<size_t, size_t>> len_extents_;  // (length, start)
+};
+
+}  // namespace dfly

src/core/extent_tree_test.cc

@@ -0,0 +1,58 @@
+// Copyright 2022, Roman Gershman.  All rights reserved.
+// See LICENSE for licensing terms.
+//
+
+#include "core/extent_tree.h"
+
+#include <gmock/gmock.h>
+
+#include "base/gtest.h"
+#include "base/logging.h"
+
+namespace dfly {
+
+using namespace std;
+
+class ExtentTreeTest : public ::testing::Test {
+ protected:
+  static void SetUpTestSuite() {
+  }
+
+  static void TearDownTestSuite() {
+  }
+
+  ExtentTree tree_;
+};
+
+TEST_F(ExtentTreeTest, Basic) {
+  tree_.Add(0, 256);
+  auto op = tree_.GetRange(64, 16);
+  EXPECT_TRUE(op);
+  EXPECT_THAT(*op, testing::Pair(0, 64));
+
+  tree_.Add(56, 8);
+  op = tree_.GetRange(64, 16);
+  EXPECT_TRUE(op);
+  EXPECT_THAT(*op, testing::Pair(64, 128));
+
+  op = tree_.GetRange(18, 2);
+  EXPECT_TRUE(op);
+  EXPECT_THAT(*op, testing::Pair(128, 146));
+
+  op = tree_.GetRange(80, 16);
+  EXPECT_TRUE(op);
+  EXPECT_THAT(*op, testing::Pair(160, 240));
+
+  op = tree_.GetRange(4, 1);
+  EXPECT_TRUE(op);
+  EXPECT_THAT(*op, testing::Pair(56, 60));
+
+  op = tree_.GetRange(32, 1);
+  EXPECT_FALSE(op);
+  tree_.Add(64, 240 - 64);
+  op = tree_.GetRange(32, 4);
+  EXPECT_TRUE(op);
+  EXPECT_THAT(*op, testing::Pair(60, 92));
+}
+
+}  // namespace dfly