tdb 0.1.0

data/ext/tdb/tdb.c

@@ -0,0 +1,690 @@
+#include "rbtdb.h"
+#include <sys/types.h>
+#include <sys/stat.h>
+#include <fcntl.h>
+#include <errno.h>
+#ifdef HAVE_RUBY_ST_H
+#  include <ruby/st.h>
+#else
+#  include <st.h>
+#endif
+
+static VALUE cTDB, cERR;
+static st_table *exc_hash;
+static VALUE hashes;
+
+/* must be a macro to prevent GC from killing converted 'val's */
+#define TO_TDB_DATA(data,val) do { \
+	StringValue(val); \
+	(data).dptr = (unsigned char *)RSTRING_PTR(val); \
+	(data).dsize = RSTRING_LEN(val); \
+} while (0)
+
+static void init_exc(enum TDB_ERROR ecode, const char *name)
+{
+	VALUE exc = rb_define_class_under(cERR, name, cERR);
+	st_insert(exc_hash, (st_data_t)ecode, (st_data_t)exc);
+}
+
+static void init_errors(void)
+{
+	cERR = rb_define_class_under(cTDB, "ERR", rb_eStandardError);
+	exc_hash = st_init_numtable();
+
+	init_exc(TDB_ERR_CORRUPT, "CORRUPT");
+	init_exc(TDB_ERR_IO, "IO");
+	init_exc(TDB_ERR_LOCK, "LOCK");
+	init_exc(TDB_ERR_OOM, "OOM");
+	init_exc(TDB_ERR_EXISTS, "EXISTS"),
+	init_exc(TDB_ERR_NOLOCK, "NOLOCK");
+	init_exc(TDB_ERR_LOCK_TIMEOUT, "LOCK_TIMEOUT");
+	init_exc(TDB_ERR_EINVAL, "EINVAL");
+	init_exc(TDB_ERR_NOEXIST, "NOEXIST");
+	init_exc(TDB_ERR_RDONLY, "RDONLY");
+#ifdef HAVE_CONST_TDB_ERR_NESTING
+	init_exc(TDB_ERR_NESTING, "NESTING");
+#endif /* HAVE_CONST_TDB_ERR_NESTING */
+}
+
+static void my_raise(struct tdb_context *tdb)
+{
+	enum TDB_ERROR ecode = tdb_error(tdb);
+	const char *str = tdb_errorstr(tdb);
+	VALUE exc;
+
+	switch (ecode) {
+	case TDB_SUCCESS:
+		rb_bug("attempted to raise with no error");
+	case TDB_ERR_CORRUPT:
+	case TDB_ERR_IO:
+	case TDB_ERR_LOCK:
+	case TDB_ERR_OOM:
+	case TDB_ERR_EXISTS:
+	case TDB_ERR_NOLOCK:
+	case TDB_ERR_LOCK_TIMEOUT:
+	case TDB_ERR_EINVAL:
+	case TDB_ERR_NOEXIST:
+	case TDB_ERR_RDONLY:
+#ifdef HAVE_CONST_TDB_ERR_NESTING
+	case TDB_ERR_NESTING:
+#endif /* HAVE_CONST_TDB_ERR_NESTING */
+		if (!st_lookup(exc_hash, (st_data_t)ecode, (st_data_t *)&exc))
+			rb_bug("no-existent exception: %s\n", str);
+	}
+	rb_raise(exc, str);
+}
+
+static void init_hashes(void)
+{
+#define HF(x) \
+rb_hash_aset(hashes,ID2SYM(rb_intern(#x)),ULONG2NUM((unsigned long)rbtdb_##x))
+	HF(murmur1);
+	HF(murmur1_aligned);
+	HF(murmur2);
+	HF(murmur2a);
+	HF(murmur2_neutral);
+	HF(murmur2_aligned);
+	HF(fnv1a);
+	HF(djb2);
+	HF(djb3);
+	HF(jenkins_lookup3);
+	HF(default);
+}
+
+#ifndef HAVE_RB_THREAD_BLOCKING_REGION
+/* (very) partial emulation of the 1.9 rb_thread_blocking_region under 1.8 */
+#  include <rubysig.h>
+typedef VALUE rb_blocking_function_t(void *);
+static VALUE my_tbr(rb_blocking_function_t *fn, void *data)
+{
+	VALUE rv;
+
+	TRAP_BEG;
+	rv = fn(data);
+	TRAP_END;
+
+	return rv;
+}
+#else
+static VALUE my_tbr(rb_blocking_function_t *fn, void *data)
+{
+	return rb_thread_blocking_region(fn, data, RUBY_UBF_IO, 0);
+}
+#endif /* HAVE_RUBY_THREAD_BLOCKING_REGION */
+
+static void gcfree(void *ptr)
+{
+	struct tdb_context *tdb = ptr;
+
+	if (tdb)
+		(void)tdb_close(tdb);
+}
+
+static VALUE alloc(VALUE klass)
+{
+	return Data_Wrap_Struct(klass, NULL, gcfree, NULL);
+}
+
+static struct tdb_context *db(VALUE self, int check_opened)
+{
+	struct tdb_context *tdb;
+
+	Data_Get_Struct(self, struct tdb_context, tdb);
+
+	if (!tdb && check_opened)
+		rb_raise(rb_eIOError, "closed database");
+
+	return tdb;
+}
+
+struct open_args {
+	const char *name;
+	int hash_size;
+	int tdb_flags;
+	int open_flags;
+	mode_t mode;
+	struct tdb_logging_context *log_ctx;
+	tdb_hash_func hash_fn;
+};
+
+static VALUE nogvl_open(void *ptr)
+{
+	struct open_args *o = ptr;
+	struct tdb_context *tdb;
+
+	tdb = tdb_open_ex(o->name, o->hash_size, o->tdb_flags,
+	                  o->open_flags, o->mode, o->log_ctx, o->hash_fn);
+
+	return (VALUE)tdb;
+}
+
+static void set_args(struct open_args *o, VALUE opts)
+{
+	VALUE tmp;
+
+	o->name = NULL;
+	o->hash_size = 0; /* default */
+	o->tdb_flags = TDB_DEFAULT;
+	o->open_flags = O_RDWR | O_CREAT;
+	o->mode = 0666;
+	o->log_ctx = NULL;
+	o->hash_fn = NULL;
+
+	if (NIL_P(opts))
+		return;
+	Check_Type(opts, T_HASH);
+
+	tmp = rb_hash_aref(opts, ID2SYM(rb_intern("hash_size")));
+	if (!NIL_P(tmp))
+		o->hash_size = NUM2INT(tmp);
+
+	tmp = rb_hash_aref(opts, ID2SYM(rb_intern("mode")));
+	if (!NIL_P(tmp))
+		o->mode = NUM2UINT(tmp);
+
+	tmp = rb_hash_aref(opts, ID2SYM(rb_intern("open_flags")));
+	if (!NIL_P(tmp))
+		o->open_flags = NUM2INT(tmp);
+
+	tmp = rb_hash_aref(opts, ID2SYM(rb_intern("tdb_flags")));
+	if (!NIL_P(tmp))
+		o->tdb_flags = NUM2INT(tmp);
+
+	tmp = rb_hash_aref(opts, ID2SYM(rb_intern("hash")));
+	if (!NIL_P(tmp)) {
+		VALUE num = rb_hash_aref(hashes, tmp);
+
+		if (NIL_P(num)) {
+			tmp = rb_inspect(tmp);
+			rb_raise(rb_eArgError,
+			         "`%s' is not a valid hash function",
+			         StringValuePtr(tmp));
+		}
+
+		o->hash_fn = (tdb_hash_func)NUM2ULONG(num);
+	}
+}
+
+static VALUE init(int argc, VALUE *argv, VALUE self)
+{
+	struct tdb_context *tdb = db(self, 0);
+	VALUE path, opts;
+	struct open_args o;
+
+	if (tdb)
+		rb_raise(rb_eRuntimeError, "TDB already initialized");
+	rb_scan_args(argc, argv, "11", &path, &opts);
+	set_args(&o, opts);
+
+	if (NIL_P(path))
+		o.tdb_flags |= TDB_INTERNAL;
+	else
+		o.name = StringValuePtr(path);
+
+	tdb = (struct tdb_context *)my_tbr(nogvl_open, &o);
+	if (!tdb) {
+		switch (errno) {
+		case ENOMEM:
+		case EMFILE:
+		case ENFILE:
+			rb_gc();
+			tdb = (struct tdb_context *)my_tbr(nogvl_open, &o);
+		}
+		if (!tdb)
+			rb_sys_fail("tdb_open_ex");
+	}
+	DATA_PTR(self) = tdb;
+
+	return self;
+}
+
+/* tdb_close can do a lot, including cancel transactions an munmap */
+static VALUE nogvl_close(void *ptr)
+{
+	struct tdb_context *tdb = ptr;
+
+	return (VALUE)tdb_close(tdb);
+}
+
+static VALUE tdbclose(VALUE self)
+{
+	struct tdb_context *tdb = db(self, 1);
+
+	DATA_PTR(self) = NULL;
+
+	if ((int)my_tbr(nogvl_close, tdb) == -1)
+		rb_sys_fail("tdb_close");
+
+	return Qnil;
+}
+
+static VALUE closed(VALUE self)
+{
+	struct tdb_context *tdb = db(self, 0);
+
+	return tdb ? Qfalse : Qtrue;
+}
+
+#ifdef HAVE_RB_THREAD_CALL_WITH_GVL
+/* missing prototype in ruby.h: */
+void *rb_thread_call_with_gvl(void *(*func)(void *), void *data);
+#else
+static void * my_rb_thread_call_with_gvl(void *(*func)(void *), void *data)
+{
+	return (*func)(data);
+}
+#define rb_thread_call_with_gvl my_rb_thread_call_with_gvl
+#endif /* !HAVE_RB_THREAD_CALL_WITH_GVL */
+
+/*
+ * We avoid the extra malloc/free pair enforced by tdb_fetch. We
+ * use tdb_parse_record to give us pointers to (hopefully) mmap-ed
+ * regions and create a String object directly off that region.
+ */
+struct fetch_parse_args {
+	struct tdb_context *tdb;
+	union {
+		TDB_DATA key;
+		TDB_DATA val;
+	} as;
+	VALUE value;
+};
+
+static VALUE str_new_tdb_data(TDB_DATA *val)
+{
+	return rb_str_new((const char *)val->dptr, val->dsize);
+}
+
+static void *gvl_str_new(void *data)
+{
+	struct fetch_parse_args *f = data;
+
+	f->value = str_new_tdb_data(&f->as.val);
+
+	return NULL;
+}
+
+static int fetch_parse(TDB_DATA key, TDB_DATA val, void *data)
+{
+	struct fetch_parse_args *f = data;
+
+	f->as.val = val;
+	(void)rb_thread_call_with_gvl(gvl_str_new, data);
+
+	return 0;
+}
+
+static VALUE nogvl_parse_record(void *ptr)
+{
+	struct fetch_parse_args *f = ptr;
+
+	if (tdb_parse_record(f->tdb, f->as.key, fetch_parse, ptr) == -1)
+		return Qnil;
+
+	return f->value;
+}
+
+static VALUE fetch(VALUE self, VALUE key)
+{
+	struct fetch_parse_args f;
+
+	f.tdb = db(self, 1);
+	TO_TDB_DATA(f.as.key, key);
+	f.value = Qnil;
+
+	return my_tbr(nogvl_parse_record, &f);
+}
+
+struct store_args {
+	struct tdb_context *tdb;
+	TDB_DATA key;
+	TDB_DATA val;
+	int flag;
+};
+
+static VALUE nogvl_store(void *ptr)
+{
+	struct store_args *s = ptr;
+
+	return (VALUE)tdb_store(s->tdb, s->key, s->val, s->flag);
+}
+
+static VALUE rbtdb_store(VALUE self, VALUE key, VALUE val, int flag, int soft)
+{
+	struct store_args s;
+
+	s.tdb = db(self, 1);
+	TO_TDB_DATA(s.key, key);
+	TO_TDB_DATA(s.val, val);
+	s.flag = flag;
+
+	if ((int)my_tbr(nogvl_store, &s) == -1) {
+		if (soft) {
+			int ecode = tdb_error(s.tdb);
+
+			if ((flag == TDB_INSERT) && (ecode == TDB_ERR_EXISTS))
+				return Qnil;
+			if ((flag == TDB_MODIFY) && (ecode == TDB_ERR_NOEXIST))
+				return Qnil;
+		}
+		my_raise(s.tdb);
+	}
+
+	return val;
+}
+
+static VALUE store(VALUE self, VALUE key, VALUE val)
+{
+	return rbtdb_store(self, key, val, 0, 0);
+}
+
+static VALUE insert_bang(VALUE self, VALUE key, VALUE val)
+{
+	return rbtdb_store(self, key, val, TDB_INSERT, 0);
+}
+
+static VALUE insert(VALUE self, VALUE key, VALUE val)
+{
+	return rbtdb_store(self, key, val, TDB_INSERT, 1);
+}
+
+static VALUE modify_bang(VALUE self, VALUE key, VALUE val)
+{
+	return rbtdb_store(self, key, val, TDB_MODIFY, 0);
+}
+
+static VALUE modify(VALUE self, VALUE key, VALUE val)
+{
+	return rbtdb_store(self, key, val, TDB_MODIFY, 1);
+}
+
+struct exists_args {
+	struct tdb_context *tdb;
+	TDB_DATA key;
+};
+
+static VALUE nogvl_exists(void *ptr)
+{
+	struct exists_args *e = ptr;
+
+	return tdb_exists(e->tdb, e->key) == 0 ? Qfalse : Qtrue;
+}
+
+static VALUE has_key(VALUE self, VALUE key)
+{
+	struct exists_args e;
+
+	e.tdb = db(self, 1);
+	TO_TDB_DATA(e.key, key);
+
+	return my_tbr(nogvl_exists, &e);
+}
+
+struct traverse_args {
+	struct tdb_context *tdb;
+	TDB_DATA key;
+	TDB_DATA val;
+	int state;
+};
+
+static VALUE protected_yield(VALUE val)
+{
+	VALUE *kv = (VALUE *)val;
+
+	return rb_yield_values(2, kv[0], kv[1]);
+}
+
+static void *my_yield(void *data)
+{
+	struct traverse_args *t = data;
+	VALUE kv[2];
+
+	kv[0] = str_new_tdb_data(&t->key);
+	kv[1] = str_new_tdb_data(&t->val);
+
+	rb_protect(protected_yield, (VALUE)kv, &t->state);
+
+	return NULL;
+}
+
+static int
+traverse_fn(struct tdb_context *tdb, TDB_DATA key, TDB_DATA val, void *data)
+{
+	struct traverse_args *t = data;
+
+	t->key = key;
+	t->val = val;
+	(void)rb_thread_call_with_gvl(my_yield, t);
+
+	return t->state;
+}
+
+static VALUE nogvl_traverse(void *ptr)
+{
+	struct traverse_args *t = ptr;
+
+	(void)tdb_traverse(t->tdb, traverse_fn, t);
+
+	return Qfalse;
+}
+
+static VALUE each(VALUE self)
+{
+	struct traverse_args t;
+
+	t.tdb = db(self, 1);
+	t.state = 0;
+
+	my_tbr(nogvl_traverse, &t);
+	if (t.state)
+		rb_jump_tag(t.state);
+	return self;
+}
+
+struct delete_args {
+	struct tdb_context *tdb;
+	TDB_DATA key;
+};
+
+static VALUE nogvl_delete(void *ptr)
+{
+	struct delete_args *d = ptr;
+
+	return tdb_delete(d->tdb, d->key) == 0 ? Qtrue : Qfalse;
+}
+
+static VALUE nuke(VALUE self, VALUE key)
+{
+	struct delete_args d;
+
+	d.tdb = db(self, 1);
+	TO_TDB_DATA(d.key, key);
+
+	return my_tbr(nogvl_delete, &d);
+}
+
+static VALUE delete(VALUE self, VALUE key)
+{
+	VALUE rc = fetch(self, key);
+
+	if (! NIL_P(rc))
+		if (nuke(self, key) == Qfalse)
+			return Qnil;
+	return rc;
+}
+
+static VALUE lockall(VALUE self)
+{
+	struct tdb_context *tdb = db(self, 1);
+	if ((int)my_tbr((rb_blocking_function_t *)tdb_lockall, tdb))
+		my_raise(tdb);
+
+	return Qtrue;
+}
+
+static VALUE trylockall(VALUE self)
+{
+	struct tdb_context *tdb = db(self, 1);
+	void *fn = tdb_lockall_nonblock;
+
+	if ((int)my_tbr((rb_blocking_function_t *)fn, tdb)) {
+		if (tdb_error(tdb) == TDB_ERR_LOCK)
+			return Qfalse;
+		my_raise(tdb);
+	}
+	return Qtrue;
+}
+
+static VALUE unlockall(VALUE self)
+{
+	struct tdb_context *tdb = db(self, 1);
+	if ((int)my_tbr((rb_blocking_function_t *)tdb_unlockall, tdb))
+		my_raise(tdb);
+	return Qtrue;
+}
+
+static VALUE lockall_read(VALUE self)
+{
+	struct tdb_context *tdb = db(self, 1);
+	if ((int)my_tbr((rb_blocking_function_t *)tdb_lockall_read, tdb))
+		my_raise(tdb);
+	return Qtrue;
+}
+
+static VALUE trylockall_read(VALUE self)
+{
+	struct tdb_context *tdb = db(self, 1);
+	void *fn = tdb_lockall_read_nonblock;
+	if ((int)my_tbr((rb_blocking_function_t *)fn, tdb)) {
+		if (tdb_error(tdb) == TDB_ERR_LOCK)
+			return Qfalse;
+		my_raise(tdb);
+	}
+	return Qtrue;
+}
+
+static VALUE unlockall_read(VALUE self)
+{
+	struct tdb_context *tdb = db(self, 1);
+	if ((int)my_tbr((rb_blocking_function_t *)tdb_unlockall_read, tdb))
+		my_raise(tdb);
+	return Qtrue;
+}
+
+static VALUE lockall_mark(VALUE self)
+{
+	struct tdb_context *tdb = db(self, 1);
+	if ((int)my_tbr((rb_blocking_function_t *)tdb_lockall_mark, tdb))
+		my_raise(tdb);
+	return Qtrue;
+}
+
+static VALUE lockall_unmark(VALUE self)
+{
+	struct tdb_context *tdb = db(self, 1);
+	if ((int)my_tbr((rb_blocking_function_t *)tdb_lockall_unmark, tdb))
+		my_raise(tdb);
+	return Qtrue;
+}
+
+void Init_tdb_ext(void)
+{
+	cTDB = rb_define_class("TDB", rb_cObject);
+
+	hashes = rb_hash_new();
+        rb_define_const(cTDB, "HASHES", hashes);
+
+	rb_define_alloc_func(cTDB, alloc);
+	rb_include_module(cTDB, rb_mEnumerable);
+
+	rb_define_method(cTDB, "initialize", init, -1);
+	rb_define_method(cTDB, "close", tdbclose, 0);
+	rb_define_method(cTDB, "closed?", closed, 0);
+
+	rb_define_method(cTDB, "fetch", fetch, 1);
+	rb_define_method(cTDB, "[]", fetch, 1);
+	rb_define_method(cTDB, "store", store, 2);
+	rb_define_method(cTDB, "[]=", store, 2);
+	rb_define_method(cTDB, "insert!", insert_bang, 2);
+	rb_define_method(cTDB, "modify!", modify_bang, 2);
+	rb_define_method(cTDB, "insert", insert, 2);
+	rb_define_method(cTDB, "modify", modify, 2);
+
+	rb_define_method(cTDB, "key?", has_key, 1);
+	rb_define_method(cTDB, "has_key?", has_key, 1);
+	rb_define_method(cTDB, "include?", has_key, 1);
+	rb_define_method(cTDB, "member?", has_key, 1);
+	rb_define_method(cTDB, "each", each, 0);
+	rb_define_method(cTDB, "nuke!", nuke, 1);
+	rb_define_method(cTDB, "delete", delete, 1);
+
+	rb_define_method(cTDB, "lockall", lockall, 0);
+	rb_define_method(cTDB, "trylockall", trylockall, 0);
+	rb_define_method(cTDB, "unlockall", unlockall, 0);
+	rb_define_method(cTDB, "lockall_read", lockall_read, 0);
+	rb_define_method(cTDB, "trylockall_read", trylockall_read, 0);
+	rb_define_method(cTDB, "unlockall_read", unlockall_read, 0);
+	rb_define_method(cTDB, "lockall_mark", lockall_mark, 0);
+	rb_define_method(cTDB, "lockall_unmark", lockall_unmark, 0);
+
+	init_errors();
+	init_hashes();
+
+#define tdb_CONST(x) rb_define_const(cTDB, #x, UINT2NUM(TDB_##x))
+
+	/* just a readability place holder */
+	tdb_CONST(DEFAULT);
+
+	/* clear database if we are the only one with it open */
+	tdb_CONST(CLEAR_IF_FIRST);
+
+	/* don't store on disk, use in-memory database */
+	tdb_CONST(INTERNAL);
+
+	/* don't do any locking */
+	tdb_CONST(NOLOCK);
+
+	/* don't use mmap */
+	tdb_CONST(NOMMAP);
+
+	/* convert endian (internal use) */
+	tdb_CONST(CONVERT);
+
+	/* header is big-endian (internal use) */
+	tdb_CONST(BIGENDIAN);
+
+	/* don't use synchronous transactions */
+	tdb_CONST(NOSYNC);
+
+	/* maintain a sequence number */
+	tdb_CONST(SEQNUM);
+
+	/* Activate the per-hashchain freelist, default 5 */
+	tdb_CONST(VOLATILE);
+
+#ifdef TDB_ALLOW_NESTING
+	/* Allow transactions to nest */
+	tdb_CONST(ALLOW_NESTING);
+#endif
+
+#ifdef TDB_DISALLOW_NESTING
+	/* Disallow transactions to nest */
+	tdb_CONST(DISALLOW_NESTING);
+#endif
+
+#ifdef TDB_INCOMPATIBLE_HASH
+	/* Better hashing: can't be opened by tdb < 1.2.6. */
+	tdb_CONST(INCOMPATIBLE_HASH);
+#endif
+}
+
+/*
+ * Document-class: TDB
+ *
+ * <code>
+ *   tdb = TDB.new("/path/to/file", flags => IO::RDWR|IO::CREAT)
+ *   tdb.store("HELLO", "world")
+ *   tdb.fetch("HELLO")  -> "world"
+ *   tdb.delete("HELLO") -> "world"
+ * </code>
+ */