Ver código fonte

fix(memory): budget version history + vectors, tighten estimateDocumentSize

estimateDocumentSize() was severely under-counting:
- Ignored JSON tree overhead (now uses 2.5x serialized size)
- Ignored Document metadata strings (updatedBy, nodeId, encryptedFields)

Version history and vector embeddings were completely off-books —
estimatedMemoryBytes_ never reflected them, so eviction didn't trigger
aggressively enough and real RSS grew far beyond the configured budget.

Changes:
- estimateDocumentSize: 2.5x JSON size + metadata string accounting
- New estimateVectorSize helper — tracks float arrays
- New estimateDocumentVersionSize helper — tracks deque entries
- storeVector/removeVector update estimatedMemoryBytes_
- saveToHistory updates counter on append and on maxVersions prune
- alterCollection history prune now decrements the counter
- dropCollection now frees history + vector memory alongside docs
- loadVersionHistory accounts for snapshot-restored entries
- clear() resets the counter (it was leaking between mode switches)

No behavior change for deployments where maxVersions is set sanely —
the counter is just honest now.
fszontagh 3 meses atrás
pai
commit
8f378fe75f
2 arquivos alterados com 111 adições e 4 exclusões
  1. 105 4
      service/src/memory_store.cpp
  2. 6 0
      service/src/memory_store.hpp

+ 105 - 4
service/src/memory_store.cpp

@@ -110,10 +110,18 @@ bool MemoryStore::dropCollection(const std::string& name) {
     {
         std::shared_lock<std::shared_mutex> collLock(it->second->mutex);
         docCount = it->second->documents.size();
-        // Calculate total memory being freed
+        // Calculate total memory being freed (documents + version history + vectors)
         for (const auto& [id, doc] : it->second->documents) {
             totalSize += estimateDocumentSize(doc);
         }
+        for (const auto& [id, history] : it->second->versionHistory) {
+            for (const auto& ver : history) {
+                totalSize += estimateDocumentVersionSize(ver);
+            }
+        }
+        for (const auto& [id, vec] : it->second->vectors) {
+            totalSize += estimateVectorSize(vec);
+        }
     }
 
     collections_.erase(it);
@@ -151,12 +159,17 @@ bool MemoryStore::alterCollection(const std::string& name, const CollectionOptio
     // If maxVersions is now limited (and was unlimited or higher before), prune all history
     if (newMaxVersions > 0 && (oldMaxVersions == 0 || newMaxVersions < oldMaxVersions)) {
         uint64_t prunedVersions = 0;
+        uint64_t prunedBytes = 0;
         for (auto& [docId, history] : it->second->versionHistory) {
             while (history.size() > newMaxVersions) {
+                prunedBytes += estimateDocumentVersionSize(history.front());
                 history.pop_front();
                 ++prunedVersions;
             }
         }
+        if (prunedBytes > 0) {
+            estimatedMemoryBytes_.fetch_sub(prunedBytes, std::memory_order_relaxed);
+        }
         if (prunedVersions > 0) {
             spdlog::info("Pruned {} version history entries from collection '{}' (maxVersions: {} -> {})",
                          prunedVersions, name, oldMaxVersions, newMaxVersions);
@@ -1249,12 +1262,33 @@ std::vector<float> MemoryStore::extractVector(CollectionData& coll, nlohmann::js
 void MemoryStore::storeVector(CollectionData& coll, const std::string& docId,
                                const std::vector<float>& vec) {
     if (vec.empty()) return;
+
+    uint64_t oldSize = 0;
+    auto existing = coll.vectors.find(docId);
+    if (existing != coll.vectors.end()) {
+        oldSize = estimateVectorSize(existing->second);
+    }
+
     coll.vectors[docId] = vec;
     coll.vectorNorms[docId] = computeNorm(vec.data(), vec.size());
+
+    uint64_t newSize = estimateVectorSize(coll.vectors[docId]);
+
+    // Update memory counter by the delta
+    if (newSize > oldSize) {
+        estimatedMemoryBytes_.fetch_add(newSize - oldSize, std::memory_order_relaxed);
+    } else if (oldSize > newSize) {
+        estimatedMemoryBytes_.fetch_sub(oldSize - newSize, std::memory_order_relaxed);
+    }
 }
 
 void MemoryStore::removeVector(CollectionData& coll, const std::string& docId) {
-    coll.vectors.erase(docId);
+    auto it = coll.vectors.find(docId);
+    if (it != coll.vectors.end()) {
+        uint64_t freed = estimateVectorSize(it->second);
+        coll.vectors.erase(it);
+        estimatedMemoryBytes_.fetch_sub(freed, std::memory_order_relaxed);
+    }
     coll.vectorNorms.erase(docId);
 }
 
@@ -1480,13 +1514,18 @@ void MemoryStore::saveToHistory(CollectionData& coll, const Document& currentDoc
     ver.createdAt = currentDoc.createdAt;
     ver.createdBy = currentDoc.createdBy;
 
+    uint64_t addedBytes = estimateDocumentVersionSize(ver);
+
     auto& history = coll.versionHistory[currentDoc.id];
     history.push_back(std::move(ver));
+    estimatedMemoryBytes_.fetch_add(addedBytes, std::memory_order_relaxed);
 
     // Prune oldest versions if maxVersions is configured
     if (coll.options.maxVersions > 0) {
         while (history.size() > coll.options.maxVersions) {
+            uint64_t freed = estimateDocumentVersionSize(history.front());
             history.pop_front();
+            estimatedMemoryBytes_.fetch_sub(freed, std::memory_order_relaxed);
         }
     }
 }
@@ -1858,7 +1897,28 @@ void MemoryStore::loadVersionHistory(const std::string& collection, const std::s
     CollectionData* coll = getOrCreateCollection(collection);
 
     std::unique_lock<std::shared_mutex> lock(coll->mutex);
+
+    // Account for any history being replaced
+    uint64_t oldBytes = 0;
+    auto existing = coll->versionHistory.find(docId);
+    if (existing != coll->versionHistory.end()) {
+        for (const auto& ver : existing->second) {
+            oldBytes += estimateDocumentVersionSize(ver);
+        }
+    }
+
+    uint64_t newBytes = 0;
+    for (const auto& ver : history) {
+        newBytes += estimateDocumentVersionSize(ver);
+    }
+
     coll->versionHistory[docId] = std::move(history);
+
+    if (newBytes > oldBytes) {
+        estimatedMemoryBytes_.fetch_add(newBytes - oldBytes, std::memory_order_relaxed);
+    } else if (oldBytes > newBytes) {
+        estimatedMemoryBytes_.fetch_sub(oldBytes - newBytes, std::memory_order_relaxed);
+    }
 }
 
 void MemoryStore::clear() {
@@ -1867,6 +1927,10 @@ void MemoryStore::clear() {
 
     std::lock_guard<std::mutex> statsLock(statsMutex_);
     stats_ = Stats{};
+
+    // Reset the memory accounting counter — every document, version entry,
+    // and vector just went away, so any pending increments are stale.
+    estimatedMemoryBytes_.store(0, std::memory_order_relaxed);
 }
 
 // ===== Private Methods =====
@@ -2143,8 +2207,45 @@ void MemoryStore::removeFromExpirationIndex(CollectionData& coll, const std::str
 }
 
 uint64_t MemoryStore::estimateDocumentSize(const Document& doc) {
-    // Rough estimate: ID + collection + JSON dump size + metadata
-    return doc.id.size() + doc.collection.size() + doc.data.dump().size() + 100;
+    uint64_t size = 0;
+    size += doc.id.size();
+    size += doc.collection.size();
+    size += doc.createdBy.size();
+    size += doc.updatedBy.size();
+    size += doc.nodeId.size();
+    // JSON tree overhead: .dump() underestimates actual heap usage.
+    // A nlohmann::json object's heap footprint is roughly 2-3x the serialized size
+    // because of node allocations, hash tables, etc. Use 2.5x as a calibrated
+    // middle ground that avoids chronic under-count without being paranoid.
+    const std::string serialized = doc.data.dump();
+    size += static_cast<uint64_t>(serialized.size() * 5 / 2);  // 2.5x
+    // Encrypted-fields vector
+    for (const auto& f : doc.encryptedFields) {
+        size += f.size();
+    }
+    // Fixed Document overhead (versions, timestamps, pointers)
+    size += 128;
+    return size;
+}
+
+uint64_t MemoryStore::estimateVectorSize(const std::vector<float>& vec) {
+    // capacity might exceed size, but the hot path uses size-based estimate
+    // since std::vector rarely over-allocates by much for in-place stored data.
+    // Includes small constant for container overhead and the paired vectorNorms entry.
+    return vec.size() * sizeof(float) + 32;
+}
+
+uint64_t MemoryStore::estimateDocumentVersionSize(const DocumentVersion& ver) {
+    uint64_t size = 0;
+    size += ver.updatedBy.size();
+    size += ver.createdBy.size();
+    const std::string serialized = ver.data.dump();
+    size += static_cast<uint64_t>(serialized.size() * 5 / 2);
+    for (const auto& f : ver.encryptedFields) {
+        size += f.size();
+    }
+    size += 64;  // struct overhead + deque node overhead
+    return size;
 }
 
 void MemoryStore::expirationLoop() {

+ 6 - 0
service/src/memory_store.hpp

@@ -546,6 +546,12 @@ private:
     // Estimate document memory size
     [[nodiscard]] static uint64_t estimateDocumentSize(const Document& doc);
 
+    // Estimate vector memory size (float array + container overhead)
+    [[nodiscard]] static uint64_t estimateVectorSize(const std::vector<float>& vec);
+
+    // Estimate version history entry memory size
+    [[nodiscard]] static uint64_t estimateDocumentVersionSize(const DocumentVersion& ver);
+
     // Background expiration thread
     void expirationLoop();