view mod_storage_lmdb/mod_storage_lmdb.lua @ 5296:0f5657db1cfc

mod_isolate_host: handle server-generated stanzas The hook for setting the no_host_isolation is only called for c2s sessions. This does not work for stanzas generated by the server, such as PEP notifications or presence probe answers. To handle that, we do per-stanza checks for the case that the origin is local.
author Jonas Schäfer <jonas@wielicki.name>
date Sat, 01 Apr 2023 12:03:08 +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