Skip to content

chore(core): Refactor SmallString #6138

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
dranikpg merged 4 commits into from
Dec 8, 2025
+72 −134
Merged

chore(core): Refactor SmallString #6138

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
a853ed1
chore(core): Refactor SmallString
dranikpg Nov 29, 2025
178aa39
more fixes
dranikpg Nov 30, 2025
a524391
update comment
dranikpg Dec 7, 2025
780b1ce
check good size
dranikpg Dec 8, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
77 changes: 28 additions & 49 deletions src/core/compact_object.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -380,7 +380,6 @@ struct TL {

thread_local TL tl;

constexpr bool kUseSmallStrings = true;
constexpr bool kUseAsciiEncoding = true;

} // namespace
Expand Down Expand Up @@ -1145,40 +1144,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;
}

Expand All @@ -1190,15 +1181,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_);
}
Expand Down Expand Up @@ -1258,7 +1242,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 {
Expand Down Expand Up @@ -1471,8 +1455,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];
Expand Down Expand Up @@ -1581,18 +1564,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()) {
Comment on lines -1591 to -1595
Copy link
Contributor Author

@dranikpg dranikpg Nov 30, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

this was the part with 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_);
Expand Down
100 changes: 35 additions & 65 deletions src/core/small_string.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand All @@ -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;
}

Expand All @@ -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());
Expand All @@ -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;
Expand Down
29 changes: 9 additions & 20 deletions src/core/small_string.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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).
Expand Down
Loading