fix: stream memory tracking (#4067)

* add object memory usage for streams
* add test
This commit is contained in:
Kostas Kyrimis 2024-11-27 11:41:08 +01:00 committed by GitHub
parent 065a63cab7
commit 66e0fd0908
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 91 additions and 16 deletions

View file

@ -125,11 +125,6 @@ size_t MallocUsedZSet(unsigned encoding, void* ptr) {
return 0;
}
size_t MallocUsedStream(unsigned encoding, void* streamv) {
// stream* str_obj = (stream*)streamv;
return 0; // TODO
}
inline void FreeObjHash(unsigned encoding, void* ptr) {
switch (encoding) {
case kEncodingStrMap2:
@ -316,7 +311,7 @@ size_t RobjWrapper::MallocUsed(bool slow) const {
case OBJ_ZSET:
return MallocUsedZSet(encoding_, inner_obj_);
case OBJ_STREAM:
return MallocUsedStream(encoding_, inner_obj_);
return sz_;
default:
LOG(FATAL) << "Not supported " << type_;
@ -370,7 +365,12 @@ size_t RobjWrapper::Size() const {
StringMap* sm = (StringMap*)inner_obj_;
return sm->UpperBoundSize();
}
default:
LOG(FATAL) << "Unexpected encoding " << encoding_;
}
case OBJ_STREAM:
// Size mean malloc bytes for streams
return sz_;
default:;
}
return 0;
@ -461,6 +461,10 @@ void RobjWrapper::SetString(string_view s, MemoryResource* mr) {
}
}
void RobjWrapper::SetSize(uint64_t size) {
sz_ = size;
}
bool RobjWrapper::DefragIfNeeded(float ratio) {
auto do_defrag = [this, ratio](auto defrag_fun) mutable {
auto [new_ptr, realloced] = defrag_fun(encoding_, inner_obj_, ratio);
@ -811,6 +815,17 @@ void CompactObj::SetJsonSize(int64_t size) {
}
}
void CompactObj::AddStreamSize(int64_t size) {
if (size < 0) {
// We might have a negative size. For example, if we remove a consumer,
// the tracker will report a negative net (since we deallocated),
// so the object now consumes less memory than it did before. This DCHECK
// is for fanity and to catch any potential issues with our tracking approach.
DCHECK(static_cast<int64_t>(u_.r_obj.Size()) >= size);
}
u_.r_obj.SetSize((u_.r_obj.Size() + size));
}
void CompactObj::SetJson(const uint8_t* buf, size_t len) {
SetMeta(JSON_TAG);
u_.json_obj.flat.flat_ptr = (uint8_t*)tl.local_mr->allocate(len, kAlignSize);