switchman 2.1.0 → 3.0.6

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,77 @@ 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
+          # Now find the actual record, if it exists
           @default = begin
-            find_cached("default_shard") { Shard.where(default: true).take } || default
+            find_cached('default_shard') { Shard.where(default: true).take } || default
+          # If we are *super* early in boot, the connection pool won't exist; we don't want to fill in the default shard yet
+          rescue ::ActiveRecord::ConnectionNotEstablished
+            nil
+          # rescue the fake default if the table doesn't exist
           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?
+          return default unless @default
 
           # 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.current_switchman_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.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 = nil
+        shards.each do |klass, shard|
+          next if klass == UnshardedRecord
+
+          next unless klass.current_shard != shard.database_server.id.to_sym ||
+                      klass.current_switchman_shard != shard
+
+          (activated_classes ||= []) << klass
+          klass.connected_to_stack << { shard: shard.database_server.id.to_sym, klasses: [klass], switchman_shard: shard }
         end
-        old_shards
+        activated_classes
+      end
+
+      def active_shards
+        sharded_models.map do |klass|
+          [klass, current(klass)]
+        end.compact.to_h
       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,6 +112,11 @@ 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
@@ -136,89 +124,47 @@ 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 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
+      # * +classes+ - an array of classes to activate
+      #    parallel: - true/false to execute in parallel, or an 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
       #                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, 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
 
-        parallel = case options[:parallel]
-                     when true
-                       1
-                     when false, nil
-                       0
-                     else
-                       options[:parallel]
-                   end
-        options.delete(:parallel)
+        parallel = [Environment.cpu_count || 2, 2].min 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 > 0
-          max_procs = determine_max_procs(options.delete(:max_procs), parallel)
+        if parallel > 1
           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
-            parallel = [(max_procs.to_f / database_servers.count).ceil, parallel].min if max_procs
-
-            scopes = Hash[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]
+            return if database_servers.count.zero?
+
+            scopes = database_servers.map do |server|
+              [server, server.shards.merge(scope)]
+            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 = Hash[scopes.map do |(server, shards)|
-                [server, shards.in_groups(parallel, false).compact]
-              end]
-            else
-              scopes = Hash[scopes.map { |(server, shards)| [server, [shards]] }]
-            end
           end
 
           exception_pipes = []
@@ -229,8 +175,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
@@ -244,99 +190,86 @@ module Switchman
           end
 
           # only one process; don't bother forking
-          if scopes.length == 1 && parallel == 1
-            return with_each_shard(scopes.first.last.first, categories, options) { yield }
-          end
+          return with_each_shard(scopes.first.last, classes, exception: exception, &block) if scopes.length == 1
 
           # clear connections prior to forking (no more queries will be executed in the parent,
           # and we want them gone so that we don't accidentally use them post-fork doing something
           # silly like dealloc'ing prepared statements)
           ::ActiveRecord::Base.clear_all_connections!
 
-          scopes.each do |server, subscopes|
-            subscopes.each_with_index do |subscope, idx|
-              if subscopes.length > 1
-                name = "#{server.id} #{idx + 1}"
-              else
-                name = server.id
+          scopes.each do |server, subscope|
+            name = server.id
+
+            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
+              # (we don't actually know the limit, depending on how the process
+              # 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)
+              max_length = 128 - bin.length - name.length - 3
+              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, classes, exception: exception, &block)
+              exception_pipe.last.close
+            rescue Exception => e # rubocop:disable Lint/RescueException
+              begin
+                dumped = Marshal.dump(e)
+                dumped = nil if dumped.length > 64 * 1024
+              rescue
+                dumped = nil
               end
 
-              exception_pipe = IO.pipe
-              exception_pipes << exception_pipe
-              pid, io_in, io_out, io_err = Open4.pfork4(lambda do
-                begin
-                  Switchman.config[:on_fork_proc]&.call
-
-                  # set a pretty name for the process title, up to 128 characters
-                  # (we don't actually know the limit, depending on how the process
-                  # 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)
-                  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(" ")
-                  Process.setproctitle(new_title)
-
-                  with_each_shard(subscope, categories, options) { yield }
-                  exception_pipe.last.close
-                rescue Exception => 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.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
-              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
+                backtrace = backtrace[0...25] + ['...'] + backtrace[-25..] if backtrace.length > 50
+                e2.set_backtrace(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
-              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(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
+              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 pids.count >= parallel
+              while out_fds.count >= parallel
+                # wait for output if we've hit the parallel 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
+            # 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
 
-          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
@@ -344,67 +277,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
 
@@ -413,22 +330,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] ||= []
@@ -437,7 +354,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
 
@@ -449,7 +366,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
@@ -457,9 +374,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
 
@@ -470,12 +388,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
+
+            ($1.to_i * IDS_PER_SHARD) + id
           end
         when Integer, /^-?\d+$/
           any_id.to_i
-        else
-          nil
         end
       end
 
@@ -487,11 +404,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
@@ -508,8 +426,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
 
@@ -520,6 +440,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
 
@@ -528,6 +449,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
@@ -540,29 +462,42 @@ 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 == 0
-        else
-          return 1 if parallel_input.nil? || parallel_input < 1
-          cpus = Environment.cpu_count
-          if cpus && cpus > 0
-            max_procs = cpus * parallel_input
+      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
-        end
 
-        return max_procs
+          # 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
 
-      private
       # in-process caching
       def cached_shards
         @cached_shards ||= {}.compare_by_identity
@@ -593,10 +528,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
@@ -610,11 +541,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)
@@ -635,23 +566,25 @@ module Switchman
       Shard.default
     end
 
-    def activate(*categories)
-      shards = hashify_categories(categories)
-      Shard.activate(shards) do
-        yield
-      end
+    def original_id
+      id
+    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)
@@ -659,45 +592,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: #{$!}"
@@ -707,8 +640,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
 
@@ -718,7 +652,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
 
@@ -727,7 +662,7 @@ 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
@@ -736,22 +671,14 @@ module Switchman
       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