summaryrefslogtreecommitdiff
path: root/lib/rbot/registry/bdb.rb
blob: 200e38ee6532c1d287bc82e5ae82a4be269a6ee5 (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
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
#-- vim:sw=2:et
#++
#
# :title: Berkeley DB interface

begin
  require 'bdb'
rescue LoadError
  fatal "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"
rescue Exception => e
  fatal "rbot couldn't load the bdb module: #{e.pretty_inspect}"
end

if not defined? BDB
  exit 2
end

if BDB::VERSION_MAJOR < 4
  fatal "Your bdb (Berkeley DB) version #{BDB::VERSION} is too old!"
  fatal "rbot will only run with bdb version 4 or higher, please upgrade."
  fatal "For maximum reliability, upgrade to version 4.2 or higher."
  raise BDB::Fatal, BDB::VERSION + " is too old"
end

if BDB::VERSION_MAJOR == 4 and BDB::VERSION_MINOR < 2
  warning "Your bdb (Berkeley DB) version #{BDB::VERSION} may not be reliable."
  warning "If possible, try upgrade version 4.2 or later."
end

# make BTree lookups case insensitive
module BDB
  class CIBtree < Btree
    def bdb_bt_compare(a, b)
      if a == nil || b == nil
        warning "CIBTree: 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
      relfilename = @bot.path key
      relfilename << '.db'
      if absfilename && File.exist?(key)
        # db already exists, use it
        @db = DBHash.open_db(key)
      elsif absfilename
        # create empty db
        @db = DBHash.create_db(key)
      elsif File.exist? relfilename
        # db already exists, use it
        @db = DBHash.open_db relfilename
      else
        # create empty db
        @db = DBHash.create_db relfilename
      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 = 8*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.pretty_inspect}. 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

      relfilename = @bot.path key
      relfilename << '.db'

      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? relfilename
        # db already exists, use it
        @db = DBTree.open_db relfilename
      else
        # create empty db
        @db = DBTree.create_db relfilename
      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 Exception => e
        debug "Failed: #{e.pretty_inspect}"
      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 Exception => e
        debug "Failed: #{e.pretty_inspect}"
      end
    end

    def DBTree.stats()
      begin
        debug "General stats:"
        debug @@env.stat
        debug "Lock stats:"
        debug @@env.lock_stat
        debug "Log stats:"
        debug @@env.log_stat
        debug "Txn stats:"
        debug @@env.txn_stat
      rescue
        debug "Couldn't dump stats"
      end
    end

    def DBTree.cleanup_env()
      begin
        debug "DBTree: checking transactions ..."
        has_active_txn = @@env.txn_stat["st_nactive"] > 0
        if has_active_txn
          warning "DBTree: 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 Exception => e
        error "failed to clean up environment: #{e.pretty_inspect}"
      end
    end

  end

end


module Irc
class Bot

  # This class is now used purely for upgrading from prior versions of rbot
  # the new registry is split into multiple DBHash objects, one per plugin
  class Registry
    def initialize(bot)
      @bot = bot
      upgrade_data
      upgrade_data2
    end

    # check for older versions of rbot with data formats that require updating
    # NB this function is called _early_ in init(), pretty much all you have to
    # work with is @bot.botclass.
    def upgrade_data
      oldreg = @bot.path 'registry.db'
      newreg = @bot.path 'plugin_registry.db'
      if File.exist?(oldreg)
        log _("upgrading old-style (rbot 0.9.5 or earlier) plugin registry to new format")
        old = BDB::Hash.open(oldreg, nil, "r+", 0600)
        new = BDB::CIBtree.open(newreg, nil, BDB::CREATE | BDB::EXCL, 0600)
        old.each {|k,v|
          new[k] = v
        }
        old.close
        new.close
        File.rename(oldreg, oldreg + ".old")
      end
    end

    def upgrade_data2
      oldreg = @bot.path 'plugin_registry.db'
      newdir = @bot.path 'registry'
      if File.exist?(oldreg)
        Dir.mkdir(newdir) unless File.exist?(newdir)
        env = BDB::Env.open(@bot.botclass, BDB::INIT_TRANSACTION | BDB::CREATE | BDB::RECOVER)# | BDB::TXN_NOSYNC)
        dbs = Hash.new
        log _("upgrading previous (rbot 0.9.9 or earlier) plugin registry to new split format")
        old = BDB::CIBtree.open(oldreg, nil, "r+", 0600, "env" => env)
        old.each {|k,v|
          prefix,key = k.split("/", 2)
          prefix.downcase!
          # subregistries were split with a +, now they are in separate folders
          if prefix.gsub!(/\+/, "/")
            # Ok, this code needs to be put in the db opening routines
            dirs = File.dirname("#{@bot.botclass}/registry/#{prefix}.db").split("/")
            dirs.length.times { |i|
              dir = dirs[0,i+1].join("/")+"/"
              unless File.exist?(dir)
                log _("creating subregistry directory #{dir}")
                Dir.mkdir(dir)
              end
            }
          end
          unless dbs.has_key?(prefix)
            log _("creating db #{@bot.botclass}/registry/#{prefix}.db")
            dbs[prefix] = BDB::CIBtree.open("#{@bot.botclass}/registry/#{prefix}.db",
              nil, BDB::CREATE | BDB::EXCL,
              0600, "env" => env)
          end
          dbs[prefix][key] = v
        }
        old.close
        File.rename(oldreg, oldreg + ".old")
        dbs.each {|k,v|
          log _("closing db #{k}")
          v.close
        }
        env.close
      end
    end


  # This class provides persistent storage for plugins via a hash interface.
  # The default mode is an object store, so you can store ruby objects and
  # reference them with hash keys. This is because the default store/restore
  # methods of the plugins' RegistryAccessor are calls to Marshal.dump and
  # Marshal.restore,
  # for example:
  #   blah = Hash.new
  #   blah[:foo] = "fum"
  #   @registry[:blah] = blah
  # then, even after the bot is shut down and disconnected, on the next run you
  # can access the blah object as it was, with:
  #   blah = @registry[:blah]
  # The registry can of course be used to store simple strings, fixnums, etc as
  # well, and should be useful to store or cache plugin data or dynamic plugin
  # configuration.
  #
  # WARNING:
  # in object store mode, don't make the mistake of treating it like a live
  # object, e.g. (using the example above)
  #   @registry[:blah][:foo] = "flump"
  # will NOT modify the object in the registry - remember that Registry#[]
  # returns a Marshal.restore'd object, the object you just modified in place
  # will disappear. You would need to:
  #   blah = @registry[:blah]
  #   blah[:foo] = "flump"
  #   @registry[:blah] = blah
  #
  # If you don't need to store objects, and strictly want a persistant hash of
  # strings, you can override the store/restore methods to suit your needs, for
  # example (in your plugin):
  #   def initialize
  #     class << @registry
  #       def store(val)
  #         val
  #       end
  #       def restore(val)
  #         val
  #       end
  #     end
  #   end
  # Your plugins section of the registry is private, it has its own namespace
  # (derived from the plugin's class name, so change it and lose your data).
  # Calls to registry.each etc, will only iterate over your namespace.
  class Accessor

    attr_accessor :recovery

    # plugins don't call this - a Registry::Accessor is created for them and
    # is accessible via @registry.
    def initialize(bot, name)
      @bot = bot
      @name = name.downcase
      @filename = @bot.path 'registry', @name
      dirs = File.dirname(@filename).split("/")
      dirs.length.times { |i|
        dir = dirs[0,i+1].join("/")+"/"
        unless File.exist?(dir)
          debug "creating subregistry directory #{dir}"
          Dir.mkdir(dir)
        end
      }
      @filename << ".db"
      @registry = nil
      @default = nil
      @recovery = nil
      # debug "initializing registry accessor with name #{@name}"
    end

    def registry
        @registry ||= DBTree.new @bot, "registry/#{@name}"
    end

    def flush
      # debug "fushing registry #{registry}"
      return if !@registry
      registry.flush
      registry.sync
    end

    def close
      # debug "closing registry #{registry}"
      return if !@registry
      registry.close
    end

    # convert value to string form for storing in the registry
    # defaults to Marshal.dump(val) but you can override this in your module's
    # registry object to use any method you like.
    # For example, if you always just handle strings use:
    #   def store(val)
    #     val
    #   end
    def store(val)
      Marshal.dump(val)
    end

    # restores object from string form, restore(store(val)) must return val.
    # If you override store, you should override restore to reverse the
    # action.
    # For example, if you always just handle strings use:
    #   def restore(val)
    #     val
    #   end
    def restore(val)
      begin
        Marshal.restore(val)
      rescue Exception => e
        error _("failed to restore marshal data for #{val.inspect}, attempting recovery or fallback to default")
        debug e
        if defined? @recovery and @recovery
          begin
            return @recovery.call(val)
          rescue Exception => ee
            error _("marshal recovery failed, trying default")
            debug ee
          end
        end
        return default
      end
    end

    # lookup a key in the registry
    def [](key)
      if File.exist?(@filename) && registry.has_key?(key)
        return restore(registry[key])
      else
        return default
      end
    end

    # set a key in the registry
    def []=(key,value)
      registry[key] = store(value)
    end

    # set the default value for registry lookups, if the key sought is not
    # found, the default will be returned. The default default (har) is nil.
    def set_default (default)
      @default = default
    end

    def default
      @default && (@default.dup rescue @default)
    end

    # just like Hash#each
    def each(set=nil, bulk=0, &block)
      return nil unless File.exist?(@filename)
      registry.each(set, bulk) {|key,value|
        block.call(key, restore(value))
      }
    end

    # just like Hash#each_key
    def each_key(set=nil, bulk=0, &block)
      return nil unless File.exist?(@filename)
      registry.each_key(set, bulk) {|key|
        block.call(key)
      }
    end

    # just like Hash#each_value
    def each_value(set=nil, bulk=0, &block)
      return nil unless File.exist?(@filename)
      registry.each_value(set, bulk) { |value|
        block.call(restore(value))
      }
    end

    # just like Hash#has_key?
    def has_key?(key)
      return false unless File.exist?(@filename)
      return registry.has_key?(key)
    end
    alias include? has_key?
    alias member? has_key?
    alias key? has_key?

    # just like Hash#has_both?
    def has_both?(key, value)
      return false unless File.exist?(@filename)
      return registry.has_both?(key, store(value))
    end

    # just like Hash#has_value?
    def has_value?(value)
      return false unless File.exist?(@filename)
      return registry.has_value?(store(value))
    end

    # just like Hash#index?
    def index(value)
      return nil unless File.exist?(@filename)
      ind = registry.index(store(value))
      if ind
        return ind
      else
        return nil
      end
    end

    # delete a key from the registry
    def delete(key)
      return default unless File.exist?(@filename)
      return registry.delete(key)
    end

    # returns a list of your keys
    def keys
      return [] unless File.exist?(@filename)
      return registry.keys
    end

    # Return an array of all associations [key, value] in your namespace
    def to_a
      return [] unless File.exist?(@filename)
      ret = Array.new
      registry.each {|key, value|
        ret << [key, restore(value)]
      }
      return ret
    end

    # Return an hash of all associations {key => value} in your namespace
    def to_hash
      return {} unless File.exist?(@filename)
      ret = Hash.new
      registry.each {|key, value|
        ret[key] = restore(value)
      }
      return ret
    end

    # empties the registry (restricted to your namespace)
    def clear
      return true unless File.exist?(@filename)
      registry.clear
    end
    alias truncate clear

    # returns an array of the values in your namespace of the registry
    def values
      return [] unless File.exist?(@filename)
      ret = Array.new
      self.each {|k,v|
        ret << restore(v)
      }
      return ret
    end

    def sub_registry(prefix)
      return Accessor.new(@bot, @name + "/" + prefix.to_s)
    end

    # returns the number of keys in your registry namespace
    def length
      self.keys.length
    end
    alias size length

  end

  end
end
end