Improved handling of database handling and cleanup during shutdown

4 files changed, 84 insertions, 27 deletions

diff --git a/lib/rbot/dbhash.rb b/lib/rbot/dbhash.rb
index 5433bf90..ed5b9e61 100644
--- a/lib/rbot/dbhash.rb
+++ b/lib/rbot/dbhash.rb
@@ -22,7 +22,7 @@ module Irc
   # 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)
@@ -50,30 +50,39 @@ module Irc
     def DBHash.create_db(name)
       debug "DBHash: creating empty db #{name}"
       return BDB::Hash.open(name, nil, 
-                             BDB::CREATE | BDB::EXCL, 0600)
+      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?
-        #@@env = BDB::Env.open("#{@bot.botclass}", BDB::INIT_TRANSACTION | BDB::CREATE |  BDB::RECOVER)
-        @@env = BDB::Env.open("#{@bot.botclass}", BDB::CREATE | BDB::INIT_MPOOL | BDB::RECOVER)
+        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)
@@ -95,17 +104,43 @@ module Irc
 
     def DBTree.create_db(name)
       debug "DBTree: creating empty db #{name}"
-      return BDB::CIBtree.open(name, nil, 
-                             BDB::CREATE | BDB::EXCL,
-                             0600, "env" => @@env)
+      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 BDB::CIBtree.open(name, nil, 
-                             "r+", 0600, "env" => @@env)
+      return @@env.open_db(BDB::CIBtree, name, nil, "r+", 0600)
     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
+        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)
+        }
+        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
diff --git a/lib/rbot/ircbot.rb b/lib/rbot/ircbot.rb
index 981a09a9..9c682949 100644
--- a/lib/rbot/ircbot.rb
+++ b/lib/rbot/ircbot.rb
@@ -463,15 +463,25 @@ class IrcBot
       debug "failed to trap signals, probably running on windows?"
     end
     message = @lang.get("quit") if (message.nil? || message.empty?)
+    debug "Clearing socket"
     @socket.clearq
+    debug "Saving"
     save
+    debug "Cleaning up"
     @plugins.cleanup
+    debug "Logging quits"
     @channels.each_value {|v|
       log "@ quit (#{message})", v.name
     }
-    @registry.close
+    # debug "Closing registries"
+    # @registry.close
+    debug "Cleaning up the db environment"
+    DBTree.cleanup_env
+    debug "Sending quit message"
     @socket.puts "QUIT :#{message}"
+    debug "Flushing socket"
     @socket.flush
+    debug "Shutting down socket"
     @socket.shutdown
     puts "rbot quit (#{message})"
   end
@@ -570,7 +580,8 @@ class IrcBot
   def status
     secs_up = Time.new - @startup_time
     uptime = Utils.secs_to_string secs_up
-    return "Uptime #{uptime}, #{@plugins.length} plugins active, #{@registry.length} items stored in registry, #{@socket.lines_sent} lines sent, #{@socket.lines_received} received."
+    # return "Uptime #{uptime}, #{@plugins.length} plugins active, #{@registry.length} items stored in registry, #{@socket.lines_sent} lines sent, #{@socket.lines_received} received."
+    return "Uptime #{uptime}, #{@plugins.length} plugins active, #{@socket.lines_sent} lines sent, #{@socket.lines_received} received."
   end
 
   # we'll ping the server every 30 seconds or so, and expect a response
diff --git a/lib/rbot/plugins.rb b/lib/rbot/plugins.rb
index 54b814e0..e0425900 100644
--- a/lib/rbot/plugins.rb
+++ b/lib/rbot/plugins.rb
@@ -16,7 +16,7 @@ module Plugins
   #      def karma_stats(m, params)
   #        m.reply "..."
   #      end
-  #      
+  #
   #      # the default action is the first component
   #      plugin.map 'karma'
   #
@@ -29,10 +29,10 @@ module Plugins
   #        item = params[:key]
   #        m.reply 'karma for #{item}'
   #      end
-  #      
+  #
   #      # you can setup defaults, to make parameters optional
   #      plugin.map 'karma :key', :defaults => {:key => 'defaultvalue'}
-  #      
+  #
   #      # the default auth check is also against the first component
   #      # but that can be changed
   #      plugin.map 'karmastats', :auth => 'karma'
@@ -46,14 +46,14 @@ module Plugins
   #    plugin.register_maps
   #    This also sets the privmsg handler to use the map lookups for
   #    handling messages. You can still use listen(), kick() etc methods
-  # 
+  #
   # listen(UserMessage)::
   #                        Called for all messages of any type. To
   #                        differentiate them, use message.kind_of? It'll be
   #                        either a PrivMessage, NoticeMessage, KickMessage,
   #                        QuitMessage, PartMessage, JoinMessage, NickMessage,
   #                        etc.
-  #                              
+  #
   # privmsg(PrivMessage)::
   #                        called for a PRIVMSG if the first word matches one
   #                        the plugin register()d for. Use m.plugin to get
@@ -70,7 +70,7 @@ module Plugins
   # part(PartMessage)::
   #                        Called when a user (or the bot) parts a channel
   #
-  # quit(QuitMessage)::    
+  # quit(QuitMessage)::
   #                        Called when a user (or the bot) quits IRC
   #
   # nick(NickMessage)::
@@ -81,7 +81,7 @@ module Plugins
   #
   # connect()::            Called when a server is joined successfully, but
   #                        before autojoin channels are joined (no params)
-  # 
+  #
   # save::                 Called when you are required to save your plugin's
   #                        state, if you maintain data between sessions
   #
@@ -100,9 +100,15 @@ module Plugins
     end
 
     def flush_registry
+      # debug "Flushing #{@registry}"
       @registry.flush
     end
 
+    def cleanup
+      # debug "Closing #{@registry}"
+      @registry.close
+    end
+
     def map(*args)
       @handler.map(*args)
       # register this map
@@ -120,7 +126,7 @@ module Plugins
     def name
       @names.join("|")
     end
-    
+
     # return a help string for your module. for complex modules, you may wish
     # to break your help into topics, and return a list of available topics if
     # +topic+ is nil. +plugin+ is passed containing the matching prefix for
@@ -129,7 +135,7 @@ module Plugins
     def help(plugin, topic)
       "no help"
     end
-    
+
     # register the plugin as a handler for messages prefixed +name+
     # this can be called multiple times for a plugin to handle multiple
     # message prefixes
@@ -164,7 +170,7 @@ module Plugins
       @dirs = dirlist
       scan
     end
-    
+
     # access to associated bot
     def Plugins.bot
       @@bot
@@ -198,7 +204,7 @@ module Plugins
             # the idea here is to prevent namespace pollution. perhaps there
             # is another way?
             plugin_module = Module.new
-            
+
             begin
               plugin_string = IO.readlines(tmpfilename).join("")
               debug "loading plugin #{tmpfilename}"
@@ -266,7 +272,7 @@ module Plugins
         end
       end
     end
-    
+
     # see if each plugin handles +method+, and if so, call it, passing
     # +message+ as a parameter
     def delegate(method, *args)
diff --git a/lib/rbot/registry.rb b/lib/rbot/registry.rb
index 24093311..580bfb5b 100644
--- a/lib/rbot/registry.rb
+++ b/lib/rbot/registry.rb
@@ -119,6 +119,11 @@ module Irc
 
     def flush
       @registry.flush
+      @registry.sync
+    end
+
+    def close
+      @registry.close
     end
 
     # convert value to string form for storing in the registry