feat(lib): add keep_prefix functionality for long keys (#67)

Introduce `keep_prefix` parameter for `get` and `set` methods. The `keep_prefix` preserves the first 479 bytes of keys longer than 479 bytes, appending the SHA-256 hash to maintain prefix search capability with `lmdb.prefix()` while complying with LMDB's 511-byte key size limit. This transformation occurs specifically when keys exceed 479 bytes. If off, when the key exceeds LMDB's 511-byte key size limit, the entire key is automatically hashed with sha-256 and used as the actual key.

Author: HanadaLee

README.md

@@ -37,7 +37,7 @@ interacting with the module to access/change data.
 
 #### get
 
-**syntax:** *value, err = lmdb.get(key, db?)*
+**syntax:** *value, err = lmdb.get(key, db?, keep_prefix?)*
 
 **context:** *any context **except** init_by_lua**
 
@@ -46,13 +46,15 @@ it defaults to `"_default"`.
 
 If the key does not exist, `nil` will be returned.
 
+The `keep_prefix` preserves the first 479 bytes of keys longer than 479 bytes, appending the SHA-256 hash to maintain prefix search capability with `lmdb.prefix()` while complying with LMDB's 511-byte key size limit. This transformation occurs specifically when keys exceed 479 bytes. If off, when the key exceeds LMDB's 511-byte key size limit, the entire key is automatically hashed with sha-256 and used as the actual key.
+
 In case of error, `nil` and a string describing the error will be returned instead.
 
 [Back to TOC](#table-of-contents)
 
 #### set
 
-**syntax:** *ok, err = lmdb.set(key, value, db?)*
+**syntax:** *ok, err = lmdb.set(key, value, db?, keep_prefix?)*
 
 **context:** *any context **except** init_by_lua**
 
@@ -61,6 +63,8 @@ it defaults to `"_default"`.
 
 Setting a key's value to `nil` will remove that key from the corresponding database.
 
+The `keep_prefix` preserves the first 479 bytes of keys longer than 479 bytes, appending the SHA-256 hash to maintain prefix search capability with `lmdb.prefix()` while complying with LMDB's 511-byte key size limit. This transformation occurs specifically when keys exceed 479 bytes. If off, when the key exceeds LMDB's 511-byte key size limit, the entire key is automatically hashed with sha-256 and used as the actual key.
+
 In case of error, `nil` and a string describing the error will be returned instead.
 
 [Back to TOC](#table-of-contents)

lib/resty/lmdb.lua

@@ -26,9 +26,9 @@ local CAN_YIELD_PHASES = {
 }
 
 
-function _M.get(key, db)
+function _M.get(key, db, keep_prefix)
     CACHED_TXN:reset()
-    CACHED_TXN:get(key, db)
+    CACHED_TXN:get(key, db, keep_prefix)
     local res, err = CACHED_TXN:commit()
     if not res then
         return nil, err
@@ -38,9 +38,9 @@ function _M.get(key, db)
 end
 
 
-function _M.set(key, value, db)
+function _M.set(key, value, db, keep_prefix)
     CACHED_TXN:reset()
-    CACHED_TXN:set(key, value, db)
+    CACHED_TXN:set(key, value, db, keep_prefix)
     local res, err = CACHED_TXN:commit()
     if not res then
         return nil, err

lib/resty/lmdb/transaction.lua

@@ -12,6 +12,7 @@ local get_string_buf = base.get_string_buf
 local get_string_buf_size = base.get_string_buf_size
 local setmetatable = setmetatable
 local assert = assert
+local string_sub = string.sub
 local math_max = math.max
 local type = type
 local C = ffi.C
@@ -51,16 +52,26 @@ do
 
     -- lmdb has 511 bytes limitation for key
     local MAX_KEY_SIZE = 511
+    -- when keep prefix, the key size is limited to 479 bytes
+    local MAX_KEY_SIZE_WHEN_KEEP_PREFIX = 511 - 32 -- it is 479
 
     local sha256 = function(str)
         resty_sha256:reset()
         resty_sha256:update(str)
         return resty_sha256:final()
     end
 
-    normalize_key = function(key)
-        if key and #key > MAX_KEY_SIZE then
-            return assert(sha256(key))
+    normalize_key = function(key, keep_prefix)
+        if key then
+            local key_len = #key
+            if keep_prefix and key_len > MAX_KEY_SIZE_WHEN_KEEP_PREFIX then
+                local prefix = string_sub(key, 1, MAX_KEY_SIZE_WHEN_KEEP_PREFIX)
+                return prefix .. assert(sha256(key))
+            end
+
+            if key_len > MAX_KEY_SIZE then
+                return assert(sha256(key))
+            end
         end
 
         return key
@@ -114,18 +125,18 @@ function _TXN_MT:_new_op()
 end
 
 
-function _TXN_MT:get(key, db)
+function _TXN_MT:get(key, db, keep_prefix)
     local op = self:_new_op()
     op.opcode = "GET"
-    op.key = normalize_key(key)
+    op.key = normalize_key(key, keep_prefix)
     op.db = db or DEFAULT_DB
 end
 
 
-function _TXN_MT:set(key, value, db)
+function _TXN_MT:set(key, value, db, keep_prefix)
     local op = self:_new_op()
     op.opcode = "SET"
-    op.key = normalize_key(key)
+    op.key = normalize_key(key, keep_prefix)
     op.value = value
     op.db = db or DEFAULT_DB
     op.flags = 0

t/08-max-key-size.t

@@ -149,3 +149,214 @@ value2
 [error]
 [warn]
 [crit]
+
+
+=== TEST 6: key size is 480 set() / get() with keep_prefix
+--- http_config eval: $::HttpConfig
+--- main_config eval: $::MainConfig
+--- config
+    location = /t {
+        content_by_lua_block {
+            local l = require("resty.lmdb")
+
+            local key = string.rep("a", 480)
+            ngx.say(l.set(key, "value", true))
+            ngx.say(l.get(key, true))
+        }
+    }
+--- request
+GET /t
+--- response_body
+true
+value
+--- no_error_log
+[error]
+[warn]
+[crit]
+
+
+=== TEST 7: key size is 479 set() / get() with keep_prefix
+--- http_config eval: $::HttpConfig
+--- main_config eval: $::MainConfig
+--- config
+    location = /t {
+        content_by_lua_block {
+            local l = require("resty.lmdb")
+
+            local key = string.rep("a", 479)
+            ngx.say(l.set(key, "value", true))
+            ngx.say(l.get(key, true))
+        }
+    }
+--- request
+GET /t
+--- response_body
+true
+value
+--- no_error_log
+[error]
+[warn]
+[crit]
+
+
+=== TEST 8: key size is 1024 set() / get() with keep_prefix
+--- http_config eval: $::HttpConfig
+--- main_config eval: $::MainConfig
+--- config
+    location = /t {
+        content_by_lua_block {
+            local l = require("resty.lmdb")
+
+            local key = string.rep("a", 1024)
+            ngx.say(l.set(key, "value", true))
+            ngx.say(l.get(key, true))
+        }
+    }
+--- request
+GET /t
+--- response_body
+true
+value
+--- no_error_log
+[error]
+[warn]
+[crit]
+
+
+=== TEST 9: key size is 416 set() / get() with keep_prefix
+--- http_config eval: $::HttpConfig
+--- main_config eval: $::MainConfig
+--- config
+    location = /t {
+        content_by_lua_block {
+            local l = require("resty.lmdb")
+
+            local key = string.rep("a", 416)
+            ngx.say(l.set(key, "value", true))
+            ngx.say(l.get(key, true))
+        }
+    }
+--- request
+GET /t
+--- response_body
+true
+value
+--- no_error_log
+[error]
+[warn]
+[crit]
+
+
+=== TEST 10: different keys are ok with keep_prefix
+--- http_config eval: $::HttpConfig
+--- main_config eval: $::MainConfig
+--- config
+    location = /t {
+        content_by_lua_block {
+            local l = require("resty.lmdb")
+
+            local key1 = string.rep("a", 480)
+            ngx.say(l.set(key1, "value1", true))
+            ngx.say(l.get(key1, true))
+
+            local key2 = string.rep("b", 480)
+            ngx.say(l.set(key2, "value2", true))
+            ngx.say(l.get(key2, true))
+        }
+    }
+--- request
+GET /t
+--- response_body
+true
+value1
+true
+value2
+--- no_error_log
+[error]
+[warn]
+[crit]
+
+
+=== TEST 11: direct get value via normalized key with keep_prefix
+--- http_config eval: $::HttpConfig
+--- main_config eval: $::MainConfig
+--- config
+    location = /t {
+        content_by_lua_block {
+            local l = require("resty.lmdb")
+            local resty_sha256 = assert(require("resty.sha256").new())
+
+            local sha256 = function(str)
+                resty_sha256:reset()
+                resty_sha256:update(str)
+                return resty_sha256:final()
+            end
+
+            local key = string.rep("a", 480)
+            ngx.say(l.set(key, "value", true))
+            ngx.say(l.get(key, true))
+
+            local normalized_key = string.sub(key, 1, 479) .. assert(sha256(key))
+            ngx.say(l.get(normalized_key))
+        }
+    }
+--- request
+GET /t
+--- response_body
+true
+value
+value
+--- no_error_log
+[error]
+[warn]
+[crit]
+
+
+=== TEST 12: get values via prefix with keep_prefix
+--- http_config eval: $::HttpConfig
+--- main_config eval: $::MainConfig
+--- config
+    location = /t {
+        content_by_lua_block {
+            local l = require("resty.lmdb")
+
+            local prefix = string.rep("a", 240)
+
+            local key1 = prefix .. string.rep("b", 240)
+            ngx.say(l.set(key1, "value1", true))
+            ngx.say(l.get(key1, true))
+
+            local key2 = prefix .. string.rep("c", 240)
+            ngx.say(l.set(key2, "value2", true))
+            ngx.say(l.get(key2, true))
+
+            local key3 = string.rep("d", 480)
+            ngx.say(l.set(key3, "value3", true))
+            ngx.say(l.get(key3, true))
+
+            local key4 = string.rep("a", 239)
+            ngx.say(l.set(key4, "value4", true))
+            ngx.say(l.get(key4, true))
+
+            for k, v in l.prefix(prefix) do
+                ngx.say(value: ", v)
+            end
+        }
+    }
+--- request
+GET /t
+--- response_body
+true
+value1
+true
+value2
+true
+value3
+true
+value4
+value: value1
+value: value2
+--- no_error_log
+[error]
+[warn]
+[crit]