dk-bdb 0.2.4

data/ext/bdb.h

@@ -0,0 +1,104 @@
+
+#ifndef BDB2_H
+#define BDB2_H
+
+#include <ruby.h>
+
+#ifdef stat
+  #undef stat
+#endif
+
+#ifdef close
+  #undef close
+#endif
+
+#ifdef rename
+  #undef rename
+#endif
+
+//#include <version.h>
+#include <db.h>
+
+#define NOTXN NULL
+
+#ifdef OPEN_MAX
+#define LMAXFD OPEN_MAX
+#else
+#ifdef FOPEN_MAX
+#define LMAXFD FOPEN_MAX
+#endif
+#endif
+#ifndef LMAXFD
+#error "No max fd define available."
+#endif
+
+#define FNLEN 40
+
+#define filename_copy(fp,fv) \
+  strncpy(fp,RSTRING_PTR(fv),FNLEN);
+
+#define filename_dup(fpd,fps) \
+  strncpy(fpd,fps,FNLEN);
+
+typedef struct s_envh {
+  VALUE self;
+  DB_ENV *env;
+  VALUE adb; /* Ruby array holding opened databases */
+  VALUE atxn; /* Ruby array holding open transactions */
+} t_envh;
+
+typedef struct s_dbh {
+  VALUE self;
+  DB *db;
+	int db_opened;
+  VALUE aproc;
+  VALUE sproc; /* key sorting callback */
+  
+  t_envh *env; /* Parent environment, NULL if not opened from one */
+  VALUE adbc; /* Ruby array holding opened cursor */
+  char filename[FNLEN+1];
+} t_dbh;
+
+typedef struct s_dbch {
+  VALUE self;
+  DBC *dbc;
+  t_dbh *db;
+  char filename[FNLEN+1];
+} t_dbch;
+
+typedef struct s_txnh {
+  VALUE self;
+  DB_TXN *txn;
+  t_envh *env;
+} t_txnh;
+
+#define cu(b,m) \
+  rb_define_const(b,#m,UINT2NUM(m))
+
+#define ci(b,m) \
+  rb_define_const(b,#m,INT2NUM(m))
+
+#define cs(b,m) \
+  rb_define_const(b,#m,rb_str_new2(m))
+
+#define simple_set(fname) \
+VALUE db_ ## fname ## _eq(VALUE obj, VALUE v) \
+{ \
+  rb_ivar_set(obj,fv_ ## fname,v); \
+  return obj; \
+} 
+
+#define attr_writer(fname) \
+  VALUE fname ## _writer(VALUE obj, VALUE v) \
+  {					     \
+    rb_ivar_set(obj,fv_ ## fname,v); \
+    return obj; \
+  }
+
+#define attr_reader(fname) \
+  VALUE fname ## _reader(VALUE obj) \
+  { \
+    return rb_ivar_get(obj,fv_ ## fname);  \
+  }
+
+#endif

data/ext/extconf.rb

@@ -0,0 +1,59 @@
+#!/usr/bin/env ruby
+require 'mkmf'
+default_dir = '/usr/local/BerkeleyDB.4.8'
+inc, lib = dir_config('db', "#{default_dir}/include", "#{default_dir}/lib")
+
+versions=%w(db-4.8 db-4.7 db-4.6 db-4.5 db-4.4 db-4.3 db-4.2)
+until versions.empty?
+  (lib_ok = have_library(versions.shift, 'db_version', 'db.h')) && break
+end
+
+def create_header
+  if File.exist?("bdb_aux._c")
+    message("Not writing bdb_aux._c (defines), already exists\n")
+    return
+  end
+  
+  message("Writing bdb_aux._c (defines), this takes a while\n")
+  db_header = $CPPFLAGS.split.select { |f| f =~ /^-I/ }.map { |e| 
+    f = File.join(e[2..-1], 'db.h')
+    File.exists?(f) ? f : nil
+  }.select { |e| e }.first
+  
+  n=0
+  defines=[]
+  File.open(db_header) {|fd|
+    File.open("bdb_aux._c","w") {|hd|
+      hd.puts("/* This file automatically generated by extconf.rb */\n")
+      fd.each_line {|l|
+        if l =~ %r{^#define\s+(DBC?_\w*)\s+([^\/]*)\s*(.*?)(\/\*.*)?$}
+          name = $1
+          value = $2
+          if macro_defined?(name,"#include <db.h>")
+            case value
+            when /^"/
+              hd.print(%Q{cs(mBdb,%s);\n}%[name])
+            when /^\(?(0x|\d)/
+              hd.print(%Q{cu(mBdb,%s);\n}%[name])
+            when /^\(?-/
+              hd.print(%Q{ci(mBdb,%s);\n}%[name])
+            else
+              $stderr.puts "don't know how to handle #{name} #{value.strip}, guessing UINT"
+              hd.print(%Q{cu(mBdb,%s);\n}%[name])
+            end
+            n+=1
+          end
+        end
+      }
+    }
+    message("\nwrote #{n} defines\n")
+  }
+end
+
+if lib_ok
+  create_header
+  create_makefile('bdb')
+else
+  $stderr.puts("cannot create Makefile")
+  exit 1
+end

data/lib/bdb/base.rb

@@ -0,0 +1,68 @@
+require 'bdb'
+require 'tuple'
+require 'bdb/environment'
+require 'bdb/result_set'
+
+class Bdb::Base
+  def initialize(opts)
+    @config  = Bdb::Environment.config.merge(opts)
+    @indexes = {}
+  end
+  attr_reader :indexes
+  
+  def config(config = {})
+    @config.merge!(config)
+  end
+    
+  def index_by(field, opts = {})
+    raise "index on #{field} already exists" if indexes[field]
+    indexes[field] = opts
+  end
+  
+  def path
+    config[:path] || Dir.pwd
+  end
+
+  def environment
+    @environment ||= Bdb::Environment.new(path, self)
+  end
+  
+  def transaction(nested = true, &block)
+    environment.transaction(nested, &block)
+  end
+  
+  def synchronize(&block)
+    environment.synchronize(&block)
+  end
+  
+  def checkpoint(opts = {})
+    environment.checkpoint(opts)
+  end
+
+  def master?
+    environment.master?
+  end
+
+private
+
+  def get_field(field, value)
+    value.kind_of?(Hash) ? value[field] : value.send(field)
+  end
+end
+
+class Object
+  attr_accessor :bdb_locator_key
+end
+
+# Array comparison should try Tuple comparison first.
+class Array
+  cmp = instance_method(:<=>)
+
+  define_method(:<=>) do |other|
+    begin
+      Tuple.dump(self) <=> Tuple.dump(other)
+    rescue TypeError => e
+      cmp.bind(self).call(other)
+    end
+  end
+end

data/lib/bdb/database.rb

@@ -0,0 +1,207 @@
+require 'bdb/base'
+
+class Bdb::Database < Bdb::Base
+  def initialize(name, opts = {})
+    @name = name
+    super(opts)
+  end
+  attr_reader :name
+
+  def db(index = nil)
+    if @db.nil?
+      @db = {}
+      open_flags = master? ? Bdb::DB_CREATE : Bdb::DB_RDONLY
+      transaction(false) do
+        primary_db = environment.env.db
+        primary_db.pagesize = config[:page_size] if config[:page_size]
+        primary_db.open(transaction, name, nil, Bdb::Db::BTREE, open_flags, 0)
+        @db[:primary_key] = primary_db
+
+        indexes.each do |field, opts|
+          index_callback = lambda do |db, key, data|
+            value = Marshal.load(data)
+            index_key = value.kind_of?(Hash) ? value[:field] : value.send(field)
+            if opts[:multi_key] and index_key.kind_of?(Array)
+              # Index multiple keys. If the key is an array, you must wrap it with an outer array.
+              index_key.collect {|k| Tuple.dump(k)}
+            elsif index_key
+              # Index a single key.
+              Tuple.dump(index_key)
+            end
+          end
+          index_db = environment.env.db
+          index_db.flags = Bdb::DB_DUPSORT unless opts[:unique]
+          index_db.pagesize = config[:page_size] if config[:page_size]
+          index_db.open(transaction, "#{name}_by_#{field}", nil, Bdb::Db::BTREE, open_flags, 0)
+          primary_db.associate(transaction, index_db, open_flags, index_callback)
+          @db[field] = index_db
+        end
+      end
+    end
+    @db[index || :primary_key]
+  rescue Bdb::DbError => e
+    # Retry if the database doesn't exist and we are a replication client.
+    if not master? and e.code == Errno::ENOENT::Errno
+      close
+      sleep 1
+      retry
+    else
+      raise(e)
+    end
+  end
+
+  def close
+    return unless @db
+    synchronize do
+      @db.each {|field, db| db.close(0)}
+      @db = nil
+    end
+  end  
+
+  def close_environment
+    environment.close
+  end
+
+  def count(field, key)
+    with_cursor(db(field)) do |cursor|
+      k, v = cursor.get(Tuple.dump(key), nil, Bdb::DB_SET)
+      k ? cursor.count : 0
+    end
+  end
+
+  def get(*keys, &block)
+    opts  = keys.last.kind_of?(Hash) ? keys.pop : {}
+    db    = db(opts[:field])
+    set   = Bdb::ResultSet.new(opts, &block)
+    flags = opts[:modify] ? Bdb::DB_RMW : 0
+    flags = 0 if environment.disable_transactions?
+
+    keys.each do |key|
+      key = get_key(key, opts)
+      if key == :all
+        with_cursor(db) do |cursor|          
+          if opts[:reverse]
+            k,v  = cursor.get(nil, nil, Bdb::DB_LAST | flags)          # Start at the last item.
+            iter = lambda {cursor.get(nil, nil, Bdb::DB_PREV | flags)} # Move backward.
+          else
+            k,v  = cursor.get(nil, nil, Bdb::DB_FIRST | flags)         # Start at the first item.
+            iter = lambda {cursor.get(nil, nil, Bdb::DB_NEXT | flags)} # Move forward.
+          end
+
+          while k
+            set << unmarshal(v, :tuple => k)
+            k,v = iter.call
+          end
+        end
+      elsif key.kind_of?(Range)
+        # Fetch a range of keys.
+        with_cursor(db) do |cursor|
+          first = Tuple.dump(key.first)
+          last  = Tuple.dump(key.last)
+
+          # Return false once we pass the end of the range.
+          cond = key.exclude_end? ? lambda {|k| k < last} : lambda {|k| k <= last}
+          if opts[:reverse]
+            iter = lambda {cursor.get(nil, nil, Bdb::DB_PREV | flags)} # Move backward.
+            
+            # Position the cursor at the end of the range.
+            k,v = cursor.get(last, nil, Bdb::DB_SET_RANGE | flags) || cursor.get(nil, nil, Bdb::DB_LAST | flags)
+            while k and not cond.call(k)
+              k,v = iter.call
+            end
+            
+            cond = lambda {|k| k >= first} # Change the condition to stop when we move past the start.
+          else
+            k,v  = cursor.get(first, nil, Bdb::DB_SET_RANGE | flags)   # Start at the beginning of the range.
+            iter = lambda {cursor.get(nil, nil, Bdb::DB_NEXT | flags)} # Move forward.
+          end
+          
+          while k and cond.call(k)
+            set << unmarshal(v, :tuple => k)
+            k,v = iter.call
+          end
+        end
+      else
+        if (db.flags & Bdb::DB_DUPSORT) == 0
+          synchronize do
+            # There can only be one item for each key.
+            data = db.get(transaction, Tuple.dump(key), nil, flags)
+            set << unmarshal(data, :key => key) if data
+          end
+        else
+          # Have to use a cursor because there may be multiple items with each key.
+          with_cursor(db) do |cursor|
+            k,v = cursor.get(Tuple.dump(key), nil, Bdb::DB_SET | flags)
+            while k
+              set << unmarshal(v, :tuple => k)
+              k,v = cursor.get(nil, nil, Bdb::DB_NEXT_DUP | flags)
+            end
+          end
+        end
+      end
+    end
+    set.results
+  rescue Bdb::ResultSet::LimitReached
+    set.results
+  end
+
+  def [](key)
+    get(key).first
+  end
+
+  def set(key, value, opts = {})
+    synchronize do
+      key   = Tuple.dump(key)
+      value = Marshal.dump(value)
+      flags = opts[:create] ? Bdb::DB_NOOVERWRITE : 0
+      db.put(transaction, key, value, flags)
+      value
+   end
+  end
+
+  def delete(key)
+    synchronize do
+      key = Tuple.dump(key)
+      db.del(transaction, key, 0)
+    end
+  end
+
+  # Deletes all records in the database. Beware!
+  def truncate!
+    synchronize do
+      db.truncate(transaction)
+    end
+  end
+
+  def sync
+    db.sync
+  end
+
+private
+
+  def get_key(key, opts)
+    if opts[:partial] and not key.kind_of?(Range) and not key == :all
+      first = [*key]
+      last  = first + [true]
+      key   = first..last
+    end
+    key
+  end
+
+  def unmarshal(value, opts = {})
+    value = Marshal.load(value)
+    value.bdb_locator_key = opts[:tuple] ? Tuple.load(opts[:tuple]) : [*opts[:key]]
+    value
+  end
+
+  def with_cursor(db)
+    synchronize do
+      begin
+        cursor = db.cursor(transaction, 0)
+        yield(cursor)
+      ensure
+        cursor.close if cursor
+      end
+    end
+  end
+end

data/lib/bdb/environment.rb

@@ -0,0 +1,135 @@
+require 'thread'
+require 'bdb/replication'
+
+class Bdb::Environment
+  @@env = {}
+  def self.new(path, database = nil)
+    # Only allow one environment per path.
+    path = File.expand_path(path)
+    @@env[path] ||= super(path)
+    @@env[path].databases << database if database
+    @@env[path]
+  end
+
+  def initialize(path)
+    @path = path
+  end
+  attr_reader :path
+
+  def self.[](path)
+    new(path)
+  end
+
+  def self.config(config = {})
+    @config ||= {
+      :max_locks    => 5000,
+      :lock_timeout => 30 * 1000 * 1000,
+      :txn_timeout  => 30 * 1000 * 1000,
+      :cache_size   => 1  * 1024 * 1024,
+    }
+    @config.merge!(config)
+  end
+
+  def config(config = {})
+    @config ||= self.class.config
+    @config.merge!(config)
+  end
+
+  include Replication
+  def self.replicate(path, opts)
+    self[path].replicate(opts)
+  end
+
+  def databases
+    @databases ||= []
+  end
+
+  def env
+    if @env.nil?
+      synchronize do
+        @env = Bdb::Env.new(0)
+        if disable_transactions?
+          env_flags = Bdb::DB_CREATE | Bdb::DB_INIT_MPOOL
+        else
+          env_flags = Bdb::DB_CREATE | Bdb::DB_INIT_TXN | Bdb::DB_INIT_LOCK |
+                      Bdb::DB_REGISTER | Bdb::DB_RECOVER | Bdb::DB_INIT_MPOOL | Bdb::DB_THREAD
+
+          env_flags |= Bdb::DB_INIT_REP if replicate?
+        end
+        @env.cachesize = config[:cache_size] if config[:cache_size]
+        @env.set_timeout(config[:txn_timeout],  Bdb::DB_SET_TXN_TIMEOUT)  if config[:txn_timeout]
+        @env.set_timeout(config[:lock_timeout], Bdb::DB_SET_LOCK_TIMEOUT) if config[:lock_timeout]
+        @env.set_lk_max_locks(config[:max_locks]) if config[:max_locks]
+        @env.set_lk_detect(Bdb::DB_LOCK_RANDOM)
+        @env.flags_on = Bdb::DB_TXN_WRITE_NOSYNC | Bdb::DB_TIME_NOTGRANTED
+        init_replication(@env) if replicate?
+
+        @env.open(path, env_flags, 0)
+        start_replication(@env) if replicate?
+        @exit_handler ||= at_exit { close }
+      end
+    end
+    @env
+  end
+
+  def close
+    return unless @env
+    synchronize do
+      databases.each {|database| database.close}
+      @env.close
+      @env = nil
+    end
+  end
+
+  def transaction(nested = true)
+    return @transaction unless block_given?
+    return yield if disable_transactions?
+
+    synchronize do
+      parent = @transaction
+      begin
+        @transaction = env.txn_begin(nested ? parent : nil, 0)
+        value = yield
+        @transaction.commit(0)
+        @transaction = nil
+        value
+      ensure
+        @transaction.abort if @transaction
+        @transaction = parent
+      end
+    end
+  end
+
+  def checkpoint(opts = {})
+    return if disable_transactions?
+    env.txn_checkpoint(opts[:kbyte] || 0, opts[:min] || 0, opts[:force] ? Bdb::DB_FORCE : 0)
+  end
+
+  def disable_transactions?
+    config[:disable_transactions]
+  end
+
+  def synchronize
+    @mutex ||= Mutex.new
+    if @thread_id == thread_id
+      yield
+    else
+      @mutex.synchronize do
+        begin
+          @thread_id = thread_id
+          Thread.exclusive { yield }
+        ensure
+          @thread_id = nil
+        end
+      end
+    end
+  rescue Bdb::DbError => e
+    exit!(9) if e.code == Bdb::DB_RUNRECOVERY    
+    retry if transaction.nil? and e.code == Bdb::DB_LOCK_DEADLOCK
+    raise e
+  end
+
+  def thread_id
+    Thread.current.object_id
+  end
+end

data/lib/bdb/partitioned_database.rb

@@ -0,0 +1,74 @@
+require 'bdb/base'
+
+class Bdb::PartitionedDatabase < Bdb::Base
+  SEPARATOR = '__'
+  PARTITION_PATTERN = /^[-\w]*$/
+
+  def initialize(base_name, opts = {})
+    @base_name    = base_name
+    @partition_by = opts.delete(:partition_by)
+    super(opts)
+  end
+  attr_reader :base_name, :partition_by, :partition
+
+  def databases
+    @databases ||= {}
+  end
+
+  def database(partition = nil)
+    partition ||= self.partition
+    raise 'partition value required' if partition.nil?
+    partition = partition.to_s
+    raise "invalid partition value: #{partition}" unless partition =~ PARTITION_PATTERN
+
+    databases[partition] ||= begin
+      name = [partition, base_name].join(SEPARATOR)
+      database = Bdb::Database.new(name, config)
+      indexes.each do |field, opts|
+        database.index_by(field, opts)
+      end
+      database
+    end
+  end
+
+  def partitions
+    Dir[environment.path + "/*#{SEPARATOR}#{base_name}"].collect do |file|
+      File.basename(file).split(SEPARATOR).first
+    end
+  end
+
+  def with_partition(partition)
+    @partition, old_partition = partition, @partition
+    yield
+  ensure
+    @partition = old_partition
+  end
+
+  def close
+    databases.each do |partition, database|
+      database.close
+    end
+    @databases.clear
+  end
+
+  def get(*keys, &block)
+    opts = keys.last.kind_of?(Hash) ? keys.last : {}
+    database(opts[partition_by]).get(*keys, &block)
+  end
+
+  def set(key, value, opts = {})
+    partition = get_field(partition_by, value)
+    database(partition).set(key, value, opts)
+  end
+
+  def delete(key, opts = {})
+    database(opts[partition_by]).delete(key)
+  end
+
+  # Deletes all records in the database. Beware!
+  def truncate!
+    partitions.each do |partition|
+      database(partition).truncate!
+    end
+  end
+end

data/lib/bdb/replication.rb

@@ -0,0 +1,68 @@
+require 'socket'
+module Replication
+  DEFAULT_PORT = 3463
+  NUM_THREADS  = 1
+
+  ACK_POLICY = {
+    :all       => Bdb::DB_REPMGR_ACKS_ALL,
+    :all_peers => Bdb::DB_REPMGR_ACKS_ALL_PEERS,
+    :none      => Bdb::DB_REPMGR_ACKS_NONE,
+    :one       => Bdb::DB_REPMGR_ACKS_ONE,
+    :one_peer  => Bdb::DB_REPMGR_ACKS_ONE_PEER,
+    :quorom    => Bdb::DB_REPMGR_ACKS_QUORUM,
+  }
+
+  def replicate?
+    not @replicate.nil?
+  end
+
+  def master?
+    not replicate? or replicate[:master]
+  end
+
+  def replicate(opts = nil)
+    return @replicate if opts.nil?
+
+    master  = normalize_host(opts.delete(:from))
+    clients = [*opts.delete(:to)].compact.collect {|h| normalize_host(h)}
+    local   = normalize_host(opts.delete(:host) || ENV['BDB_REPLICATION_HOST'], opts.delete(:port))
+    remote  = clients + [master] - [local]
+
+    opts[:master] = (local == master)
+    opts[:local]  = local
+    opts[:remote] = remote
+    opts[:num_threads] ||= NUM_THREADS
+    @replicate = opts
+
+    env
+  end
+
+private
+
+  def init_replication(env)
+    env.set_verbose(Bdb::DB_VERB_REPLICATION, true) if replicate[:verbose]
+    env.rep_priority = replicate[:master] ? 1 : 0
+    env.repmgr_ack_policy = ACK_POLICY[replicate[:ack_policy]] if replicate[:ack_policy]
+    env.repmgr_set_local_site(*replicate[:local])
+    replicate[:remote].each do |s|
+      env.repmgr_add_remote_site(*s)
+    end
+    env.rep_nsites = replicate[:remote].size + 1
+  end
+
+  def start_replication(env)
+    env.repmgr_start(replicate[:num_threads], replicate[:master] ? Bdb::DB_REP_MASTER : Bdb::DB_REP_CLIENT)
+  end
+
+  def normalize_host(*host)
+    host = host.compact.join(':')
+    host, port = host.split(':')
+    host ||= Socket.gethostname
+    port ||= DEFAULT_PORT
+    port = port.to_i
+
+    addr_info = Socket.getaddrinfo(host.strip, port)
+    ip = addr_info.detect {|i| i[0] == 'AF_INET'}[3]
+    [ip, port]
+  end
+end