switchman 3.0.2 → 3.1.0

data/{app/models → lib}/switchman/shard.rb

@@ -1,10 +1,5 @@
 # frozen_string_literal: true
 
-require 'switchman/database_server'
-require 'switchman/default_shard'
-require 'switchman/environment'
-require 'switchman/errors'
-
 module Switchman
   class Shard < UnshardedRecord
     # ten trillion possible ids per shard. yup.
@@ -39,7 +34,7 @@ module Switchman
           # 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
           rescue
@@ -59,34 +54,38 @@ module Switchman
 
       def current(klass = ::ActiveRecord::Base)
         klass ||= ::ActiveRecord::Base
-        klass.connection_pool.shard
+        klass.current_switchman_shard
       end
 
       def activate(shards)
         activated_classes = activate!(shards)
         yield
       ensure
-        activated_classes.each do |klass|
-          klass.connection_pool.shard_stack.pop
+        activated_classes&.each do |klass|
           klass.connected_to_stack.pop
         end
       end
 
       def activate!(shards)
-        activated_classes = []
+        activated_classes = nil
         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
+                      klass.current_switchman_shard != shard
 
-          activated_classes << klass
-          klass.connected_to_stack << { shard: shard.database_server.id.to_sym, klasses: [klass] }
-          klass.connection_pool.shard_stack << shard
+          (activated_classes ||= []) << klass
+          klass.connected_to_stack << { shard: shard.database_server.id.to_sym, klasses: [klass], switchman_shard: shard }
         end
         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'
@@ -103,6 +102,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
@@ -111,14 +115,12 @@ module Switchman
       #
       # * +shards+ - an array or relation of Shards to iterate over
       # * +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
+      #    parallel: - true/false to execute in parallel, or an integer of how many
+      #                sub-processes. Note that parallel invocation currently uses
+      #                forking.
       #    exception: - :ignore, :raise, :defer (wait until the end and raise the first
       #                error), or a proc
-      def with_each_shard(*args, parallel: false, max_procs: nil, exception: :raise, &block)
+      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
 
         return Array.wrap(yield) unless default.is_a?(Shard)
@@ -133,14 +135,13 @@ module Switchman
           scope, classes = args
         end
 
-        parallel = 1 if parallel == true
+        parallel = [Environment.cpu_count || 2, 2].min if parallel == true
         parallel = 0 if parallel == false || parallel.nil?
 
         scope ||= Shard.all
         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.positive?
-          max_procs = determine_max_procs(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.
@@ -148,67 +149,11 @@ module Switchman
             # nothing to do
             return if database_servers.count.zero?
 
-            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
+            scopes = database_servers.to_h do |server|
+              [server, scope.merge(server.shards)]
+            end
           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 = []
-          pids = []
-          out_fds = []
-          err_fds = []
-          pid_to_name_map = {}
-          fd_to_name_map = {}
-          errors = []
-
-          wait_for_output = lambda do
-            ready, = IO.select(out_fds + err_fds)
-            ready.each do |fd|
-              if fd.eof?
-                fd.close
-                out_fds.delete(fd)
-                err_fds.delete(fd)
-                next
-              end
-              line = fd.readline
-              puts "#{fd_to_name_map[fd]}: #{line}"
-            end
-          end
-
-          # only one process; don't bother forking
-          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,
@@ -216,95 +161,36 @@ module Switchman
           # silly like dealloc'ing prepared statements)
           ::ActiveRecord::Base.clear_all_connections!
 
-          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
-                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
-                  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
-              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
+          parent_process_name = `ps -ocommand= -p#{Process.pid}`.slice(/#{$0}.*/)
+          ret = ::Parallel.map(scopes, in_processes: scopes.length > 1 ? parallel : 0) do |server, subscope|
+            name = server.id
+            # rubocop:disable Style/GlobalStdStream
+            $stdout = Parallel::PrefixingIO.new(name, STDOUT)
+            $stderr = Parallel::PrefixingIO.new(name, STDERR)
+            # rubocop:enable Style/GlobalStdStream
+            begin
+              max_length = 128 - name.length - 3
+              short_parent_name = parent_process_name[0..max_length] if max_length >= 0
+              new_title = [short_parent_name, name].join(' ')
+              Process.setproctitle(new_title)
+              Switchman.config[:on_fork_proc]&.call
+              with_each_shard(subscope, classes, exception: exception, &block).map { |result| Parallel::ResultWrapper.new(result) }
+            rescue => e
+              logger.error e.full_message
+              Parallel::QuietExceptionWrapper.new(name, ::Parallel::ExceptionWrapper.new(e))
             end
-          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
-          end
-
-          # check for an exception; we only re-raise the first one
-          exception_pipes.each do |exception_pipe|
-            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
+          end.flatten
 
+          errors = ret.select { |val| val.is_a?(Parallel::QuietExceptionWrapper) }
           unless errors.empty?
-            raise ParallelShardExecError,
-                  "The following subprocesses did not exit cleanly: #{errors.sort.join(', ')}"
+            raise errors.first.exception if errors.length == 1
+
+            raise Errors::ParallelShardExecError,
+                  "The following database server(s) did not finish processing cleanly: #{errors.map(&:name).sort.join(', ')}",
+                  cause: errors.first.exception
           end
 
-          return
+          return ret.map(&:result)
         end
 
         classes ||= []
@@ -402,7 +288,7 @@ module Switchman
           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
@@ -480,32 +366,15 @@ module Switchman
         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
+        configure_connects_to
       end
 
-      def initialize_sharding
-        full_connects_to_hash = DatabaseServer.all.map { |db| [db.id.to_sym, db.connects_to_hash] }.to_h
+      def configure_connects_to
+        full_connects_to_hash = DatabaseServer.all.to_h { |db| [db.id.to_sym, db.connects_to_hash] }
         sharded_models.each do |klass|
           connects_to_hash = full_connects_to_hash.deep_dup
           if klass == UnshardedRecord
@@ -596,6 +465,10 @@ module Switchman
       Shard.default
     end
 
+    def original_id
+      id
+    end
+
     def activate(*classes, &block)
       shards = hashify_classes(classes)
       Shard.activate(shards, &block)
@@ -668,7 +541,7 @@ module Switchman
       return nil unless local_id
 
       self.class.signed_id_operation(local_id) do |abs_id|
-        abs_id + id * IDS_PER_SHARD
+        abs_id + (id * IDS_PER_SHARD)
       end
     end
 

data/lib/switchman/sharded_instrumenter.rb

@@ -16,7 +16,7 @@ module Switchman
         payload[:shard] = {
           database_server_id: shard.database_server.id,
           id: shard.id,
-          env: shard.database_server.guard_rail_environment
+          env: ::Rails.version < '7.0' ? @shard_host.pool.connection_klass&.current_role : @shard_host.pool.connection_class&.current_role
         }
       end
       super name, payload

data/lib/switchman/standard_error.rb

@@ -3,20 +3,19 @@
 module Switchman
   module StandardError
     def initialize(*args)
-      # Shard.current can throw this when switchman isn't working right; if we try to
-      # do our stuff here, it'll cause a SystemStackError, which is a pain to deal with
-      if is_a?(::ActiveRecord::ConnectionNotEstablished)
-        super
-        return
-      end
+      super
+      # These seem to get themselves into a bad state if we try to lookup shards while processing
+      return if is_a?(IO::EAGAINWaitReadable)
 
-      if defined?(Shard)
-        @active_shards = Shard.sharded_models.map do |klass|
-          [klass, Shard.current(klass)]
-        end.compact.to_h
-      end
+      return if Thread.current[:switchman_error_handler]
 
-      super
+      begin
+        Thread.current[:switchman_error_handler] = true
+
+        @active_shards ||= Shard.active_shards
+      ensure
+        Thread.current[:switchman_error_handler] = nil
+      end
     end
 
     def current_shard(klass = ::ActiveRecord::Base)

data/{app/models → lib}/switchman/unsharded_record.rb

@@ -2,6 +2,6 @@
 
 module Switchman
   class UnshardedRecord < ::ActiveRecord::Base
-    sharded_model
+    self.abstract_class = true
   end
 end

data/lib/switchman/version.rb

@@ -1,5 +1,5 @@
 # frozen_string_literal: true
 
 module Switchman
-  VERSION = '3.0.2'
+  VERSION = '3.1.0'
 end

data/lib/switchman.rb

@@ -1,8 +1,21 @@
 # frozen_string_literal: true
 
 require 'guard_rail'
-require 'switchman/open4'
-require 'switchman/engine'
+require 'zeitwerk'
+
+class SwitchmanInflector < Zeitwerk::GemInflector
+  def camelize(basename, abspath)
+    if basename =~ /\Apostgresql_(.*)/
+      'PostgreSQL' + super($1, abspath)
+    else
+      super
+    end
+  end
+end
+
+loader = Zeitwerk::Loader.for_gem
+loader.inflector = SwitchmanInflector.new(__FILE__)
+loader.setup
 
 module Switchman
   def self.config
@@ -18,5 +31,12 @@ module Switchman
     @cache = cache
   end
 
+  def self.foreign_key_check(name, type, limit: nil)
+    puts "WARNING: All foreign keys need to be 8-byte integers. #{name} looks like a foreign key. If so, please add the option: `:limit => 8`" if name.to_s =~ /_id\z/ && type.to_s == 'integer' && limit.to_i < 8
+  end
+
   class OrderOnMultiShardQuery < RuntimeError; end
 end
+
+# Load the engine and everything associated at gem load time
+Switchman::Engine

data/lib/tasks/switchman.rake

@@ -1,5 +1,16 @@
 # frozen_string_literal: true
 
+# In rails 7.0+ if you have only 1 db in the env it doesn't try to do explicit activation
+# (and for rails purposes we only have one db per env because each database server is a separate env)
+if Rails.version < '7.0'
+  task_prefix = ::Rake::Task.task_defined?('app:db:migrate') ? 'app:db' : 'db'
+  ::Rake::Task["#{task_prefix}:migrate"].clear_actions.enhance do
+    ::ActiveRecord::Tasks::DatabaseTasks.migrate
+    # Ensure this doesn't blow up when running inside the dummy app
+    Rake::Task["#{task_prefix}:_dump"].invoke
+  end
+end
+
 module Switchman
   module Rake
     def self.filter_database_servers(&block)
@@ -35,9 +46,9 @@ module Switchman
 
       scope = base_scope.order(::Arel.sql('database_server_id IS NOT NULL, database_server_id, id'))
       if servers != DatabaseServer.all
-        conditions = ['database_server_id IN (?)', servers.map(&:id)]
-        conditions.first << ' OR database_server_id IS NULL' if servers.include?(Shard.default.database_server)
-        scope = scope.where(conditions)
+        database_server_ids = servers.map(&:id)
+        database_server_ids << nil if servers.include?(Shard.default.database_server)
+        scope = scope.where(database_server_id: database_server_ids)
       end
 
       scope = shard_scope(scope, shard) if shard
@@ -46,13 +57,11 @@ module Switchman
     end
 
     def self.options
-      { parallel: ENV['PARALLEL'].to_i, max_procs: ENV['MAX_PARALLEL_PROCS'] }
+      { parallel: ENV['PARALLEL'].to_i }
     end
 
     # classes - an array or proc, to activate as the current shard during the
-    # task. tasks which modify the schema may want to pass all categories in
-    # so that schema updates for non-default tables happen against all shards.
-    # this is handled automatically for the default migration tasks, below.
+    # task.
     def self.shardify_task(task_name, classes: [::ActiveRecord::Base])
       old_task = ::Rake::Task[task_name]
       old_actions = old_task.actions.dup
@@ -77,17 +86,15 @@ module Switchman
               nil
             end
           rescue => e
-            puts "Exception from #{e.current_shard.id}: #{e.current_shard.description}" if options[:parallel] != 0
+            warn "Exception from #{e.current_shard.id}: #{e.current_shard.description}:\n#{e.full_message}" if options[:parallel] != 0
             raise
-
-            #::ActiveRecord::Base.configurations = old_configurations
           end
         end
       end
     end
 
     %w[db:migrate db:migrate:up db:migrate:down db:rollback].each do |task_name|
-      shardify_task(task_name, classes: -> { Shard.sharded_models })
+      shardify_task(task_name)
     end
 
     def self.shard_scope(scope, raw_shard_ids)
@@ -201,14 +208,18 @@ module Switchman
     module PostgreSQLDatabaseTasks
       def structure_dump(filename, extra_flags = nil)
         set_psql_env
-        args = ['-s', '-x', '-O', '-f', filename]
+        args = ['--schema-only', '--no-privileges', '--no-owner', '--file', filename]
         args.concat(Array(extra_flags)) if extra_flags
         shard = Shard.current.name
         serialized_search_path = shard
         args << "--schema=#{Shellwords.escape(shard)}"
 
-        args << configuration['database']
+        ignore_tables = ::ActiveRecord::SchemaDumper.ignore_tables
+        args += ignore_tables.flat_map { |table| ['-T', table] } if ignore_tables.any?
+
+        args << db_config.database
         run_cmd('pg_dump', args, 'dumping')
+        remove_sql_header_comments(filename)
         File.open(filename, 'a') { |f| f << "SET search_path TO #{serialized_search_path};\n\n" }
       end
     end