switchman 2.2.3 → 3.0.0

data/app/models/switchman/shard.rb

@@ -6,44 +6,24 @@ require 'switchman/environment'
 require 'switchman/errors'
 
 module Switchman
-  class Shard < ::ActiveRecord::Base
+  class Shard < UnshardedRecord
     # ten trillion possible ids per shard. yup.
     IDS_PER_SHARD = 10_000_000_000_000
 
-    CATEGORIES =
-      {
-          # special cased to mean all other models
-          :primary => nil,
-          # special cased to not allow activating a shard other than the default
-          :unsharded => [Shard]
-      }
-    private_constant :CATEGORIES
-    @connection_specification_name = 'unsharded'
-
-    if defined?(::ProtectedAttributes)
-      attr_accessible :default, :name, :database_server
-    end
-
     # only allow one default
-    validates_uniqueness_of :default, :if => lambda { |s| s.default? }
+    validates_uniqueness_of :default, if: ->(s) { s.default? }
 
     after_save :clear_cache
     after_destroy :clear_cache
 
-    after_rollback :on_rollback
-
     scope :primary, -> { where(name: nil).order(:database_server_id, :id).distinct_on(:database_server_id) }
 
     class << self
-      def categories
-        CATEGORIES.keys
+      def sharded_models
+        @sharded_models ||= [::ActiveRecord::Base, UnshardedRecord].freeze
       end
 
-      def default(reload_deprecated = false, reload: false, with_fallback: false)
-        if reload_deprecated
-          reload = reload_deprecated
-          ::ActiveSupport::Deprecation.warn("positional reload parameter to Switchman::Shard.default is deprecated; use `reload: true`")
-        end
+      def default(reload: false, with_fallback: false)
         if !@default || reload
           # Have to create a dummy object so that several key methods still work
           # (it's easier to do this in one place here, and just assume that sharding
@@ -54,74 +34,68 @@ module Switchman
 
           # if we already have a default shard in place, and the caller wants
           # to use it as a fallback, use that instead of the dummy instance
-          if with_fallback && @default
-            default = @default
-          end
+          default = @default if with_fallback && @default
 
           # the first time we need a dummy dummy for re-entrancy to avoid looping on ourselves
           @default ||= default
 
           # Now find the actual record, if it exists; rescue the fake default if the table doesn't exist
           @default = begin
-            find_cached("default_shard") { Shard.where(default: true).take } || default
+            find_cached('default_shard') { Shard.where(default: true).take } || default
           rescue
             default
           end
 
-          # rebuild current shard activations - it might have "another" default shard serialized there
-          active_shards.replace(active_shards.dup.map do |category, shard|
-            shard = Shard.lookup((!shard || shard.default?) ? 'default' : shard.id)
-            [category, shard]
-          end.to_h)
-
-          activate!(primary: @default) if active_shards.empty?
-
           # make sure this is not erroneously cached
-          if @default.database_server.instance_variable_defined?(:@primary_shard)
-            @default.database_server.remove_instance_variable(:@primary_shard)
-          end
+          @default.database_server.remove_instance_variable(:@primary_shard) if @default.database_server.instance_variable_defined?(:@primary_shard)
 
           # and finally, check for cached references to the default shard on the existing connection
-          if ::ActiveRecord::Base.connected? && ::ActiveRecord::Base.connection.shard.default?
-            ::ActiveRecord::Base.connection.shard = @default
+          sharded_models.each do |klass|
+            klass.connection.shard = @default if klass.connected? && klass.connection.shard.default?
           end
         end
         @default
       end
 
-      def current(category = :primary)
-        active_shards[category] || Shard.default
+      def current(klass = ::ActiveRecord::Base)
+        klass ||= ::ActiveRecord::Base
+        klass.connection_pool.shard
       end
 
       def activate(shards)
-        old_shards = activate!(shards)
+        activated_classes = activate!(shards)
         yield
       ensure
-        active_shards.merge!(old_shards) if old_shards
+        activated_classes.each do |klass|
+          klass.connection_pool.shard_stack.pop
+          klass.connected_to_stack.pop
+        end
       end
 
       def activate!(shards)
-        old_shards = nil
-        currently_active_shards = active_shards
-        shards.each do |category, shard|
-          next if category == :unsharded
-          unless currently_active_shards[category] == shard
-            old_shards ||= {}
-            old_shards[category] = currently_active_shards[category]
-            currently_active_shards[category] = shard
-          end
+        activated_classes = []
+        shards.each do |klass, shard|
+          next if klass == UnshardedRecord
+
+          next unless klass.current_shard != shard.database_server.id.to_sym ||
+                      klass.connection_pool.shard != shard
+
+          activated_classes << klass
+          klass.connected_to_stack << { shard: shard.database_server.id.to_sym, klasses: [klass] }
+          klass.connection_pool.shard_stack << shard
         end
-        old_shards
+        activated_classes
       end
 
       def lookup(id)
         id_i = id.to_i
         return current if id_i == current.id || id == 'self'
         return default if id_i == default.id || id.nil? || id == 'default'
+
         id = id_i
-        raise ArgumentError if id == 0
+        raise ArgumentError if id.zero?
 
-        unless cached_shards.has_key?(id)
+        unless cached_shards.key?(id)
           cached_shards[id] = Shard.default.activate do
             find_cached(['shard', id]) { find_by(id: id) }
           end
@@ -129,11 +103,6 @@ module Switchman
         cached_shards[id]
       end
 
-      def preload_cache
-        cached_shards.reverse_merge!(active_shards.values.index_by(&:id))
-        cached_shards.reverse_merge!(all.index_by(&:id))
-      end
-
       def clear_cache
         cached_shards.clear
       end
@@ -141,63 +110,77 @@ module Switchman
       # ==== Parameters
       #
       # * +shards+ - an array or relation of Shards to iterate over
-      # * +categories+ - an array of categories to activate
-      # * +options+ -
-      #    :parallel - true/false to execute in parallel, or a integer of how many
-      #                sub-processes. Note that parallel invocation currently uses
-      #                forking, so should be used sparingly because you cannot get
-      #                results back
-      #    :exception - :ignore, :raise, :defer (wait until the end and raise the first
+      # * +classes+ - an array of classes to activate
+      #    parallel: - true/false to execute in parallel, or a integer of how many
+      #                sub-processes per database server. Note that parallel
+      #                invocation currently uses forking, so should be used sparingly
+      #                because errors are not raised, and you cannot get results back
+      #    max_procs: - only run this many parallel processes at a time
+      #    exception: - :ignore, :raise, :defer (wait until the end and raise the first
       #                error), or a proc
-      def with_each_shard(*args)
-        raise ArgumentError, "wrong number of arguments (#{args.length} for 0...3)" if args.length > 3
+      def with_each_shard(*args, parallel: false, max_procs: nil, exception: :raise, &block)
+        raise ArgumentError, "wrong number of arguments (#{args.length} for 0...2)" if args.length > 2
 
-        unless default.is_a?(Shard)
-          return Array.wrap(yield)
-        end
+        return Array.wrap(yield) unless default.is_a?(Shard)
 
-        options = args.extract_options!
         if args.length == 1
-          if Array === args.first && args.first.first.is_a?(Symbol)
-            categories = args.first
+          if Array === args.first && args.first.first.is_a?(Class)
+            classes = args.first
           else
             scope = args.first
           end
         else
-          scope, categories = args
+          scope, classes = args
         end
 
-        # back-compat
-        options[:parallel] = options.delete(:max_procs) if options.key?(:max_procs)
-
-        parallel = case options[:parallel]
-                     when true
-                       [Environment.cpu_count || 2, 2].min
-                     when false, nil
-                       1
-                     else
-                       options[:parallel]
-                   end
-        options.delete(:parallel)
+        parallel = 1 if parallel == true
+        parallel = 0 if parallel == false || parallel.nil?
 
         scope ||= Shard.all
-        if ::ActiveRecord::Relation === scope && scope.order_values.empty?
-          scope = scope.order(::Arel.sql("database_server_id IS NOT NULL, database_server_id, id"))
-        end
+        scope = scope.order(::Arel.sql('database_server_id IS NOT NULL, database_server_id, id')) if ::ActiveRecord::Relation === scope && scope.order_values.empty?
 
-        if parallel > 1
+        if parallel.positive?
+          max_procs = determine_max_procs(max_procs, parallel)
           if ::ActiveRecord::Relation === scope
             # still need a post-uniq, cause the default database server could be NULL or Rails.env in the db
             database_servers = scope.reorder('database_server_id').select(:database_server_id).distinct.
-                map(&:database_server).compact.uniq
+                               map(&:database_server).compact.uniq
             # nothing to do
-            return if database_servers.count == 0
+            return if database_servers.count.zero?
 
-            scopes = Hash[database_servers.map do |server|
-              [server, server.shards.merge(scope)]
-            end]
+            parallel = [(max_procs.to_f / database_servers.count).ceil, parallel].min if max_procs
+
+            scopes = database_servers.map do |server|
+              server_scope = server.shards.merge(scope)
+              if parallel == 1
+                subscopes = [server_scope]
+              else
+                subscopes = []
+                total = server_scope.count
+                ranges = []
+                server_scope.find_ids_in_ranges(batch_size: (total.to_f / parallel).ceil) do |min, max|
+                  ranges << [min, max]
+                end
+                # create a half-open range on the last one
+                ranges.last[1] = nil
+                ranges.each do |min, max|
+                  subscope = server_scope.where('id>=?', min)
+                  subscope = subscope.where('id<=?', max) if max
+                  subscopes << subscope
+                end
+              end
+              [server, subscopes]
+            end.to_h
           else
             scopes = scope.group_by(&:database_server)
+            if parallel > 1
+              parallel = [(max_procs.to_f / scopes.count).ceil, parallel].min if max_procs
+              scopes = scopes.map do |(server, shards)|
+                [server, shards.in_groups(parallel, false).compact]
+              end.to_h
+            else
+              scopes = scopes.map { |(server, shards)| [server, [shards]] }.to_h
+            end
           end
 
           exception_pipes = []
@@ -208,8 +191,8 @@ module Switchman
           fd_to_name_map = {}
           errors = []
 
-          wait_for_output = lambda do |out_fds, err_fds, fd_to_name_map|
-            ready, _ = IO.select(out_fds + err_fds)
+          wait_for_output = lambda do
+            ready, = IO.select(out_fds + err_fds)
             ready.each do |fd|
               if fd.eof?
                 fd.close
@@ -223,8 +206,9 @@ module Switchman
           end
 
           # only one process; don't bother forking
-          if scopes.length == 1
-            return with_each_shard(scopes.first.last, categories, options) { yield }
+          if scopes.length == 1 && parallel == 1
+            return with_each_shard(scopes.first.last.first, classes, exception: exception,
+                                   &block)
           end
 
           # clear connections prior to forking (no more queries will be executed in the parent,
@@ -232,13 +216,17 @@ module Switchman
           # silly like dealloc'ing prepared statements)
           ::ActiveRecord::Base.clear_all_connections!
 
-          scopes.each do |server, subscope|
-            name = server.id
+          scopes.each do |server, subscopes|
+            subscopes.each_with_index do |subscope, idx|
+              name = if subscopes.length > 1
+                       "#{server.id} #{idx + 1}"
+                     else
+                       server.id
+                     end
 
-            exception_pipe = IO.pipe
-            exception_pipes << exception_pipe
-            pid, io_in, io_out, io_err = Open4.pfork4(lambda do
-              begin
+              exception_pipe = IO.pipe
+              exception_pipes << exception_pipe
+              pid, io_in, io_out, io_err = Open4.pfork4(lambda do
                 Switchman.config[:on_fork_proc]&.call
 
                 # set a pretty name for the process title, up to 128 characters
@@ -246,70 +234,55 @@ module Switchman
                 # was started)
                 # first, simplify the binary name by stripping directories,
                 # then truncate arguments as necessary
-                bin = File.basename($0)  # Process.argv0 doesn't work on Ruby 2.5 (https://bugs.ruby-lang.org/issues/15887)
+                bin = File.basename($0) # Process.argv0 doesn't work on Ruby 2.5 (https://bugs.ruby-lang.org/issues/15887)
                 max_length = 128 - bin.length - name.length - 3
-                args = ARGV.join(" ")
-                if max_length >= 0
-                  args = args[0..max_length]
-                end
-                new_title = [bin, args, name].join(" ")
+                args = ARGV.join(' ')
+                args = args[0..max_length] if max_length >= 0
+                new_title = [bin, args, name].join(' ')
                 Process.setproctitle(new_title)
 
-                with_each_shard(subscope, categories, options) { yield }
+                with_each_shard(subscope, classes, exception: exception, &block)
                 exception_pipe.last.close
-              rescue Exception => e
+              rescue => e
                 begin
                   dumped = Marshal.dump(e)
-                  dumped = nil if dumped.length > 64 * 1024
                 rescue
-                  dumped = nil
-                end
-
-                if dumped.nil?
                   # couldn't dump the exception; create a copy with just
                   # the message and the backtrace
                   e2 = e.class.new(e.message)
-                  backtrace = e.backtrace
-                  # truncate excessively long backtraces
-                  if backtrace.length > 50
-                    backtrace = backtrace[0...25] + ['...'] + backtrace[-25..-1]
-                  end
-                  e2.set_backtrace(backtrace)
+                  e2.set_backtrace(e.backtrace)
                   e2.instance_variable_set(:@active_shards, e.instance_variable_get(:@active_shards))
                   dumped = Marshal.dump(e2)
                 end
-
                 exception_pipe.last.set_encoding(dumped.encoding)
                 exception_pipe.last.write(dumped)
                 exception_pipe.last.flush
                 exception_pipe.last.close
                 exit! 1
+              end)
+              exception_pipe.last.close
+              pids << pid
+              io_in.close # don't care about writing to stdin
+              out_fds << io_out
+              err_fds << io_err
+              pid_to_name_map[pid] = name
+              fd_to_name_map[io_out] = name
+              fd_to_name_map[io_err] = name
+
+              while max_procs && pids.count >= max_procs
+                while max_procs && out_fds.count >= max_procs
+                  # wait for output if we've hit the max_procs limit
+                  wait_for_output.call
+                end
+                # we've gotten all the output from one fd so wait for its child process to exit
+                found_pid, status = Process.wait2
+                pids.delete(found_pid)
+                errors << pid_to_name_map[found_pid] if status.exitstatus != 0
               end
-            end)
-            exception_pipe.last.close
-            pids << pid
-            io_in.close # don't care about writing to stdin
-            out_fds << io_out
-            err_fds << io_err
-            pid_to_name_map[pid] = name
-            fd_to_name_map[io_out] = name
-            fd_to_name_map[io_err] = name
-
-            while pids.count >= parallel
-              while out_fds.count >= parallel
-                # wait for output if we've hit the parallel limit
-                wait_for_output.call(out_fds, err_fds, fd_to_name_map)
-              end
-              # we've gotten all the output from one fd so wait for its child process to exit
-              found_pid, status = Process.wait2
-              pids.delete(found_pid)
-              errors << pid_to_name_map[found_pid] if status.exitstatus != 0
             end
           end
 
-          while out_fds.any? || err_fds.any?
-            wait_for_output.call(out_fds, err_fds, fd_to_name_map)
-          end
+          wait_for_output.call while out_fds.any? || err_fds.any?
           pids.each do |pid|
             _, status = Process.waitpid2(pid)
             errors << pid_to_name_map[pid] if status.exitstatus != 0
@@ -317,67 +290,51 @@ module Switchman
 
           # check for an exception; we only re-raise the first one
           exception_pipes.each do |exception_pipe|
-            begin
-              serialized_exception = exception_pipe.first.read
-              next if serialized_exception.empty?
-              exception = Marshal.load(serialized_exception)
-              raise exception
-            ensure
-              exception_pipe.first.close
-            end
+            serialized_exception = exception_pipe.first.read
+            next if serialized_exception.empty?
+
+            ex = Marshal.load(serialized_exception) # rubocop:disable Security/MarshalLoad
+            raise ex
+          ensure
+            exception_pipe.first.close
           end
 
           unless errors.empty?
-            raise ParallelShardExecError.new("The following subprocesses did not exit cleanly: #{errors.sort.join(", ")}")
+            raise ParallelShardExecError,
+                  "The following subprocesses did not exit cleanly: #{errors.sort.join(', ')}"
           end
+
           return
         end
 
-        categories ||= []
+        classes ||= []
 
         previous_shard = nil
-        close_connections_if_needed = lambda do |shard|
-          # prune the prior connection unless it happened to be the same
-          if previous_shard && shard != previous_shard && !previous_shard.database_server.shareable?
-            previous_shard.activate do
-              ::GuardRail.activated_environments.each do |env|
-                ::GuardRail.activate(env) do
-                  if ::ActiveRecord::Base.connected? && ::ActiveRecord::Base.connection.open_transactions == 0
-                    ::ActiveRecord::Base.connection_pool.current_pool.disconnect!
-                  end
-                end
-              end
-            end
-          end
-        end
-
         result = []
-        exception = nil
+        ex = nil
         scope.each do |shard|
           # shard references a database server that isn't configured in this environment
           next unless shard.database_server
-          close_connections_if_needed.call(shard)
-          shard.activate(*categories) do
-            begin
-              result.concat Array.wrap(yield)
-            rescue
-              case options[:exception]
-              when :ignore
-              when :defer
-                exception ||= $!
-              when Proc
-                options[:exception].call
-              when :raise
-                raise
-              else
-                raise
-              end
+
+          shard.activate(*classes) do
+            result.concat Array.wrap(yield)
+          rescue
+            case exception
+            when :ignore
+              # ignore
+            when :defer
+              ex ||= $!
+            when Proc
+              exception.call
+            # when :raise
+            else
+              raise
             end
           end
           previous_shard = shard
         end
-        close_connections_if_needed.call(Shard.current)
-        raise exception if exception
+        raise ex if ex
+
         result
       end
 
@@ -386,22 +343,22 @@ module Switchman
         array.each do |object|
           partition_object = partition_proc ? partition_proc.call(object) : object
           case partition_object
-            when Shard
-              shard = partition_object
-            when ::ActiveRecord::Base
-              if partition_object.respond_to?(:associated_shards)
-                partition_object.associated_shards.each do |a_shard|
-                  shard_arrays[a_shard] ||= []
-                  shard_arrays[a_shard] << object
-                end
-                next
-              else
-                shard = partition_object.shard
+          when Shard
+            shard = partition_object
+          when ::ActiveRecord::Base
+            if partition_object.respond_to?(:associated_shards)
+              partition_object.associated_shards.each do |a_shard|
+                shard_arrays[a_shard] ||= []
+                shard_arrays[a_shard] << object
               end
-            when Integer, /^\d+$/, /^(\d+)~(\d+)$/
-              local_id, shard = Shard.local_id_for(partition_object)
-              local_id ||= partition_object
-              object = local_id if !partition_proc
+              next
+            else
+              shard = partition_object.shard
+            end
+          when Integer, /^\d+$/, /^(\d+)~(\d+)$/
+            local_id, shard = Shard.local_id_for(partition_object)
+            local_id ||= partition_object
+            object = local_id unless partition_proc
           end
           shard ||= Shard.current
           shard_arrays[shard] ||= []
@@ -410,7 +367,7 @@ module Switchman
         # TODO: use with_each_shard (or vice versa) to get
         # connection management and parallelism benefits
         shard_arrays.inject([]) do |results, (shard, objects)|
-          results.concat shard.activate { Array.wrap(yield objects) }
+          results.concat(shard.activate { Array.wrap(yield objects) })
         end
       end
 
@@ -422,7 +379,7 @@ module Switchman
       # stay as provided.  This assumes no consumer
       # will return a nil value from the block.
       def signed_id_operation(input_id)
-        sign = input_id < 0 ? -1 : 1
+        sign = input_id.negative? ? -1 : 1
         output = yield input_id.abs
         output * sign
       end
@@ -430,9 +387,10 @@ module Switchman
       # converts an AR object, integral id, string id, or string short-global-id to a
       # integral id. nil if it can't be interpreted
       def integral_id_for(any_id)
-        if any_id.is_a?(::Arel::Nodes::Casted)
-          any_id = any_id.val
-        elsif any_id.is_a?(::Arel::Nodes::BindParam) && ::Rails.version >= "5.2"
+        case any_id
+        when ::Arel::Nodes::Casted
+          any_id = any_id.value
+        when ::Arel::Nodes::BindParam
           any_id = any_id.value.value_before_type_cast
         end
 
@@ -443,12 +401,11 @@ module Switchman
           local_id = $2.to_i
           signed_id_operation(local_id) do |id|
             return nil if id > IDS_PER_SHARD
+
             $1.to_i * IDS_PER_SHARD + id
           end
         when Integer, /^-?\d+$/
           any_id.to_i
-        else
-          nil
         end
       end
 
@@ -460,11 +417,12 @@ module Switchman
       def local_id_for(any_id)
         id = integral_id_for(any_id)
         return NIL_NIL_ID unless id
+
         return_shard = nil
         local_id = signed_id_operation(id) do |abs_id|
           if abs_id < IDS_PER_SHARD
             abs_id
-          elsif return_shard = lookup(abs_id / IDS_PER_SHARD)
+          elsif (return_shard = lookup(abs_id / IDS_PER_SHARD))
             abs_id % IDS_PER_SHARD
           else
             return NIL_NIL_ID
@@ -481,8 +439,10 @@ module Switchman
         integral_id = integral_id_for(any_id)
         local_id, shard = local_id_for(integral_id)
         return integral_id unless local_id
+
         shard ||= source_shard
         return local_id if shard == target_shard
+
         shard.global_id_for(local_id)
       end
 
@@ -493,6 +453,7 @@ module Switchman
         local_id, shard = local_id_for(any_id)
         return any_id unless local_id
         return local_id unless shard
+
         "#{shard.id}~#{local_id}"
       end
 
@@ -501,6 +462,7 @@ module Switchman
       def global_id_for(any_id, source_shard = nil)
         id = integral_id_for(any_id)
         return any_id unless id
+
         signed_id_operation(id) do |abs_id|
           if abs_id >= IDS_PER_SHARD
             abs_id
@@ -513,11 +475,59 @@ module Switchman
 
       def shard_for(any_id, source_shard = nil)
         return any_id.shard if any_id.is_a?(::ActiveRecord::Base)
+
         _, shard = local_id_for(any_id)
         shard || source_shard || Shard.current
       end
 
+      # given the provided option, determines whether we need to (and whether
+      # it's possible) to determine a reasonable default.
+      def determine_max_procs(max_procs_input, parallel_input = 2)
+        max_procs = nil
+        if max_procs_input
+          max_procs = max_procs_input.to_i
+          max_procs = nil if max_procs.zero?
+        else
+          return 1 if parallel_input.nil? || parallel_input < 1
+
+          cpus = Environment.cpu_count
+          max_procs = cpus * parallel_input if cpus&.positive?
+        end
+
+        max_procs
+      end
+
       private
+
+      def add_sharded_model(klass)
+        @sharded_models = (sharded_models + [klass]).freeze
+        initialize_sharding
+      end
+
+      def initialize_sharding
+        full_connects_to_hash = DatabaseServer.all.map { |db| [db.id.to_sym, db.connects_to_hash] }.to_h
+        sharded_models.each do |klass|
+          connects_to_hash = full_connects_to_hash.deep_dup
+          if klass == UnshardedRecord
+            # no need to mention other databases for the unsharded category
+            connects_to_hash = { ::Rails.env.to_sym => DatabaseServer.find(nil).connects_to_hash }
+          end
+
+          # prune things we're already connected to
+          if klass.connection_specification_name == klass.name
+            connects_to_hash.each do |(db_name, role_hash)|
+              role_hash.each_key do |role|
+                role_hash.delete(role) if klass.connection_handler.retrieve_connection_pool(
+                  klass.connection_specification_name, role: role, shard: db_name
+                )
+              end
+            end
+          end
+
+          klass.connects_to shards: connects_to_hash
+        end
+      end
+
       # in-process caching
       def cached_shards
         @cached_shards ||= {}.compare_by_identity
@@ -548,10 +558,6 @@ module Switchman
         shard = nil unless shard.database_server
         shard
       end
-
-      def active_shards
-        Thread.current[:active_shards] ||= {}.compare_by_identity
-      end
     end
 
     def name
@@ -565,11 +571,11 @@ module Switchman
 
     def name=(name)
       write_attribute(:name, @name = name)
-      remove_instance_variable(:@name) if name == nil
+      remove_instance_variable(:@name) if name.nil?
     end
 
     def database_server
-      @database_server ||= DatabaseServer.find(self.database_server_id)
+      @database_server ||= DatabaseServer.find(database_server_id)
     end
 
     def database_server=(database_server)
@@ -590,23 +596,21 @@ module Switchman
       Shard.default
     end
 
-    def activate(*categories)
-      shards = hashify_categories(categories)
-      Shard.activate(shards) do
-        yield
-      end
+    def activate(*classes, &block)
+      shards = hashify_classes(classes)
+      Shard.activate(shards, &block)
     end
 
     # for use from console ONLY
-    def activate!(*categories)
-      shards = hashify_categories(categories)
+    def activate!(*classes)
+      shards = hashify_classes(classes)
       Shard.activate!(shards)
       nil
     end
 
     # custom serialization, since shard is self-referential
-    def _dump(depth)
-      self.id.to_s
+    def _dump(_depth)
+      id.to_s
     end
 
     def self._load(str)
@@ -614,45 +618,45 @@ module Switchman
     end
 
     def drop_database
-      raise("Cannot drop the database of the default shard") if self.default?
+      raise('Cannot drop the database of the default shard') if default?
       return unless read_attribute(:name)
 
       begin
-        adapter = self.database_server.config[:adapter]
+        adapter = database_server.config[:adapter]
         sharding_config = Switchman.config || {}
         drop_statement = sharding_config[adapter]&.[](:drop_statement)
         drop_statement ||= sharding_config[:drop_statement]
         if drop_statement
           drop_statement = Array(drop_statement).dup.
-              map { |statement| statement.gsub('%{name}', self.name) }
+                           map { |statement| statement.gsub('%{name}', name) }
         end
 
         case adapter
-          when 'mysql', 'mysql2'
-            self.activate do
-              ::GuardRail.activate(:deploy) do
-                drop_statement ||= "DROP DATABASE #{self.name}"
-                Array(drop_statement).each do |stmt|
-                  ::ActiveRecord::Base.connection.execute(stmt)
-                end
+        when 'mysql', 'mysql2'
+          activate do
+            ::GuardRail.activate(:deploy) do
+              drop_statement ||= "DROP DATABASE #{name}"
+              Array(drop_statement).each do |stmt|
+                ::ActiveRecord::Base.connection.execute(stmt)
               end
             end
-          when 'postgresql'
-            self.activate do
-              ::GuardRail.activate(:deploy) do
-                # Shut up, Postgres!
-                conn = ::ActiveRecord::Base.connection
-                old_proc = conn.raw_connection.set_notice_processor {}
-                begin
-                  drop_statement ||= "DROP SCHEMA #{self.name} CASCADE"
-                  Array(drop_statement).each do |stmt|
-                    ::ActiveRecord::Base.connection.execute(stmt)
-                  end
-                ensure
-                  conn.raw_connection.set_notice_processor(&old_proc) if old_proc
+          end
+        when 'postgresql'
+          activate do
+            ::GuardRail.activate(:deploy) do
+              # Shut up, Postgres!
+              conn = ::ActiveRecord::Base.connection
+              old_proc = conn.raw_connection.set_notice_processor {}
+              begin
+                drop_statement ||= "DROP SCHEMA #{name} CASCADE"
+                Array(drop_statement).each do |stmt|
+                  ::ActiveRecord::Base.connection.execute(stmt)
                 end
+              ensure
+                conn.raw_connection.set_notice_processor(&old_proc) if old_proc
               end
             end
+          end
         end
       rescue
         logger.info "Drop failed: #{$!}"
@@ -662,8 +666,9 @@ module Switchman
     # takes an id local to this shard, and returns a global id
     def global_id_for(local_id)
       return nil unless local_id
+
       self.class.signed_id_operation(local_id) do |abs_id|
-        abs_id + self.id * IDS_PER_SHARD
+        abs_id + id * IDS_PER_SHARD
       end
     end
 
@@ -673,7 +678,8 @@ module Switchman
     end
 
     def destroy
-      raise("Cannot destroy the default shard") if self.default?
+      raise('Cannot destroy the default shard') if default?
+
       super
     end
 
@@ -682,31 +688,22 @@ module Switchman
     def clear_cache
       Shard.default.activate do
         Switchman.cache.delete(['shard', id].join('/'))
-        Switchman.cache.delete("default_shard") if default?
+        Switchman.cache.delete('default_shard') if default?
       end
-      self.class.clear_cache
     end
 
     def default_name
       database_server.shard_name(self)
     end
 
-    def on_rollback
-      # make sure all connection pool proxies are referencing valid pools
-      ::ActiveRecord::Base.connection_handler.connection_pools.each do |pool|
-        next unless pool.is_a?(ConnectionPoolProxy)
-
-        pool.remove_shard!(self)
-      end
-    end
-
-    def hashify_categories(categories)
-      if categories.empty?
-        { :primary => self }
+    def hashify_classes(classes)
+      if classes.empty?
+        { ::ActiveRecord::Base => self }
       else
-        categories.inject({}) { |h, category| h[category] = self; h }
+        classes.each_with_object({}) do |klass, h|
+          h[klass] = self
+        end
       end
     end
-
   end
 end