From a9ef9a86c08e9b24351c21280005d527356de163 Mon Sep 17 00:00:00 2001
From: Vitaliy Filippov <vitalif@yourcmc.ru>
Date: Fri, 5 Jan 2024 01:31:29 +0300
Subject: [PATCH] Add update() API to kv_db

---
 src/kv_db.cpp | 49 ++++++++++++++++++++++++++++++++++++++++---------
 src/kv_db.h   |  3 +++
 2 files changed, 43 insertions(+), 9 deletions(-)

diff --git a/src/kv_db.cpp b/src/kv_db.cpp
index 42b89a2b..f883949f 100644
--- a/src/kv_db.cpp
+++ b/src/kv_db.cpp
@@ -1537,7 +1537,7 @@ void kv_op_t::resume_split()
 }
 
 void kv_op_t::update_block(int path_pos, bool is_delete, const std::string & key,
-    const std::string & value, std::function<void(int)> cb)
+    const std::string & arg_value, std::function<void(int)> cb)
 {
     auto blk_it = db->block_cache.find(path[path_pos].offset);
     if (blk_it == db->block_cache.end())
@@ -1552,7 +1552,7 @@ void kv_op_t::update_block(int path_pos, bool is_delete, const std::string & key
     if (blk->updating)
     {
         // Wait if block is being modified
-        db->continue_update.emplace(blk->offset, [=]() { update_block(path_pos, is_delete, key, value, cb); });
+        db->continue_update.emplace(blk->offset, [=]() { update_block(path_pos, is_delete, key, arg_value, cb); });
         return;
     }
     if (db->known_versions[blk->offset/db->ino_block_size] != block_ver || blk->invalidated)
@@ -1564,6 +1564,18 @@ void kv_op_t::update_block(int path_pos, bool is_delete, const std::string & key
     }
     uint32_t rm_size = 0;
     auto d_it = blk->data.find(key);
+    if (cas_cb && path_pos == path.size()-1 && !cas_cb(d_it != blk->data.end() ? 0 : -ENOENT, d_it != blk->data.end() ? d_it->second : ""))
+    {
+        // CAS failure
+        db->run_continue_update(blk->offset);
+        cb(-EAGAIN);
+        return;
+    }
+    if (path_pos == path.size()-1)
+    {
+        is_delete = op->opcode == KV_DEL;
+    }
+    const std::string & value = path_pos == path.size()-1 ? this->value : arg_value;
     if (d_it != blk->data.end())
     {
         if (!is_delete && d_it->second == value)
@@ -1582,13 +1594,6 @@ void kv_op_t::update_block(int path_pos, bool is_delete, const std::string & key
         cb(0);
         return;
     }
-    if (cas_cb && path_pos == path.size()-1 && !cas_cb(d_it != blk->data.end() ? 0 : -ENOENT, d_it != blk->data.end() ? d_it->second : ""))
-    {
-        // CAS failure
-        db->run_continue_update(blk->offset);
-        cb(-EAGAIN);
-        return;
-    }
     // This condition means this is a block split during previous updates
     // Its parent is already updated because prev_key_lt <= blk->right_half
     // That means we can safely remove the "split reference"
@@ -2007,6 +2012,32 @@ void kv_dbw_t::del(const std::string & key, std::function<void(int res)> cb,
     op->exec();
 }
 
+void kv_dbw_t::update(const std::string & key, std::function<int(int res, const std::string & old_value, std::string & new_value)> cas_cb,
+    std::function<void(int res)> cb)
+{
+    auto *op = new kv_op_t;
+    op->db = db;
+    op->opcode = KV_SET;
+    op->key = key;
+    op->cas_cb = [cas_cb, op](int res, const std::string & old_value)
+    {
+        int action = cas_cb(res, old_value, op->value);
+        if (action == 1)
+            op->opcode = KV_SET;
+        else if (action == 2 && res != -ENOENT)
+            op->opcode = KV_DEL;
+        else
+            return false;
+        return true;
+    };
+    op->callback = [cb](kv_op_t *op)
+    {
+        cb(op->res);
+        delete op;
+    };
+    op->exec();
+}
+
 void* kv_dbw_t::list_start(const std::string & start)
 {
     auto *op = new kv_op_t;
diff --git a/src/kv_db.h b/src/kv_db.h
index 79d0b97c..33e15462 100644
--- a/src/kv_db.h
+++ b/src/kv_db.h
@@ -27,6 +27,9 @@ struct kv_dbw_t
         std::function<bool(int res, const std::string & value)> cas_compare = NULL);
     void del(const std::string & key, std::function<void(int res)> cb,
         std::function<bool(int res, const std::string & value)> cas_compare = NULL);
+    void update(const std::string & key,
+        std::function<int(int res, const std::string & old_value, std::string & new_value)> cas_compare,
+        std::function<bool(int res)> cb);
 
     void* list_start(const std::string & start);
     void list_next(void *handle, std::function<void(int res, const std::string & key, const std::string & value)> cb);