switchman 2.0.10 → 3.0.3

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: bce5f0819ee2570b0bf51ba3872af1f34c23273712cc33c35559225723fb13a0
-  data.tar.gz: 11e7c3d612263f94ea17bc5213580ae42e2c12743a13ff36e9f988feddda0a7e
+  metadata.gz: e5760f08a540406e8e7d546f0104c1ff8833ce7610daba34810b53ffe5caa295
+  data.tar.gz: 4bf65b183425addbec13ae77d9e4f1e3b89244afbcf9625a7e86717e216ce778
 SHA512:
-  metadata.gz: 0fad57a01e86f34533c5cf539cf92360605b129bdfdf77878fec72663c3783bad56e7a7a324f3b388e37a89296d05f309dc2c84084e30c1c44d20917e15d35e3
-  data.tar.gz: 002bcec40325a407c8be94432d2396c8ae812ad2828695266e591956d1e1e8195628906fce710c4681a3c72cc5f6ea8f1859a99aac22cbd3462be5dd6118bbdf
+  metadata.gz: e3e6467f46cb98d78e53b96b1661dbb0cfb59efd25cade9d5f1f2bde4ba067478871c42f91bd06c11c42ed6153c37fca9ef0a20555f3598f910dbda17ce9b42c
+  data.tar.gz: eba93dc04390850eba1a7224bcf5f8fbd5b272133cbac94211c512df023a3e9a7d3951267b55bbc831012f03d443ff6705f74c601bc41a05efb389ecae4240c2

data/Rakefile

@@ -1,4 +1,6 @@
 #!/usr/bin/env rake
+# frozen_string_literal: true
+
 begin
   require 'bundler/setup'
 rescue LoadError
@@ -19,7 +21,7 @@ RDoc::Task.new(:rdoc) do |rdoc|
   rdoc.rdoc_files.include('lib/**/*.rb')
 end
 
-APP_RAKEFILE = File.expand_path("../spec/dummy/Rakefile", __FILE__)
+APP_RAKEFILE = File.expand_path('spec/dummy/Rakefile', __dir__)
 load 'rails/tasks/engine.rake'
 
 Bundler::GemHelper.install_tasks
@@ -27,4 +29,10 @@ Bundler::GemHelper.install_tasks
 require 'rspec/core/rake_task'
 RSpec::Core::RakeTask.new
 
-task :default => :spec
+require 'rubocop/rake_task'
+
+RuboCop::RakeTask.new do |task|
+  task.options = ['-S']
+end
+
+task default: %i[spec rubocop]

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 = @shard_category = :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
@@ -136,59 +110,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
+      # * +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
+      #    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
 
-        parallel = case options[:parallel]
-                     when true
-                       1
-                     when false, nil
-                       0
-                     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 > 0
-          max_procs = determine_max_procs(options.delete(:max_procs), parallel)
+        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?
+
             parallel = [(max_procs.to_f / database_servers.count).ceil, parallel].min if max_procs
 
-            scopes = Hash[database_servers.map do |server|
+            scopes = database_servers.map do |server|
               server_scope = server.shards.merge(scope)
               if parallel == 1
                 subscopes = [server_scope]
@@ -196,28 +158,28 @@ module Switchman
                 subscopes = []
                 total = server_scope.count
                 ranges = []
-                server_scope.find_ids_in_ranges(:batch_size => (total.to_f / parallel).ceil) do |min, max|
+                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
+                  subscope = server_scope.where('id>=?', min)
+                  subscope = subscope.where('id<=?', max) if max
                   subscopes << subscope
                 end
               end
               [server, subscopes]
-            end]
+            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)|
+              scopes = scopes.map do |(server, shards)|
                 [server, shards.in_groups(parallel, false).compact]
-              end]
+              end.to_h
             else
-              scopes = Hash[scopes.map { |(server, shards)| [server, [shards]] }]
+              scopes = scopes.map { |(server, shards)| [server, [shards]] }.to_h
             end
           end
 
@@ -229,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
@@ -245,7 +207,8 @@ module Switchman
 
           # only one process; don't bother forking
           if scopes.length == 1 && parallel == 1
-            return with_each_shard(scopes.first.last.first, categories, options) { yield }
+            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,
@@ -255,51 +218,47 @@ module Switchman
 
           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
-              end
+              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
+                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 => e
                 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 => e
-                  begin
-                    dumped = Marshal.dump(e)
-                  rescue
-                    # couldn't dump the exception; create a copy with just
-                    # the message and the backtrace
-                    e2 = e.class.new(e.message)
-                    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
+                  dumped = Marshal.dump(e)
+                rescue
+                  # couldn't dump the exception; create a copy with just
+                  # the message and the backtrace
+                  e2 = e.class.new(e.message)
+                  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
@@ -313,7 +272,7 @@ module Switchman
               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)
+                  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
@@ -323,9 +282,7 @@ module Switchman
             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
@@ -333,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
 
@@ -402,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] ||= []
@@ -426,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
 
@@ -438,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
@@ -446,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
 
@@ -459,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
 
@@ -476,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
@@ -497,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
 
@@ -509,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
 
@@ -517,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
@@ -529,29 +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)
+      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
+          max_procs = nil if max_procs.zero?
         else
           return 1 if parallel_input.nil? || parallel_input < 1
+
           cpus = Environment.cpu_count
-          if cpus && cpus > 0
-            max_procs = cpus * parallel_input
-          end
+          max_procs = cpus * parallel_input if cpus&.positive?
         end
 
-        return max_procs
+        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
@@ -582,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
@@ -599,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)
@@ -624,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)
@@ -648,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: #{$!}"
@@ -696,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
 
@@ -707,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
 
@@ -716,7 +688,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
@@ -725,22 +697,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