json: Adds ref-counted string

The string class wraps the previously added blob class. It decides at
creation time whether to create a new blob, or bump up the refcount of
an existing blob in the pool.

The pool is thread local, and is used by the strings on that thread to
store pointers to allocated blobs.

When a string is destroyed it decrements the ref-count, dropping the
blob if necessary (ie refcount == 0).

A lot of boilerplate in this class is from trial and error, adding
methods that jsoncons required to be able to use this class as a key.

Signed-off-by: Abhijat Malviya <[email protected]>

Author: abhijat

src/core/json/CMakeLists.txt

@@ -6,7 +6,8 @@ cur_gen_dir(gen_dir)
 add_library(jsonpath lexer_impl.cc driver.cc path.cc
             ${gen_dir}/jsonpath_lexer.cc ${gen_dir}/jsonpath_grammar.cc json_object.cc
             detail/jsoncons_dfs.cc detail/flat_dfs.cc
-            detail/interned_blob.cc)
+            detail/interned_blob.cc
+            detail/interned_string.cc)
 target_link_libraries(jsonpath base absl::strings TRDP::reflex TRDP::jsoncons TRDP::flatbuffers dfly_page_usage)
 
 helio_cxx_test(jsonpath_test jsonpath LABELS DFLY)

src/core/json/detail/interned_string.cc

@@ -0,0 +1,85 @@
+// Copyright 2025, DragonflyDB authors.  All rights reserved.
+// See LICENSE for licensing terms.
+
+#include "core/json/detail/interned_string.h"
+
+namespace dfly::detail {
+
+InternedString& InternedString::operator=(const InternedString& other) {
+  if (this != &other) {
+    Release();
+    entry_ = other.entry_;
+    Acquire();
+  }
+  return *this;
+}
+
+InternedString& InternedString::operator=(InternedString&& other) noexcept {
+  if (this != &other) {
+    Release();
+    entry_ = other.entry_;
+    other.entry_ = {};
+  }
+  return *this;
+}
+
+int InternedString::compare(const InternedString& other) const {
+  return std::string_view{*this}.compare(other);
+}
+
+int InternedString::compare(std::string_view other) const {
+  return std::string_view{*this}.compare(other);
+}
+
+void InternedString::ResetPool() {
+  InternedBlobPool& pool = GetPoolRef();
+  for (InternedBlobHandle handle : pool) {
+    InternedBlobHandle::Destroy(handle);
+  }
+  pool.clear();
+}
+
+InternedBlobHandle InternedString::Intern(const std::string_view sv) {
+  if (sv.empty())
+    return {};
+
+  InternedBlobPool& pool_ref = GetPoolRef();
+  if (const auto it = pool_ref.find(sv); it != pool_ref.end()) {
+    InternedBlobHandle blob = *it;
+    blob.IncrRefCount();
+    return blob;
+  }
+
+  InternedBlobHandle handle = InternedBlobHandle::Create(sv);
+  pool_ref.emplace(handle);
+  return handle;
+}
+
+void InternedString::Acquire() {  // NOLINT
+  if (entry_) {
+    entry_.IncrRefCount();
+  }
+}
+
+void InternedString::Release() {
+  if (!entry_)
+    return;
+
+  entry_.DecrRefCount();
+
+  if (entry_.RefCount() == 0) {
+    GetPoolRef().erase(entry_);
+    InternedBlobHandle::Destroy(entry_);
+  }
+}
+
+InternedBlobPool& InternedString::GetPoolRef() {
+  // Note on lifetimes: this pool is thread local and depends on the thread local memory resource
+  // defined in the stateless allocator in src/core/detail/stateless_allocator.h. Since there is no
+  // well-defined order of destruction, this pool must be manually reset before the memory resource
+  // destruction.
+  thread_local InternedBlobPool pool;
+  return pool;
+}
+
+}  // namespace dfly::detail

src/core/json/detail/interned_string.h

@@ -0,0 +1,127 @@
+// Copyright 2025, DragonflyDB authors.  All rights reserved.
+// See LICENSE for licensing terms.
+
+#pragma once
+
+#include "core/detail/stateless_allocator.h"
+#include "core/json/detail/interned_blob.h"
+
+namespace dfly::detail {
+
+// InternedString handles incrementing and decrementing reference counts of the blobs tied to its
+// own lifecycle. It deletes the blob from a shard local pool when refcount is 0.
+// TODO examine cross shard json object interactions. Can a pool end up access from another shard?
+class InternedString {
+ public:
+  using allocator_type = StatelessAllocator<char>;
+
+  InternedString() = default;
+
+  explicit InternedString(const std::string_view sv) : entry_(Intern(sv)) {
+  }
+
+  // The following constructors and members are added because they are required by jsoncons for
+  // keys. Each of these is added in response to compiler errors and should not be removed, even if
+  // they are seemingly a no-op or duplicated.
+
+  // jsoncons sometimes creates empty obj with custom allocator. If it creates an object with any
+  // other allocator, we should fail during compilation.
+  template <typename T> explicit InternedString(StatelessAllocator<T> /*unused*/) {
+  }
+
+  template <typename Alloc> InternedString(const char* data, size_t size, Alloc alloc);
+
+  template <std::contiguous_iterator It> InternedString(It begin, It end);
+
+  InternedString(const InternedString& other) : entry_{other.entry_} {
+    Acquire();
+  }
+
+  InternedString(InternedString&& other) noexcept : entry_{other.entry_} {
+    other.entry_ = {};
+  }
+
+  InternedString& operator=(const InternedString& other);
+  InternedString& operator=(InternedString&& other) noexcept;
+
+  ~InternedString() {
+    Release();
+  }
+
+  operator std::string_view() const {
+    return entry_;
+  }
+
+  [[nodiscard]] const char* data() const {
+    return entry_ ? entry_.Data() : "";
+  }
+
+  [[nodiscard]] const char* c_str() const {
+    return data();
+  }
+
+  void swap(InternedString& other) noexcept {
+    std::swap(entry_, other.entry_);
+  }
+
+  [[nodiscard]] size_t length() const {
+    return size();
+  }
+
+  [[nodiscard]] size_t size() const {
+    return entry_.Size();
+  }
+
+  [[nodiscard]] int compare(const InternedString& other) const;
+  [[nodiscard]] int compare(std::string_view) const;
+
+  // lex. comparison
+  auto operator<=>(const InternedString& other) const {
+    return std::string_view{*this} <=> std::string_view{other};
+  }
+
+  bool operator==(const InternedString& other) const = default;
+
+  void shrink_to_fit() {  // NOLINT (must be non-const to align with jsoncons usage)
+  }
+
+  // Destroys all strings in the pool. Must be called on process shutdown before the backing memory
+  // resource is destroyed.
+  static void ResetPool();
+  static InternedBlobPool& GetPoolRef();
+
+  [[nodiscard]] size_t MemUsed() const {
+    return entry_ ? entry_.MemUsed() : 0;
+  }
+
+ private:
+  // If a string exists in the pool, increments its refcount. If not, adds the string to the pool.
+  // Returns a handle wrapping the string.
+  static InternedBlobHandle Intern(std::string_view sv);
+
+  // Increments the refcount if the entry is not null
+  void Acquire();
+
+  // Decrements the refcount, removes entry from the pool if necessary, destroying the interned blob
+  void Release();
+
+  // Wraps a null pointer by default
+  InternedBlobHandle entry_;
+};
+
+template <typename Alloc>
+InternedString::InternedString(const char* data, size_t size, Alloc /*unused*/)
+    : InternedString(std::string_view{data, size}) {
+}
+
+template <std::contiguous_iterator It> InternedString::InternedString(It begin, It end) {
+  if (begin == end) {
+    return;
+  }
+
+  const auto size = std::distance(begin, end);
+  const auto data_ptr = &*begin;
+  entry_ = Intern(std::string_view(data_ptr, size));
+}
+
+}  // namespace dfly::detail

src/core/json/interned_blob_test.cc

@@ -1,10 +1,9 @@
 // Copyright 2025, DragonflyDB authors.  All rights reserved.
 // See LICENSE for licensing terms.
 
-#include "core/json/detail/interned_blob.h"
-
 #include "base/gtest.h"
 #include "core/detail/stateless_allocator.h"
+#include "core/json/detail/interned_string.h"
 #include "core/mi_memory_resource.h"
 
 using namespace std::literals;
@@ -99,3 +98,135 @@ TEST_F(InternedBlobTest, Pool) {
   EXPECT_TRUE(pool.contains("foo"));
   InternedBlobHandle::Destroy(b1);
 }
+
+using detail::InternedString;
+
+namespace {
+
+void StringCheck(const InternedString& s, const char* ptr) {
+  std::string_view sv{ptr};
+
+  EXPECT_STREQ(s.data(), ptr);
+  EXPECT_STREQ(s.c_str(), ptr);
+
+  EXPECT_EQ(s.size(), sv.size());
+  EXPECT_EQ(s.length(), sv.size());
+
+  EXPECT_EQ(std::string_view(s), sv);
+  EXPECT_EQ(std::string_view(s.data(), s.size()), sv);
+  EXPECT_EQ(std::string_view(s.c_str(), s.size()), sv);
+}
+
+}  // namespace
+
+TEST_F(InternedBlobTest, StringPool) {
+  const auto& pool = InternedString::GetPoolRef();
+  EXPECT_TRUE(pool.empty());
+  {
+    const InternedString s1{"foobar"};
+    StringCheck(s1, "foobar");
+    EXPECT_EQ(pool.size(), 1);
+    {
+      const InternedString s2{"foobar"};
+      StringCheck(s2, "foobar");
+      EXPECT_EQ(pool.size(), 1);
+    }
+    EXPECT_EQ(pool.size(), 1);
+  }
+  EXPECT_TRUE(pool.empty());
+
+  std::vector<InternedString> strings;
+  for (auto i = 0; i < 1000; ++i) {
+    strings.emplace_back(std::to_string(i));
+  }
+
+  EXPECT_EQ(pool.size(), 1000);
+  strings.clear();
+  EXPECT_TRUE(pool.empty());
+
+  for (auto i = 0; i < 1000; ++i) {
+    strings.emplace_back("zyx");
+  }
+  EXPECT_EQ(pool.size(), 1);
+  strings.clear();
+  EXPECT_TRUE(pool.empty());
+
+  InternedString empty;
+  EXPECT_TRUE(pool.empty());
+}
+
+TEST_F(InternedBlobTest, StringApi) {
+  InternedString s1{"foobar"};
+  EXPECT_EQ(std::string_view{s1}, "foobar"sv);
+  StringCheck(s1, "foobar");
+
+  const auto& pool = InternedString::GetPoolRef();
+  InternedString s2{"psi"};
+  StringCheck(s2, "psi");
+
+  EXPECT_EQ(pool.size(), 2);
+
+  // swap pointers into the pool
+  s1.swap(s2);
+
+  EXPECT_EQ(pool.size(), 2);
+
+  StringCheck(s1, "psi");
+  StringCheck(s2, "foobar");
+
+  EXPECT_NE(s1, s2);
+  EXPECT_EQ(s1, s1);
+  // foobar < psi lexicographically
+  EXPECT_LT(s2, s1);
+}
+
+TEST_F(InternedBlobTest, StringCtors) {
+  const auto& pool = InternedString::GetPoolRef();
+  InternedString s1{"foobar"};
+  EXPECT_EQ(pool.size(), 1);
+
+  // move ctor
+  auto to = std::move(s1);
+  EXPECT_EQ(pool.size(), 1);
+
+  StringCheck(to, "foobar");
+  StringCheck(s1, "");
+
+  // These tests exercise self-move and self-copy behavior. This causes errors on newer GCC when
+  // warnings are treated as errors (on CI). We need to version gate this because on older GCC this
+  // check is not present.
+#if defined(__GNUC__) && !defined(__clang__) && __GNUC__ >= 13
+#pragma GCC diagnostic push
+#pragma GCC diagnostic ignored "-Wself-move"
+#endif
+  to = std::move(to);
+  StringCheck(to, "foobar");
+
+  auto copied = to;
+  EXPECT_EQ(pool.size(), 1);
+
+  StringCheck(to, "foobar");
+  StringCheck(copied, "foobar");
+
+  copied = copied;
+#if defined(__GNUC__) && !defined(__clang__) && __GNUC__ >= 13
+#pragma GCC diagnostic pop
+#endif
+  StringCheck(copied, "foobar");
+  EXPECT_EQ(pool.size(), 1);
+
+  const auto* mr = MemoryResource();
+  const auto before = mr->used();
+
+  std::string_view sv{"......."};
+  // ptr and size with some allocator, allocator will be ignored
+  InternedString x{sv.data(), sv.size(), std::allocator<char>{}};
+  StringCheck(x, ".......");
+  EXPECT_EQ(pool.size(), 2);
+
+  EXPECT_GE(mr->used(), before + x.MemUsed());
+
+  InternedString k{sv.begin(), sv.end()};
+  StringCheck(k, ".......");
+  EXPECT_EQ(pool.size(), 2);
+}