call bucket_ref/bucket_unref on crud operations

Satbek · Satbek · commit d82469cf64db · 2025-11-21T14:19:02.000+05:00
added module bucket_ref_unref.lua with functions bucket_refrw, bucket_unrefrw,
bucket_refro, etc ...

module looks at rebalancer.safe_mode flag. if it's disabled module does nothing,
otherwise module call ref fucnctions

bucket_ref fail indicated bucket movement.
it will help with problem with double buckets.
diff --git a/crud/common/sharding/bucket_ref_unref.lua b/crud/common/sharding/bucket_ref_unref.lua
@@ -0,0 +1,89 @@
+--- module to call vshard.storage.bucket_ref / vshard.storage.bucket_unref
+--- on write requests
+--- there are two modes: safe and fast. on safe mode module
+--- calls vshard.storage.bucket_ref / vshard.storage.bucket_unref
+--- on fast mode it does nothing.
+--- default is fast mode.
+
+--- bucket_refw and bucket_unrefrw must be called in one transaction in order to prevent
+--- safe_mode change during execution.
+
+local vshard = require('vshard')
+
+local M = {}
+
+--- on module initialization safe_mode_status func must be set
+--- it's rebalance.safe_mode_status
+function M.safe_mode_status()
+    error('safe_mode_status not set')
+end
+
+--- set safe mode func
+--- from rabalance.safe_mode_status
+function M.set_safe_mode_status(safe_mode_status)
+    M.safe_mode_status = safe_mode_status
+end
+
+--- bucket_refrw calls vshard.storage.bucket_refrw if safe mode enabled
+--- otherwise returns true
+--- must be called with bucket_unrefrw in transaction
+function M.bucket_refrw(bucket_id)
+    if not M.safe_mode_status() then
+        return true
+    end
+    return vshard.storage.bucket_refrw(bucket_id)
+end
+
+--- bucket_unrefrw calls vshard.storage.bucket_unrefrw if safe mode enabled
+--- otherwise returns true
+--- must be called with bucket_refrw in transaction
+function M.bucket_unrefrw(bucket_id)
+    if not M.safe_mode_status() then
+        return true
+    end
+    return vshard.storage.bucket_unrefrw(bucket_id)
+end
+
+--- bucket_refro calls vshard.storage.bucket_refro if safe mode enabled
+--- otherwise returns true
+function M.bucket_refro(bucket_id)
+    if not M.safe_mode_status() then
+        return true
+    end
+    return vshard.storage.bucket_refro(bucket_id)
+end
+
+--- bucket_unrefro calls vshard.storage.bucket_unrefro if safe mode enabled
+--- otherwise returns true
+function M.bucket_unrefro(bucket_id)
+    if not M.safe_mode_status() then
+        return true
+    end
+    return vshard.storage.bucket_unrefro(bucket_id)
+end
+
+--- bucket_unrefro calls vshard.storage.bucket_unrefrw for every bucket if safe mode enabled
+--- otherwise returns true
+--- @param bucket_ids table<number, true>
+function M.bucket_unrefrw_many(bucket_ids)
+    if not M.safe_mode_status() then
+        return true
+    end
+
+    local unref_all_ok = true
+    local unref_last_err
+    for reffed_bucket_id in pairs(bucket_ids) do
+        local unref_ok, unref_err = M.bucket_unrefrw(reffed_bucket_id)
+        if not unref_ok then
+            unref_all_ok = nil
+            unref_last_err = unref_err
+        end
+    end
+
+    if not unref_all_ok then
+        return nil, unref_last_err
+    end
+    return true
+end
+
+return M
diff --git a/crud/delete.lua b/crud/delete.lua
@@ -9,6 +9,7 @@ local sharding_key_module = require('crud.common.sharding.sharding_key')
 local sharding_metadata_module = require('crud.common.sharding.sharding_metadata')
 local dev_checks = require('crud.common.dev_checks')
 local schema = require('crud.common.schema')
+local bucket_ref_unref = require('crud.common.sharding.bucket_ref_unref')
 
 local DeleteError = errors.new_class('DeleteError', {capture_stack = false})
 
@@ -19,6 +20,7 @@ local CRUD_DELETE_FUNC_NAME = utils.get_storage_call(DELETE_FUNC_NAME)
 
 local function delete_on_storage(space_name, key, field_names, opts)
     dev_checks('string', '?', '?table', {
+        bucket_id = 'number|cdata',
         sharding_key_hash = '?number',
         sharding_func_hash = '?number',
         skip_sharding_hash_check = '?boolean',
@@ -42,14 +44,29 @@ local function delete_on_storage(space_name, key, field_names, opts)
         return nil, err
     end
 
-    -- add_space_schema_hash is false because
-    -- reloading space format on router can't avoid delete error on storage
-    return schema.wrap_box_space_func_result(space, 'delete', {key}, {
-        add_space_schema_hash = false,
-        field_names = field_names,
-        noreturn = opts.noreturn,
-        fetch_latest_metadata = opts.fetch_latest_metadata,
-    })
+    local function make_delete()
+        local ref_ok, err_ref = bucket_ref_unref.bucket_refrw(opts.bucket_id)
+        if not ref_ok then
+            return nil, err_ref
+        end
+
+        -- add_space_schema_hash is false because
+        -- reloading space format on router can't avoid delete error on storage
+        local result = schema.wrap_box_space_func_result(space, 'delete', {key}, {
+            add_space_schema_hash = false,
+            field_names = field_names,
+            noreturn = opts.noreturn,
+            fetch_latest_metadata = opts.fetch_latest_metadata,
+        })
+        local unref_ok, err_unref = bucket_ref_unref.bucket_unrefrw(opts.bucket_id)
+        if not unref_ok then
+            return nil, err_unref
+        end
+
+        return result
+    end
+
+    return box.atomic(make_delete)
 end
 
 delete.storage_api = {[DELETE_FUNC_NAME] = delete_on_storage}
@@ -116,6 +133,7 @@ local function call_delete_on_router(vshard_router, space_name, key, opts)
     sharding.fill_bucket_id_pk(space, key, bucket_id_data.bucket_id)
 
     local delete_on_storage_opts = {
+        bucket_id = bucket_id_data.bucket_id,
         sharding_func_hash = bucket_id_data.sharding_func_hash,
         sharding_key_hash = sharding_key_hash,
         skip_sharding_hash_check = skip_sharding_hash_check,
diff --git a/crud/get.lua b/crud/get.lua
@@ -9,6 +9,7 @@ local sharding_key_module = require('crud.common.sharding.sharding_key')
 local sharding_metadata_module = require('crud.common.sharding.sharding_metadata')
 local dev_checks = require('crud.common.dev_checks')
 local schema = require('crud.common.schema')
+local bucket_ref_unref = require('crud.common.sharding.bucket_ref_unref')
 
 local GetError = errors.new_class('GetError', {capture_stack = false})
 
@@ -19,6 +20,7 @@ local CRUD_GET_FUNC_NAME = utils.get_storage_call(GET_FUNC_NAME)
 
 local function get_on_storage(space_name, key, field_names, opts)
     dev_checks('string', '?', '?table', {
+        bucket_id = 'number|cdata',
         sharding_key_hash = '?number',
         sharding_func_hash = '?number',
         skip_sharding_hash_check = '?boolean',
@@ -41,13 +43,23 @@ local function get_on_storage(space_name, key, field_names, opts)
         return nil, err
     end
 
+    local ref_ok, err_ref = bucket_ref_unref.bucket_refro(opts.bucket_id)
+    if not ref_ok then
+        return nil, err_ref
+    end
     -- add_space_schema_hash is false because
     -- reloading space format on router can't avoid get error on storage
-    return schema.wrap_box_space_func_result(space, 'get', {key}, {
+    local result = schema.wrap_box_space_func_result(space, 'get', {key}, {
         add_space_schema_hash = false,
         field_names = field_names,
         fetch_latest_metadata = opts.fetch_latest_metadata,
     })
+    local unref_ok, err_unref = bucket_ref_unref.bucket_unrefro(opts.bucket_id)
+    if not unref_ok then
+        return nil, err_unref
+    end
+
+    return result
 end
 
 get.storage_api = {[GET_FUNC_NAME] = get_on_storage}
@@ -114,6 +126,7 @@ local function call_get_on_router(vshard_router, space_name, key, opts)
     sharding.fill_bucket_id_pk(space, key, bucket_id_data.bucket_id)
 
     local get_on_storage_opts = {
+        bucket_id = bucket_id_data.bucket_id,
         sharding_func_hash = bucket_id_data.sharding_func_hash,
         sharding_key_hash = sharding_key_hash,
         skip_sharding_hash_check = skip_sharding_hash_check,
diff --git a/crud/replace.lua b/crud/replace.lua
@@ -7,6 +7,7 @@ local utils = require('crud.common.utils')
 local sharding = require('crud.common.sharding')
 local dev_checks = require('crud.common.dev_checks')
 local schema = require('crud.common.schema')
+local bucket_ref_unref = require('crud.common.sharding.bucket_ref_unref')
 
 local ReplaceError = errors.new_class('ReplaceError', { capture_stack = false })
 
@@ -42,15 +43,29 @@ local function replace_on_storage(space_name, tuple, opts)
         return nil, err
     end
 
-    -- add_space_schema_hash is true only in case of replace_object
-    -- the only one case when reloading schema can avoid insert error
-    -- is flattening object on router
-    return schema.wrap_box_space_func_result(space, 'replace', {tuple}, {
-        add_space_schema_hash = opts.add_space_schema_hash,
-        field_names = opts.fields,
-        noreturn = opts.noreturn,
-        fetch_latest_metadata = opts.fetch_latest_metadata,
-    })
+    local function make_replace()
+        local bucket_id = tuple[utils.get_bucket_id_fieldno(space)]
+        local ref_ok, err_ref = bucket_ref_unref.bucket_refrw(bucket_id)
+        if not ref_ok then
+            return nil, err_ref
+        end
+        -- add_space_schema_hash is true only in case of replace_object
+        -- the only one case when reloading schema can avoid insert error
+        -- is flattening object on router
+        local result = schema.wrap_box_space_func_result(space, 'replace', {tuple}, {
+            add_space_schema_hash = opts.add_space_schema_hash,
+            field_names = opts.fields,
+            noreturn = opts.noreturn,
+            fetch_latest_metadata = opts.fetch_latest_metadata,
+        })
+        local unref_ok, err_unref = bucket_ref_unref.bucket_unrefrw(bucket_id)
+        if not unref_ok then
+            return nil, err_unref
+        end
+        return result
+    end
+
+    return box.atomic(make_replace)
 end
 
 replace.storage_api = {[REPLACE_FUNC_NAME] = replace_on_storage}
diff --git a/crud/storage.lua b/crud/storage.lua
@@ -23,6 +23,7 @@ local count = require('crud.count')
 local borders = require('crud.borders')
 local readview = require('crud.readview')
 local storage_info = require('crud.storage_info')
+local bucket_ref_unref = require('crud.common.sharding.bucket_ref_unref')
 
 local storage = {}
 
@@ -106,6 +107,7 @@ local function init_impl()
     end
 
     rebalance.init()
+    bucket_ref_unref.set_safe_mode_status(rebalance.safe_mode_status)
 
     for _, module in ipairs(modules_with_storage_api) do
         init_storage_call(user, module.storage_api)
