summaryrefslogtreecommitdiff
path: root/lib/rbot/dbhash.rb
blob: 3168f208f1a30c6601310ccc54b1bb7577fd3967 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
begin
  require 'bdb'
rescue Exception => e
  puts "Got exception: "+e
  puts "rbot couldn't load the bdb module, perhaps you need to install it? try: http://www.ruby-lang.org/en/raa-list.rhtml?name=bdb"
  exit 2
end

# make BTree lookups case insensitive
module BDB
  class CIBtree < Btree
    def bdb_bt_compare(a, b)
      if a == nil || b == nil
        debug "CIBTree: WARNING: comparing #{a.inspect} (#{self[a].inspect}) with #{b.inspect} (#{self[b].inspect})"
      end
      (a||'').downcase <=> (b||'').downcase
    end
  end
end

module Irc

  # DBHash is for tying a hash to disk (using bdb).
  # Call it with an identifier, for example "mydata". It'll look for
  # mydata.db, if it exists, it will load and reference that db.
  # Otherwise it'll create and empty db called mydata.db
  class DBHash

    # absfilename:: use +key+ as an actual filename, don't prepend the bot's
    #               config path and don't append ".db"
    def initialize(bot, key, absfilename=false)
      @bot = bot
      @key = key
      if absfilename && File.exist?(key)
        # db already exists, use it
        @db = DBHash.open_db(key)
      elsif File.exist?(@bot.botclass + "/#{key}.db")
        # db already exists, use it
        @db = DBHash.open_db(@bot.botclass + "/#{key}.db")
      elsif absfilename
        # create empty db
        @db = DBHash.create_db(key)
      else
        # create empty db
        @db = DBHash.create_db(@bot.botclass + "/#{key}.db")
      end
    end

    def method_missing(method, *args, &block)
      return @db.send(method, *args, &block)
    end

    def DBHash.create_db(name)
      debug "DBHash: creating empty db #{name}"
      return BDB::Hash.open(name, nil, 
      BDB::CREATE | BDB::EXCL, 0600)
    end

    def DBHash.open_db(name)
      debug "DBHash: opening existing db #{name}"
      return BDB::Hash.open(name, nil, "r+", 0600)
    end

  end


  # DBTree is a BTree equivalent of DBHash, with case insensitive lookups.
  class DBTree
    @@env=nil
    # TODO: make this customizable
    # Note that it must be at least four times lg_bsize
    @@lg_max = 2*1024*1024
    # absfilename:: use +key+ as an actual filename, don't prepend the bot's
    #               config path and don't append ".db"
    def initialize(bot, key, absfilename=false)
      @bot = bot
      @key = key
      if @@env.nil?
        begin
          @@env = BDB::Env.open("#{@bot.botclass}", BDB::INIT_TRANSACTION | BDB::CREATE | BDB::RECOVER, "set_lg_max" => @@lg_max)
          debug "DBTree: environment opened with max log size #{@@env.conf['lg_max']}"
        rescue => e
          debug "DBTree: failed to open environment: #{e}. Retrying ..."
          @@env = BDB::Env.open("#{@bot.botclass}", BDB::INIT_TRANSACTION | BDB::CREATE |  BDB::RECOVER)
        end
        #@@env = BDB::Env.open("#{@bot.botclass}", BDB::CREATE | BDB::INIT_MPOOL | BDB::RECOVER)
      end

      if absfilename && File.exist?(key)
        # db already exists, use it
        @db = DBTree.open_db(key)
      elsif absfilename
        # create empty db
        @db = DBTree.create_db(key)
      elsif File.exist?(@bot.botclass + "/#{key}.db")
        # db already exists, use it
        @db = DBTree.open_db(@bot.botclass + "/#{key}.db")
      else
        # create empty db
        @db = DBTree.create_db(@bot.botclass + "/#{key}.db")
      end
    end

    def method_missing(method, *args, &block)
      return @db.send(method, *args, &block)
    end

    def DBTree.create_db(name)
      debug "DBTree: creating empty db #{name}"
      return @@env.open_db(BDB::CIBtree, name, nil, BDB::CREATE | BDB::EXCL, 0600)
    end

    def DBTree.open_db(name)
      debug "DBTree: opening existing db #{name}"
      return @@env.open_db(BDB::CIBtree, name, nil, "r+", 0600)
    end

    def DBTree.cleanup_logs()
      begin
        debug "DBTree: checkpointing ..."
        @@env.checkpoint
      rescue => e
        debug "Failed: #{e}"
      end
      begin
        debug "DBTree: flushing log ..."
        @@env.log_flush
        logs = @@env.log_archive(BDB::ARCH_ABS)
        debug "DBTree: deleting archivable logs: #{logs.join(', ')}."
        logs.each { |log|
          File.delete(log)
        }
      rescue => e
        debug "Failed: #{e}"
      end
    end

    def DBTree.cleanup_env()
      begin
        debug "DBTree: checking transactions ..."
        has_active_txn = @@env.txn_stat["st_nactive"] > 0
        if has_active_txn
          debug "DBTree: WARNING: not all transactions completed!"
        end
        DBTree.cleanup_logs
        debug "DBTree: closing environment #{@@env}"
        path = @@env.home
        @@env.close
        @@env = nil
        if has_active_txn
          debug "DBTree: keeping file because of incomplete transactions"
        else
          debug "DBTree: cleaning up environment in #{path}"
          BDB::Env.remove("#{path}")
        end
      rescue => e
        debug "Failed: #{e}"
      end
    end

  end

end