fix(server): small string allocations only under 256 bytes str (#2991)

Signed-off-by: adi_holden <[email protected]>

Author: adiholden

src/core/compact_object.cc

@@ -790,8 +790,8 @@ void CompactObj::SetString(std::string_view str) {
     }
   }
 
-  if (kUseSmallStrings) {
-    if ((taglen_ == 0 && encoded.size() < (1 << 13))) {
+  if (kUseSmallStrings && SmallString::CanAllocate(encoded.size())) {
+    if (taglen_ == 0) {
       SetMeta(SMALL_TAG, mask);
       tl.small_str_bytes += u_.small_str.Assign(encoded);
       return;

src/core/segment_allocator.cc

@@ -25,4 +25,8 @@ void SegmentAllocator::ValidateMapSize() {
   }
 }
 
+bool SegmentAllocator::CanAllocate() {
+  return address_table_.size() < (1u << kSegmentIdBits);
+}
+
 }  // namespace dfly

src/core/segment_allocator.h

@@ -30,6 +30,7 @@ class SegmentAllocator {
   using Ptr = uint32_t;
 
   SegmentAllocator(mi_heap_t* heap);
+  bool CanAllocate();
 
   uint8_t* Translate(Ptr p) const {
     return address_table_[p & kSegmentIdMask] + Offset(p);

src/core/small_string.cc

@@ -45,6 +45,10 @@ void SmallString::InitThreadLocal(void* heap) {
   XXH3_64bits_reset_withSeed(tl.xxh_state.get(), kHashSeed);
 }
 
+bool SmallString::CanAllocate(size_t size) {
+  return size <= kMaxSize && tl.seg_alloc->CanAllocate();
+}
+
 size_t SmallString::UsedThreadLocal() {
   return tl.seg_alloc ? tl.seg_alloc->used() : 0;
 }

src/core/small_string.h

@@ -10,16 +10,18 @@
 
 namespace dfly {
 
-// blob strings of upto ~64KB. Small sizes are probably predominant
+// 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.
 class SmallString {
   static constexpr unsigned kPrefLen = 10;
+  static constexpr unsigned kMaxSize = (1 << 8) - 1;
 
  public:
   static void InitThreadLocal(void* heap);
   static size_t UsedThreadLocal();
+  static bool CanAllocate(size_t size);
 
   void Reset() {
     size_ = 0;