view mod_storage_lmdb/mod_storage_lmdb.lua @ 5354:39d59d857bfb

mod_http_oauth2: Use new mod_cron API for periodic cleanup Less frequent but this isn't that important after all since, as the comment states, expired codes are not usable anyway. They're also not that large so memory usage probably doesn't matter.
author Kim Alvefur <zash@zash.se>
date Sat, 22 Apr 2023 11:59:52 +0200
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