feat(prefix): prefix based lookup feature (#48)

This facilitates efficient enumeration of keys starting with a common prefix.

Previously a index key will need to be created, with this PR it is no longer
needed and it is possible to lookup all keys starting with a common prefix
using an iterator.

KAG-2984

README.md

@@ -5,8 +5,7 @@ inside the Nginx worker process. It has two parts, a core module built into Ngin
 controls the life cycle of the database environment, and a FFI based Lua binding for
 interacting with the module to access/change data.
 
-Table of Contents
-=================
+# Table of Contents
 
 * [lua-resty-lmdb](#lua-resty-lmdb)
     * [APIs](#apis)
@@ -15,14 +14,17 @@ Table of Contents
             * [set](#set)
             * [get_env_info](#get_env_info)
         * [db\_drop](#db_drop)
+        * [prefix](#prefix)
         * [resty.lmdb.transaction](#restylmdbtransaction)
             * [reset](#reset)
             * [get](#get)
             * [set](#set)
             * [db\_open](#db_open)
             * [db\_drop](#db_drop)
             * [commit](#commit)
-    * [Directives](#Directives)
+        * [resty.lmdb.prefix](#restylmdbprefix)
+            * [reset](#reset)
+    * [Directives](#directives)
         * [lmdb_environment_path](#lmdb_environment_path)
         * [lmdb_max_databases](#lmdb_max_databases)
         * [lmdb_map_size](#lmdb_map_size)
@@ -99,18 +101,42 @@ In case of error, `nil` and a string describing the error will be returned inste
 
 [Back to TOC](#table-of-contents)
 
-### resty.lmdb.transaction
+### prefix
 
-**syntax:** *local txn = transaction.begin(hint?)*
+**syntax:** *for key, value in lmdb.prefix(prefix) do*
 
 **context:** *any context*
 
-Creates a new LMDB transaction object. This does not actually starts the transaction, but only creates
-a Lua table that stores the operations for execution later. If `hint` is provided then the Lua table holding
-the operations will be pre-allocated to store `hint` operations.
+Returns all key and their associated value for keys starting with `prefix`.
+For example, if the database contains:
+
+```
+key1: value1
+key11: value11
+key2: value2
+```
+
+Then a call of `lmdb.prefix("key")` will yield `key1`, `key11` and `key2` respectively.
+
+In case of errors while fetching from LMDB, `key` will be `nil` and `value` will be
+a string describing the error. The caller must anticipate this happening and check each return
+value carefully before consuming.
+
+**Warning on transaction safety:** Since the number of keys that could potentially
+be returned with this method could be very large, this method does not return all
+results inside a single transaction as this will be very expensive. Instead, this
+method gets keys from LMDB in batches using different read transaction. Therefore, it
+is possible that the database content has changed between batches. We may introduce a
+mechanism for detecting this case in the future, but for now there is a small opportunity
+for this to happen and you should guard your application for concurrent writes if this
+is a huge concern. This function makes best effort to detect when database content
+definitely changed between iterations, in this case `nil, "DB content changed while iterating"`
+will be returned from the iterator.
 
 [Back to TOC](#table-of-contents)
 
+### resty.lmdb.transaction
+
 #### reset
 
 **syntax:** *txn:reset()*
@@ -197,6 +223,26 @@ from the `txn` table when `commit()` returned an error is undefined.
 
 [Back to TOC](#table-of-contents)
 
+### resty.lmdb.prefix
+
+#### page
+
+**syntax:** *res, err = prefix.page(start, prefix, db?)*
+
+**context:** *any context*
+
+Return all keys `>= start` and starts with `prefix`. If `db` is omitted,
+it defaults to `"_default"`.
+
+The return value of this function is a table `res` where `res[1].key` and `res[1].value`
+corresponds to the first key and value, `res[2].key` and `res[2].value` corresponds to the
+second and etc. If no keys matched the provided criteria, then an empty table will be
+returned.
+
+In case of errors, `nil` and an string describing the reason of the failure will be returned.
+
+[Back to TOC](#table-of-contents)
+
 ## Directives
 
 ### lmdb_environment_path

config

@@ -2,6 +2,7 @@ ngx_module_type=CORE
 ngx_module_name=ngx_lua_resty_lmdb_module
 ngx_module_srcs="$ngx_addon_dir/src/ngx_lua_resty_lmdb_module.c
                  $ngx_addon_dir/src/ngx_lua_resty_lmdb_transaction.c
+                 $ngx_addon_dir/src/ngx_lua_resty_lmdb_prefix.c
                  $ngx_addon_dir/src/ngx_lua_resty_lmdb_status.c
                  "
 ngx_module_incs="$ngx_addon_dir/lmdb/libraries/liblmdb $ngx_addon_dir/src"

lib/resty/lmdb.lua

@@ -2,47 +2,128 @@ local _M = {}
 
 
 local transaction = require("resty.lmdb.transaction")
+local prefix = require("resty.lmdb.prefix")
 local status = require("resty.lmdb.status")
 
-do
-    local CACHED_TXN = transaction.begin(1)
 
+local assert = assert
+local prefix_page = prefix.page
+local get_phase = ngx.get_phase
+local ngx_sleep = ngx.sleep
+
+
+local CACHED_TXN = transaction.begin(1)
+local CAN_YIELD_PHASES = {
+    rewrite = true,
+    server_rewrite = true,
+    access = true,
+    content = true,
+    timer = true,
+    ssl_client_hello = true,
+    ssl_certificate = true,
+    ssl_session_fetch = true,
+    preread = true,
+}
+
+
+function _M.get(key, db)
+    CACHED_TXN:reset()
+    CACHED_TXN:get(key, db)
+    local res, err = CACHED_TXN:commit()
+    if not res then
+        return nil, err
+    end
+
+    return CACHED_TXN[1].result
+end
 
-    function _M.get(key, db)
-        CACHED_TXN:reset()
-        CACHED_TXN:get(key, db)
-        local res, err = CACHED_TXN:commit()
-        if not res then
-            return nil, err
-        end
 
-        return CACHED_TXN[1].result
+function _M.set(key, value, db)
+    CACHED_TXN:reset()
+    CACHED_TXN:set(key, value, db)
+    local res, err = CACHED_TXN:commit()
+    if not res then
+        return nil, err
     end
 
+    return true
+end
 
-    function _M.set(key, value, db)
-        CACHED_TXN:reset()
-        CACHED_TXN:set(key, value, db)
-        local res, err = CACHED_TXN:commit()
-        if not res then
-            return nil, err
-        end
 
-        return true
+function _M.db_drop(delete, db)
+    delete = not not delete
+
+    CACHED_TXN:reset()
+    CACHED_TXN:db_drop(delete, db)
+    local res, err = CACHED_TXN:commit()
+    if not res then
+        return nil, err
     end
 
+    return true
+end
 
-    function _M.db_drop(delete, db)
-        delete = not not delete
 
-        CACHED_TXN:reset()
-        CACHED_TXN:db_drop(delete, db)
-        local res, err = CACHED_TXN:commit()
+function _M.prefix(prefix, db)
+    local res, i, res_n, err_or_more
+    local last = prefix
+    local can_yield = CAN_YIELD_PHASES[get_phase()]
+
+    return function()
+        ::more::
         if not res then
-            return nil, err
+            -- need to fetch more data
+            res, err_or_more = prefix_page(last, prefix, db)
+            if not res then
+                return nil, err_or_more
+            end
+
+            res_n = #res
+            if res_n == 0 or (i and res_n == 1) then
+                return nil
+            end
+
+            if i then
+                -- not the first call to prefix_page
+                if res[1].key ~= last then
+                    return nil, "DB content changed while iterating"
+                end
+
+                -- this is not sufficient to prove DB content did not change,
+                -- but at least the resume point did not change.
+                -- skip the first key
+                i = 2
+
+            else
+                -- first call to prefix_page
+                i = 1
+            end
         end
 
-        return true
+        assert(res_n > 0)
+
+        if i > res_n then
+            if err_or_more then
+                last = res[i - 1].key
+                res = nil
+
+                if can_yield then
+                    ngx_sleep(0)
+                end
+
+                goto more
+            end
+
+            -- more = false
+
+            return nil
+        end
+
+        local key = res[i].key
+        local value = res[i].value
+        i = i + 1
+
+        return key, value
     end
 end
 

lib/resty/lmdb/cdefs.lua

@@ -0,0 +1,47 @@
+local ffi = require("ffi")
+local base = require("resty.core.base")
+
+
+local DEFAULT_VALUE_BUF_SIZE = 512 * 2048 -- 1MB
+base.set_string_buf_size(DEFAULT_VALUE_BUF_SIZE)
+
+
+ffi.cdef([[
+typedef unsigned int 	MDB_dbi;
+
+
+typedef enum {
+    NGX_LMDB_OP_GET = 0,
+    NGX_LMDB_OP_PREFIX,
+    NGX_LMDB_OP_SET,
+    NGX_LMDB_OP_DB_OPEN,
+    NGX_LMDB_OP_DB_DROP
+} ngx_lua_resty_lmdb_operation_e;
+
+
+typedef struct {
+    ngx_lua_resty_lmdb_operation_e  opcode;
+    ngx_str_t                       key;   /* GET, SET */
+    ngx_str_t                       value; /* GET, SET */
+    MDB_dbi                         dbi;   /* ALL OPS */
+    unsigned int                    flags; /* SET, DROP */
+} ngx_lua_resty_lmdb_operation_t;
+
+typedef struct {
+    size_t          map_size;        /**< Size of the data memory map */
+    unsigned int    page_size;       /**< Size of a database page. */
+    unsigned int    max_readers;     /**< max reader slots in the environment */
+    unsigned int    num_readers;     /**< max reader slots used in the environment */
+    unsigned int    allocated_pages; /**< number of pages allocated */
+    size_t          in_use_pages;    /**< number of pages currently in-use */
+    unsigned int    entries;         /**< the number of entries (key/value pairs) in the environment */
+} ngx_lua_resty_lmdb_ffi_status_t;
+
+int ngx_lua_resty_lmdb_ffi_env_info(ngx_lua_resty_lmdb_ffi_status_t *lst, char **err);
+
+
+int ngx_lua_resty_lmdb_ffi_execute(ngx_lua_resty_lmdb_operation_t *ops,
+    size_t n, int need_write, unsigned char *buf, size_t buf_len, char **err);
+int ngx_lua_resty_lmdb_ffi_prefix(ngx_lua_resty_lmdb_operation_t *ops,
+    size_t n, unsigned char *buf, size_t buf_len, char **err);
+]])