#include #include #include #include "ruby.h" #include "cdb/alloc.h" #include "cdb/uint32.h" #include "cdb/cdb.h" #include "cdb/cdb_make.h" #define Get_CDB(obj, var) Data_Get_Struct(obj, struct cdb, var) #define Get_CDBMake(obj, var) Data_Get_Struct(obj, struct cdb_make, var) #define Check_FD(cdb) if (cdb->fd == -1) rb_raise(rb_eCDB_Error, "file already closed") VALUE rb_eCDB_Error; static void _xread (fd, buf, len) int fd; char *buf; unsigned int len; { int ret; while (len > 0) { do ret = read(fd, buf, len); while ((ret == -1) && (errno == EINTR)); if (ret == -1) rb_sys_fail(0); if (ret == 0) rb_raise(rb_eCDB_Error, "data format error"); len -= ret; buf += ret; } } /* * class CDB */ VALUE rb_cCDB; static void _cdb_free (cdb) struct cdb *cdb; { cdb_free(cdb); if (cdb->fd != -1) close(cdb->fd); free(cdb); } static VALUE rb_cdb_new (klass, file) VALUE klass, file; { int fd; VALUE obj; struct cdb *cdb; Check_Type(file, T_STRING); fd = open(StringValuePtr(file), O_RDONLY, 0); if (fd == -1) rb_sys_fail(0); obj = Data_Make_Struct(klass, struct cdb, 0, _cdb_free, cdb); cdb_init(cdb, fd); return obj; } static VALUE rb_cdb_close (obj) VALUE obj; { struct cdb *cdb; Get_CDB(obj, cdb); Check_FD(cdb); cdb_free(cdb); close(cdb->fd); cdb->fd = -1; return Qnil; } static VALUE rb_cdb_open (klass, file) VALUE klass, file; { VALUE obj; obj = rb_cdb_new(klass, file); return rb_ensure(rb_yield, obj, rb_cdb_close, obj); } static VALUE _cdb_read (cdb, pos, len) struct cdb *cdb; uint32 pos; unsigned int len; { if (cdb->map) { if ((pos > cdb->size) || (cdb->size - pos < len)) rb_raise(rb_eCDB_Error, "data format error"); return rb_str_new(cdb->map + pos, len); } else { int ret; VALUE s; if (seek_set(cdb->fd, pos) == -1) rb_sys_fail(0); s = rb_str_new(0, len); _xread(cdb->fd, RSTRING(s)->ptr, len); return s; } } static VALUE rb_cdb_read (obj, p, l) VALUE obj, p, l; { struct cdb *cdb; uint32 pos = NUM2INT(p); unsigned int len = NUM2INT(l); Get_CDB(obj, cdb); Check_FD(cdb); return _cdb_read(cdb, pos, len); } static VALUE rb_cdb_find (obj, key) VALUE obj, key; { struct cdb *cdb; uint32 pos; Get_CDB(obj, cdb); Check_FD(cdb); StringValue(key); switch (cdb_find(cdb, RSTRING(key)->ptr, RSTRING(key)->len)) { case 1: return _cdb_read(cdb, cdb_datapos(cdb), cdb_datalen(cdb)); case 0: return Qnil; case -1: rb_sys_fail(0); default: rb_raise(rb_eCDB_Error, "cdb_find returned unexpected value"); } } static void _read_uint32 (fd, num) int fd; uint32 *num; { char buf[4]; _xread(fd, buf, 4); uint32_unpack(buf, num); } static VALUE rb_cdb_each (argc, argv, obj) int argc; VALUE *argv, obj; { int ret; VALUE key; struct cdb *cdb; Get_CDB(obj, cdb); Check_FD(cdb); if (rb_scan_args(argc, argv, "01", &key) == 1) { StringValue(key); cdb_findstart(cdb); while ((ret = cdb_findnext(cdb, RSTRING(key)->ptr, RSTRING(key)->len)) == 1) rb_yield(_cdb_read(cdb, cdb_datapos(cdb), cdb_datalen(cdb))); switch (ret) { case 0: return obj; case -1: rb_sys_fail(0); default: rb_raise(rb_eCDB_Error, "cdb_find returned unexpected value"); } } else { int fd; VALUE data; struct stat sb; uint32 eod, klen, dlen, pos; fd = cdb->fd; if (fstat(fd, &sb) == -1) rb_sys_fail(0); if (cdb->map) { char *eodp, *p = cdb->map; uint32_unpack(p, &eod); if (sb.st_size < eod) rb_raise(rb_eCDB_Error, "data format error"); eodp = p + eod; p += 2048; while (p < eodp) { uint32_unpack(p, &klen); p += 4; uint32_unpack(p, &dlen); p += 4; rb_yield(rb_assoc_new(rb_str_new(p, klen), rb_str_new(p + klen, dlen))); p += klen + dlen; } } else { if (seek_set(fd, 0) == -1) rb_sys_fail(0); _read_uint32(fd, &eod); if (sb.st_size < eod) rb_raise(rb_eCDB_Error, "data format error"); if (seek_set(fd, 2048) == -1) rb_sys_fail(0); pos = 2048; while (pos < eod) { _read_uint32(fd, &klen); _read_uint32(fd, &dlen); key = rb_str_new(0, klen); data = rb_str_new(0, dlen); _xread(fd, RSTRING(key)->ptr, klen); _xread(fd, RSTRING(data)->ptr, dlen); rb_yield(rb_assoc_new(key, data)); pos += 8 + klen + dlen; } } return obj; } } /* * class CDBMake */ VALUE rb_cCDBMake; static void _cdbmake_free (cdb) struct cdb_make *cdb; { struct cdb_hplist *head, *next; head = cdb->head; while (head) { next = head->next; alloc_free(head); head = next; } if (cdb->split) alloc_free(cdb->split); if (cdb->fd != -1) close(cdb->fd); free(cdb); } static VALUE rb_cdbmake_new (argc, argv, klass) int argc; VALUE *argv, klass; { int fd; mode_t mode = 0644; VALUE obj, arg1, arg2; struct cdb_make *cdb; if (rb_scan_args(argc, argv, "11", &arg1, &arg2) == 2) mode = NUM2INT(arg2); Check_Type(arg1, T_STRING); fd = open(StringValuePtr(arg1), O_WRONLY|O_CREAT|O_TRUNC, mode); if (fd == -1) rb_sys_fail(0); obj = Data_Make_Struct(klass, struct cdb_make, 0, _cdbmake_free, cdb); if (cdb_make_start(cdb, fd) == -1) rb_sys_fail(0); return obj; } static VALUE rb_cdbmake_finish (obj) VALUE obj; { struct cdb_make *cdb; Get_CDBMake(obj, cdb); Check_FD(cdb); if (cdb_make_finish(cdb) == -1) rb_sys_fail(0); if (fsync(cdb->fd) == -1) rb_sys_fail(0); close(cdb->fd); cdb->fd = -1; return Qnil; } static VALUE rb_cdbmake_open (argc, argv, klass) int argc; VALUE *argv, klass; { VALUE obj; obj = rb_cdbmake_new(argc, argv, klass); return rb_ensure(rb_yield, obj, rb_cdbmake_finish, obj); } static VALUE rb_cdbmake_add (obj, key, data) VALUE obj, key, data; { struct cdb_make *cdb; Get_CDBMake(obj, cdb); Check_FD(cdb); StringValue(key); StringValue(data); if (cdb_make_add(cdb, RSTRING(key)->ptr, RSTRING(key)->len, RSTRING(data)->ptr, RSTRING(data)->len) == -1) rb_sys_fail(0); return obj; } /* * Module Initialization */ void Init_cdb () { rb_cCDB = rb_define_class("CDB", rb_cObject); rb_cCDBMake = rb_define_class("CDBMake", rb_cObject); rb_eCDB_Error = rb_define_class("CDBError", rb_eRuntimeError); rb_include_module(rb_cCDB, rb_mEnumerable); rb_define_singleton_method(rb_cCDB, "new", rb_cdb_new, 1); rb_define_singleton_method(rb_cCDB, "open", rb_cdb_open, 1); rb_define_singleton_method(rb_cCDBMake, "new", rb_cdbmake_new, -1); rb_define_singleton_method(rb_cCDBMake, "open", rb_cdbmake_open, -1); rb_define_method(rb_cCDB, "find", rb_cdb_find, 1); rb_define_method(rb_cCDB, "[]", rb_cdb_find, 1); rb_define_method(rb_cCDB, "has", rb_cdb_find, 1); rb_define_method(rb_cCDB, "each", rb_cdb_each, -1); rb_define_method(rb_cCDB, "read", rb_cdb_read, 2); rb_define_method(rb_cCDB, "close", rb_cdb_close, 0); rb_define_method(rb_cCDBMake, "add", rb_cdbmake_add, 2); rb_define_method(rb_cCDBMake, "[]=", rb_cdbmake_add, 2); rb_define_method(rb_cCDBMake, "finish", rb_cdbmake_finish, 0); }