Fix log error of radix tree deserialization. (#1811)

- Fix log error of radix tree deserialization. - Fix insert twice root node when deserialization. Fixes #1810 Signed-off-by: vegetableysm <[email protected]>
v6d-io · Mar 11, 2024 · db102b9 · db102b9
1 parent a5c8a67
commit db102b9
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 19 deletions.
diff --git a/modules/llm-cache/radix-tree/radix-tree.cc b/modules/llm-cache/radix-tree/radix-tree.cc
@@ -21,28 +21,30 @@ limitations under the License.
 
 namespace vineyard {
 
-RadixTree::RadixTree(int cacheCapacity) {
+RadixTree::RadixTree(int cacheCapacity, bool withRoot) {
   this->tree = raxNew();
   // add one to the cache capacity because we insert a root node to the tree.
   this->cacheCapacity = cacheCapacity + 1;
   this->nodeCount = 0;
 
   // prepare root node
-  std::vector<int> rootToken = {INT32_MAX};
-  std::shared_ptr<NodeData> evictedNode;
-  this->InsertInternal(rootToken, evictedNode);
-  if (VLOG_IS_ON(100)) {
-    VLOG(100) << raxShow(this->tree);
+  if (withRoot) {
+    std::vector<int> rootToken = {INT32_MAX};
+    std::shared_ptr<NodeData> evictedNode;
+    this->InsertInternal(rootToken, evictedNode);
+    if (VLOG_IS_ON(100)) {
+      VLOG(100) << raxShow(this->tree);
+    }
+    raxNode* dataNode =
+        raxFindAndReturnDataNode(this->tree, rootToken, NULL, false);
+    DataWrapper* data = new DataWrapper();
+    data->data = nullptr;
+    data->dataLength = 0;
+    dataNode->custom_data = data;
+    VLOG(100) << "root data wrapper:" << data;
+    dataNode->issubtree = true;
+    this->rootToken = rootToken;
   }
-  raxNode* dataNode =
-      raxFindAndReturnDataNode(this->tree, rootToken, NULL, false);
-  DataWrapper* data = new DataWrapper();
-  data->data = nullptr;
-  data->dataLength = 0;
-  dataNode->custom_data = data;
-  VLOG(100) << "root data wrapper:" << data;
-  dataNode->issubtree = true;
-  this->rootToken = rootToken;
 }
 
 RadixTree::~RadixTree() {
@@ -430,7 +432,7 @@ std::shared_ptr<RadixTree> RadixTree::Deserialize(std::string data) {
   // is created by upper layer. Here just recover it from serialized
   // string.
   std::shared_ptr<RadixTree> radixTree =
-      std::make_shared<RadixTree>(cacheCapacity);
+      std::make_shared<RadixTree>(cacheCapacity, false);
   radixTree->nodeCount = tokenList.size();
 
   if (VLOG_IS_ON(100)) {
@@ -450,8 +452,8 @@ std::shared_ptr<RadixTree> RadixTree::Deserialize(std::string data) {
     int retval =
         raxInsertAndReturnDataNode(radixTree->tree, tokenList[i], data,
                                    reinterpret_cast<void**>(&dataNode), NULL);
-    if (retval == 0 || dataNode == NULL) {
-      LOG(ERROR) << "Insert token list failed";
+    if (retval == 0) {
+      LOG(WARNING) << "Overwrite the data of the node";
     }
     dataNode->timestamp = timestampList[i];
     dataNode->numnodes = subTreeSizeList[i];

diff --git a/modules/llm-cache/radix-tree/radix-tree.h b/modules/llm-cache/radix-tree/radix-tree.h
@@ -76,7 +76,7 @@ class RadixTree : public std::enable_shared_from_this<RadixTree> {
       const std::vector<int>& tokens, std::shared_ptr<NodeData>& header);
 
  public:
-  explicit RadixTree(int cacheCapacity);
+  explicit RadixTree(int cacheCapacity, bool withRoot = true);
 
   ~RadixTree();