view mod_storage_lmdb/mod_storage_lmdb.lua @ 5139:449e4ca4de32

mod_unified_push: Remove dependency on trunk util.jwt (0.12 compat) This should allow the module to work on 0.12, while preserving expiry checking (which was not built in to 0.12's util.jwt).
author Matthew Wild <mwild1@gmail.com>
date Tue, 10 Jan 2023 16:34:21 +0000
parents 2ed6f6eeeaba
children
line wrap: on
line source

-- mod_storage_lmdb
-- Copyright (C) 2015 Kim Alvefur
--
-- This file is MIT/X11 licensed.
-- 
-- Depends on lightningdbm
-- https://github.com/shmul/lightningdbm
--
-- luacheck: globals prosody open

local assert = assert;
local select = select;
local xpcall = xpcall;
local traceback = debug.traceback;

local lmdb = require"lightningmdb";
local lfs = require"lfs";
local path = require"util.paths";
local serialization = require"util.serialization";
local serialize = serialization.serialize;
local deserialize = serialization.deserialize;

local function transaction(env, flag, func, ...)
	local args, n_args = {...}, select("#", ...);
	local t = env:txn_begin(nil, flag);
	local function f() return func(t, unpack(args, 1, n_args)); end
	local success, a, b, c = xpcall(f, traceback);
	if not success then
		t:abort();
		return success, a;
	end
	local ok, err = t:commit();
	if not ok then
		return ok, err;
	end
	return success, a, b, c;
end

local function keyvalue_set(t, db, key, value)
	if value ~= nil then
		return assert(t:put(db, key, value, 0));
	else
		return t:del(db, key, value);
	end
end

local function keyvalue_get(t, db, key)
	local data, err = t:get(db, key, 0);
	assert(data or not err, err);
	return data;
end

local drivers = {};
local provider = {};

local keyval = {};
local keyval_mt = { __index = keyval, flags = lmdb.MDB_CREATE };
drivers.keyval = keyval_mt;

function keyval:set(key, value)
	if type(value) == "table" and next(value) == nil then
		value = nil;
	end
	if value ~= nil then
		value = serialize(value);
	end
	return transaction(self.env, 0, keyvalue_set, self.db, key, value);
end

function keyval:get(key)
	local ok, data = transaction(self.env, lmdb.MDB_RDONLY, keyvalue_get, self.db, key);
	if not ok then return ok, data; end
	return deserialize(data);
end

function provider:init(config)
	if config.base_path then
		lfs.mkdir(config.base_path);
	end
	local env = lmdb.env_create();
	env:set_maxdbs(config.maxdbs or 20);
	local env_flags = 0;
	if config.flags then
		for flag in config.flags do
			env_flags = env_flags + assert(lmdb["MDB_"..flag:upper()], "No such flag "..flag);
		end
	end
	env:open(config.base_path or ".", env_flags, tonumber("640", 8));
	self.env = env;
end

function provider:open(store, typ)
	typ = typ or "keyval";
	local driver_mt = drivers[typ];
	if not driver_mt then
		return nil, "unsupported-store";
	end
	local env = self.env;
	local t = env:txn_begin(nil, 0);
	local db = t:dbi_open(store.."_"..typ, driver_mt.flags);
	local ok, err = t:commit();
	if not ok then
		module:log("error", "Could not open database %s_%s: %s", store, typ, tostring(err));
		return ok, err;
	end

	return setmetatable({ env = env, store = store, type = typ, db = db }, driver_mt);
end

if prosody then
	provider:init({
		base_path = path.resolve_relative_path(prosody.paths.data, module.host);
		flags = module:get_option_set("lmdb_flags", {});
		maxdbs = module:get_option_number("lmdb_maxdbs", 20);
	});

	function module.unload() --luacheck: ignore
		provider.env:sync(1);
		provider.env:close();
	end

	module:provides("storage", provider);
else
	return provider;
end