chore(core): Refactor SmallString (#6138)

Author: dranikpg

src/core/compact_object.cc

@@ -380,7 +380,6 @@ struct TL {
 
 thread_local TL tl;
 
-constexpr bool kUseSmallStrings = true;
 constexpr bool kUseAsciiEncoding = true;
 
 }  // namespace
@@ -1148,40 +1147,32 @@ void CompactObj::GetString(char* dest) const {
   }
 
   if (mask_bits_.encoding) {
+    StrEncoding str_encoding = GetStrEncoding();
+    string_view decode_blob;
+
     if (taglen_ == ROBJ_TAG) {
       CHECK_EQ(OBJ_STRING, u_.r_obj.type());
       DCHECK_EQ(OBJ_ENCODING_RAW, u_.r_obj.encoding());
-      string_view blob{(const char*)u_.r_obj.inner_obj(), u_.r_obj.Size()};
-      GetStrEncoding().Decode(blob, dest);
-      return;
+      decode_blob = {(const char*)u_.r_obj.inner_obj(), u_.r_obj.Size()};
     } else {
       CHECK_EQ(SMALL_TAG, taglen_);
-      string_view slices[2];
-      unsigned num = u_.small_str.GetV(slices);
-      DCHECK_EQ(2u, num);
-      size_t decoded_len = GetStrEncoding().DecodedSize(u_.small_str.size(), slices[0][0]);
-
-      if (mask_bits_.encoding == HUFFMAN_ENC) {
-        tl.tmp_buf.resize(slices[0].size() + slices[1].size() - 1);
-        uint8_t* next = tl.tmp_buf.data();
-        memcpy(next, slices[0].data() + 1, slices[0].size() - 1);
-        next += slices[0].size() - 1;
-        memcpy(next, slices[1].data(), slices[1].size());
-        string_view src(reinterpret_cast<const char*>(tl.tmp_buf.data()), tl.tmp_buf.size());
-        const auto& decoder = tl.GetHuffmanDecoder(huffman_domain_);
-        CHECK(decoder.Decode(src, decoded_len, dest));
-        return;
-      }
+      auto& ss = u_.small_str;
 
-      // we left some space on the left to allow inplace ascii unpacking.
-      size_t space_left = decoded_len - u_.small_str.size();
+      char* copy_dest;
+      if (str_encoding.enc_ == HUFFMAN_ENC) {
+        tl.tmp_buf.resize(ss.size());
+        copy_dest = reinterpret_cast<char*>(tl.tmp_buf.data());
+      } else {
+        // Write to rightmost location of dest buffer to leave some bytes for inline unpacking
+        size_t decoded_len = str_encoding.DecodedSize(ss.size(), ss.first_byte());
+        copy_dest = dest + (decoded_len - ss.size());
+      }
 
-      char* next = dest + space_left;
-      memcpy(next, slices[0].data(), slices[0].size());
-      next += slices[0].size();
-      memcpy(next, slices[1].data(), slices[1].size());
-      detail::ascii_unpack_simd(reinterpret_cast<uint8_t*>(dest + space_left), decoded_len, dest);
+      ss.Get(copy_dest);
+      decode_blob = {copy_dest, ss.size()};
     }
+
+    str_encoding.Decode(decode_blob, dest);
     return;
   }
 
@@ -1193,15 +1184,8 @@ void CompactObj::GetString(char* dest) const {
     return;
   }
 
-  if (taglen_ == SMALL_TAG) {
-    string_view slices[2];
-    unsigned num = u_.small_str.GetV(slices);
-    DCHECK_EQ(2u, num);
-    memcpy(dest, slices[0].data(), slices[0].size());
-    dest += slices[0].size();
-    memcpy(dest, slices[1].data(), slices[1].size());
-    return;
-  }
+  if (taglen_ == SMALL_TAG)
+    return u_.small_str.Get(dest);
 
   LOG(FATAL) << "Bad tag " << int(taglen_);
 }
@@ -1267,7 +1251,7 @@ void CompactObj::Materialize(std::string_view blob, bool is_raw) {
   DCHECK_GT(blob.size(), kInlineLen);  // There are no mutable commands that shrink strings
 
   if (is_raw) {
-    if (kUseSmallStrings && SmallString::CanAllocate(blob.size())) {
+    if (SmallString::CanAllocate(blob.size())) {
       SetMeta(SMALL_TAG, mask_);
       tl.small_str_bytes += u_.small_str.Assign(blob);
     } else {
@@ -1481,8 +1465,7 @@ bool CompactObj::CmpEncoded(string_view sv) const {
     DCHECK_GT(sv.size(), 16u);  // we would not be in SMALL_TAG, otherwise.
 
     string_view slice[2];
-    unsigned num = u_.small_str.GetV(slice);
-    DCHECK_EQ(2u, num);
+    u_.small_str.Get(slice);
     DCHECK_LT(slice[0].size(), 14u);
 
     uint8_t tmpbuf[14];
@@ -1591,18 +1574,14 @@ void CompactObj::EncodeString(string_view str, bool is_key) {
 
   DCHECK_GT(encoded.size(), kInlineLen);
 
-  if (kUseSmallStrings && SmallString::CanAllocate(encoded.size())) {
-    if (taglen_ == 0) {
+  if (SmallString::CanAllocate(encoded.size())) {
+    if (taglen_ == SMALL_TAG)
+      tl.small_str_bytes -= u_.small_str.MallocUsed();
+    else
       SetMeta(SMALL_TAG, mask_);
-      tl.small_str_bytes += u_.small_str.Assign(encoded);
-      return;
-    }
 
-    if (taglen_ == SMALL_TAG && encoded.size() <= u_.small_str.size()) {
-      tl.small_str_bytes -= u_.small_str.MallocUsed();
-      tl.small_str_bytes += u_.small_str.Assign(encoded);
-      return;
-    }
+    tl.small_str_bytes += u_.small_str.Assign(encoded);
+    return;
   }
 
   SetMeta(ROBJ_TAG, mask_);

src/core/small_string.cc

@@ -54,55 +54,39 @@ size_t SmallString::UsedThreadLocal() {
 
 static_assert(sizeof(SmallString) == 16);
 
-// we should use only for sizes greater than kPrefLen
 size_t SmallString::Assign(std::string_view s) {
   DCHECK_GT(s.size(), kPrefLen);
-
+  DCHECK(CanAllocate(s.size()));
   uint8_t* realptr = nullptr;
 
-  if (size_ == 0) {
-    // packed structs can not be tied here.
-    auto [sp, rp] = tl.seg_alloc->Allocate(s.size() - kPrefLen);
+  // reallocate if we need a larger allocation or it becomes space-inefficient
+  size_t heap_len = s.size() - kPrefLen;
+  if (size_t available = MallocUsed(); available < heap_len || heap_len * 2 < available) {
+    Free();
+
+    auto [sp, rp] = tl.seg_alloc->Allocate(heap_len);
     small_ptr_ = sp;
     realptr = rp;
-    size_ = s.size();
-  } else if (s.size() <= size_) {
-    realptr = tl.seg_alloc->Translate(small_ptr_);
-
-    if (s.size() < size_) {
-      size_t capacity = mi_usable_size(realptr);
-      if (s.size() * 2 < capacity) {
-        tl.seg_alloc->Free(small_ptr_);
-        auto [sp, rp] = tl.seg_alloc->Allocate(s.size() - kPrefLen);
-        small_ptr_ = sp;
-        realptr = rp;
-      }
-      size_ = s.size();
-    }
   } else {
-    LOG(FATAL) << "TBD: Bad usage";
+    realptr = tl.seg_alloc->Translate(small_ptr_);
   }
 
+  size_ = s.size();
   memcpy(prefix_, s.data(), kPrefLen);
-  memcpy(realptr, s.data() + kPrefLen, s.size() - kPrefLen);
-
+  memcpy(realptr, s.data() + kPrefLen, heap_len);
   return mi_malloc_usable_size(realptr);
 }
 
 void SmallString::Free() {
-  if (size_ <= kPrefLen)
-    return;
-
-  tl.seg_alloc->Free(small_ptr_);
+  if (size_)
+    tl.seg_alloc->Free(small_ptr_);
   size_ = 0;
 }
 
 uint16_t SmallString::MallocUsed() const {
-  if (size_ <= kPrefLen)
-    return 0;
-  auto* realptr = tl.seg_alloc->Translate(small_ptr_);
-
-  return mi_malloc_usable_size(realptr);
+  if (size_)
+    return mi_malloc_usable_size(tl.seg_alloc->Translate(small_ptr_));
+  return 0;
 }
 
 bool SmallString::Equal(std::string_view o) const {
@@ -112,13 +96,10 @@ bool SmallString::Equal(std::string_view o) const {
   if (size_ == 0)
     return true;
 
-  DCHECK_GT(size_, kPrefLen);
-
   if (memcmp(prefix_, o.data(), kPrefLen) != 0)
     return false;
 
   uint8_t* realp = tl.seg_alloc->Translate(small_ptr_);
-
   return memcmp(realp, o.data() + kPrefLen, size_ - kPrefLen) == 0;
 }
 
@@ -127,21 +108,16 @@ bool SmallString::Equal(const SmallString& os) const {
     return false;
 
   string_view me[2], other[2];
-  unsigned n1 = GetV(me);
-  unsigned n2 = os.GetV(other);
-
-  if (n1 != n2)
-    return false;
+  Get(me);
+  os.Get(other);
 
   return me[0] == other[0] && me[1] == other[1];
 }
 
 uint64_t SmallString::HashCode() const {
-  DCHECK_GT(size_, kPrefLen);
-
   string_view slice[2];
+  Get(slice);
 
-  GetV(slice);
   XXH3_state_t* state = tl.xxh_state.get();
   XXH3_64bits_reset_withSeed(state, kHashSeed);
   XXH3_64bits_update(state, slice[0].data(), slice[0].size());
@@ -150,41 +126,35 @@ uint64_t SmallString::HashCode() const {
   return XXH3_64bits_digest(state);
 }
 
-void SmallString::Get(std::string* dest) const {
-  dest->resize(size_);
-  if (size_) {
-    DCHECK_GT(size_, kPrefLen);
-    memcpy(dest->data(), prefix_, kPrefLen);
-    uint8_t* ptr = tl.seg_alloc->Translate(small_ptr_);
-    memcpy(dest->data() + kPrefLen, ptr, size_ - kPrefLen);
-  }
-}
-
-unsigned SmallString::GetV(string_view dest[2]) const {
-  DCHECK_GT(size_, kPrefLen);
-  if (size_ <= kPrefLen) {
-    dest[0] = string_view{prefix_, size_};
-    return 1;
-  }
+void SmallString::Get(string_view dest[2]) const {
+  DCHECK(size_);
 
   dest[0] = string_view{prefix_, kPrefLen};
   uint8_t* ptr = tl.seg_alloc->Translate(small_ptr_);
   dest[1] = string_view{reinterpret_cast<char*>(ptr), size_ - kPrefLen};
-  return 2;
 }
 
-bool SmallString::DefragIfNeeded(PageUsage* page_usage) {
-  DCHECK_GT(size_, kPrefLen);
-  if (size_ <= kPrefLen) {
-    return false;
-  }
+void SmallString::Get(char* out) const {
+  string_view strs[2];
+  Get(strs);
+  memcpy(out, strs[0].data(), strs[0].size());
+  memcpy(out + strs[0].size(), strs[1].data(), strs[1].size());
+}
 
+void SmallString::Get(std::string* dest) const {
+  dest->resize(size_);
+  Get(dest->data());
+}
+
+bool SmallString::DefragIfNeeded(PageUsage* page_usage) {
   uint8_t* cur_real_ptr = tl.seg_alloc->Translate(small_ptr_);
   if (!page_usage->IsPageForObjectUnderUtilized(tl.seg_alloc->heap(), cur_real_ptr))
     return false;
 
-  auto [sp, rp] = tl.seg_alloc->Allocate(size_ - kPrefLen);
+  if (!CanAllocate(size_ - kPrefLen))  // Forced
+    return false;
 
+  auto [sp, rp] = tl.seg_alloc->Allocate(size_ - kPrefLen);
   memcpy(rp, cur_real_ptr, size_ - kPrefLen);
   tl.seg_alloc->Free(small_ptr_);
   small_ptr_ = sp;

src/core/small_string.h

@@ -10,10 +10,8 @@ namespace dfly {
 
 class PageUsage;
 
-// blob strings of upto ~256B. Small sizes are probably predominant
-// for in-memory workloads, especially for keys.
-// Please note that this class does not have automatic constructors and destructors, therefore
-// it requires explicit management.
+// Efficient storage of strings longer than 10 bytes.
+// Requires explicit memory management
 class SmallString {
   static constexpr unsigned kPrefLen = 10;
   static constexpr unsigned kMaxSize = (1 << 8) - 1;
@@ -23,41 +21,32 @@ class SmallString {
   static size_t UsedThreadLocal();
   static bool CanAllocate(size_t size);
 
-  void Reset() {
-    size_ = 0;
-  }
-
   // Returns malloc used.
   size_t Assign(std::string_view s);
   void Free();
 
   bool Equal(std::string_view o) const;
   bool Equal(const SmallString& mps) const;
 
-  uint16_t size() const {
-    return size_;
-  }
-
   uint64_t HashCode() const;
-
-  // I am lying here. we should use mi_malloc_usable size really.
   uint16_t MallocUsed() const;
 
+  void Get(std::string_view dest[2]) const;
+  void Get(char* out) const;
   void Get(std::string* dest) const;
 
-  // returns 1 or 2 slices representing this small string.
-  // Guarantees zero copy, i.e. dest will not point to any of external buffers.
-  // With current implementation, it will return 2 slices for a non-empty string.
-  unsigned GetV(std::string_view dest[2]) const;
-
   bool DefragIfNeeded(PageUsage* page_usage);
 
+  size_t size() const {
+    return size_;
+  }
+
   uint8_t first_byte() const {
     return prefix_[0];
   }
 
  private:
-  // prefix of the string that is broken down into 2 parts.
+  // The string is stored broken up into two parts, the first one - in this array
   char prefix_[kPrefLen];
 
   uint32_t small_ptr_;  // 32GB capacity because we ignore 3 lsb bits (i.e. x8).