switchman 3.0.14 → 3.0.17

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 148d79c6d4c79ef51204c2d32718b30463814b56849effaaf9cacaf86789261d
-  data.tar.gz: 1db4039008bd8958e0dc918b0de81a85294f795f6aea23a1e68e592d6ddc691f
+  metadata.gz: 213adc61b124f8ba48e52259c47702cb601fdc9e0ef71c452c1eeee23dc94e46
+  data.tar.gz: f7982e2e7cb3efce13291382b7fa49f9cb0da87cfd5939581368ac5735ed3e00
 SHA512:
-  metadata.gz: c2ff7f714441d360a38125d7eea85ac3aba981d76fa8a8fceca527f499b7c1a90eb4dd2c91fcf6e292569089e35b077b7ad248cb069a3a9a0579071b93f54a0d
-  data.tar.gz: 7fc113c8b0a1ecdd6d5590401af399fc8f7dd220bc655c190e962870f5e64e2d1e922879e3515693713a9cc77b290913d38740102a29000c19071dab21c10844
+  metadata.gz: e3651c9824e8258c8854e0281ecbc081b42f3fa97524d90eb41ac701f9a5eb0bef53ea77a700c682ffa5913ea71aa7816044cfa408b051f20f4b96b86a4e5307
+  data.tar.gz: 8fac15d77c092790b6410bf15aaccef4c90243308970eddd29e2c9145281cde67e4d6572439b97966ecc2c5924557c7916aca8665e722530c2ead0d43dbcdab2

data/app/models/switchman/shard.rb

@@ -125,8 +125,7 @@ module Switchman
       # * +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
+      #                forking.
       #    exception: - :ignore, :raise, :defer (wait until the end and raise the first
       #                error), or a proc
       def with_each_shard(*args, parallel: false, exception: :raise, &block)
@@ -159,137 +158,47 @@ module Switchman
             return if database_servers.count.zero?
 
             scopes = database_servers.to_h do |server|
-              [server, server.shards.merge(scope)]
+              [server, scope.merge(server.shards)]
             end
           else
             scopes = scope.group_by(&:database_server)
           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
-          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, subscope|
+          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
-
-            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(' ')
+            # 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)
-
-              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
-
-              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
-              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
+              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
-            # 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
-
-          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 errors.first.exception if errors.length == 1
+
             raise ParallelShardExecError,
-                  "The following subprocesses did not exit cleanly: #{errors.sort.join(', ')}"
+                  "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 ||= []
@@ -498,6 +407,7 @@ module Switchman
 
           klass.connects_to shards: connects_to_hash
         end
+        DatabaseServer.all.each { |db| db.guard! if db.config[:prefer_secondary] } unless @sharding_initialized
 
         @sharding_initialized = true
       end

data/lib/switchman/active_record/base.rb

@@ -28,19 +28,11 @@ module Switchman
           if self != ::ActiveRecord::Base && current_scope
             current_scope.activate do
               db = Shard.current(connection_classes).database_server
-              if ::GuardRail.environment == db.guard_rail_environment
-                super
-              else
-                db.unguard { super }
-              end
+              db.unguard { super }
             end
           else
             db = Shard.current(connection_classes).database_server
-            if ::GuardRail.environment == db.guard_rail_environment
-              super
-            else
-              db.unguard { super }
-            end
+            db.unguard { super }
           end
         end
 
@@ -68,6 +60,28 @@ module Switchman
           end
         end
 
+        def current_role_overriden?
+          current_role != current_role(without_overrides: true)
+        end
+
+        # significant change: Allow per-shard roles
+        def current_role(without_overrides: false)
+          return super() if without_overrides
+
+          sharded_role = nil
+          connected_to_stack.reverse_each do |hash|
+            shard_role = hash.dig(:shard_roles, current_shard)
+            if shard_role && (hash[:klasses].include?(Base) || hash[:klasses].include?(connection_classes))
+              sharded_role = shard_role
+              break
+            end
+          end
+          # Allow a shard-specific role to be reverted to regular inheritance
+          return sharded_role if sharded_role && sharded_role != :_switchman_inherit
+
+          super()
+        end
+
         # significant change: _don't_ check if klasses.include?(Base)
         # i.e. other sharded models don't inherit the current shard of Base
         def current_shard
@@ -146,7 +160,8 @@ module Switchman
 
       def with_transaction_returning_status
         shard.activate(self.class.connection_classes) do
-          super
+          db = Shard.current(self.class.connection_classes).database_server
+          db.unguard { super }
         end
       end
 
@@ -167,9 +182,7 @@ module Switchman
 
       def update_columns(*)
         db = shard.database_server
-        return db.unguard { super } if ::GuardRail.environment != db.guard_rail_environment
-
-        super
+        db.unguard { super }
       end
 
       def id_for_database

data/lib/switchman/active_record/persistence.rb

@@ -14,9 +14,7 @@ module Switchman
 
       def delete
         db = shard.database_server
-        return db.unguard { super } unless ::GuardRail.environment == db.guard_rail_environment
-
-        super
+        db.unguard { super }
       end
     end
   end

data/lib/switchman/database_server.rb

@@ -129,27 +129,26 @@ module Switchman
         end
     end
 
-    def guard_rail_environment
-      @guard_rail_environment || ::GuardRail.environment
-    end
-
     # locks this db to a specific environment, except for
     # when doing writes (then it falls back to the current
     # value of GuardRail.environment)
     def guard!(environment = :secondary)
-      @guard_rail_environment = environment
+      ::ActiveRecord::Base.connected_to_stack << { shard_roles: { id.to_sym => environment }, klasses: [::ActiveRecord::Base] }
     end
 
     def unguard!
-      @guard_rail_environment = nil
+      ::ActiveRecord::Base.connected_to_stack << { shard_roles: { id.to_sym => :_switchman_inherit }, klasses: [::ActiveRecord::Base] }
     end
 
     def unguard
-      old_env = @guard_rail_environment
-      unguard!
-      yield
-    ensure
-      guard!(old_env)
+      return yield unless ::ActiveRecord::Base.current_role_overriden?
+
+      begin
+        unguard!
+        yield
+      ensure
+        ::ActiveRecord::Base.connected_to_stack.pop
+      end
     end
 
     def shards

data/lib/switchman/guard_rail/relation.rb

@@ -6,20 +6,17 @@ module Switchman
       def exec_queries(*args)
         if lock_value
           db = Shard.current(connection_classes).database_server
-          return db.unguard { super } if ::GuardRail.environment != db.guard_rail_environment
+          db.unguard { super }
+        else
+          super
         end
-        super
       end
 
       %w[update_all delete_all].each do |method|
         class_eval <<-RUBY, __FILE__, __LINE__ + 1
           def #{method}(*args)
             db = Shard.current(connection_classes).database_server
-            if ::GuardRail.environment != db.guard_rail_environment
-              db.unguard { super }
-            else
-              super
-            end
+            db.unguard { super }
           end
         RUBY
       end

data/lib/switchman/guard_rail.rb

@@ -3,6 +3,11 @@
 module Switchman
   module GuardRail
     module ClassMethods
+      def environment
+        # no overrides so we get the global role, not the role for the default shard
+        ::ActiveRecord::Base.current_role(without_overrides: true)
+      end
+
       def activate(role)
         DatabaseServer.send(:reference_role, role)
         super

data/lib/switchman/parallel.rb

@@ -0,0 +1,68 @@
+# frozen_string_literal: true
+
+require 'parallel'
+
+module Switchman
+  module Parallel
+    module UndumpableException
+      def initialize(original)
+        super
+        @active_shards = original.instance_variable_get(:@active_shards)
+        current_shard
+      end
+    end
+
+    class QuietExceptionWrapper
+      attr_accessor :name
+
+      def initialize(name, wrapper)
+        @name = name
+        @wrapper = wrapper
+      end
+
+      def exception
+        @wrapper.exception
+      end
+    end
+
+    class UndumpableResult
+      attr_reader :name
+
+      def initialize(result)
+        @name = result.inspect
+      end
+
+      def inspect
+        "#<UndumpableResult:#{name}>"
+      end
+    end
+
+    class ResultWrapper
+      attr_reader :result
+
+      def initialize(result)
+        @result =
+          begin
+            Marshal.dump(result) && result
+          rescue
+            UndumpableResult.new(result)
+          end
+      end
+    end
+
+    class PrefixingIO
+      delegate_missing_to :@original_io
+
+      def initialize(prefix, original_io)
+        @prefix = prefix
+        @original_io = original_io
+      end
+
+      def puts(*args)
+        args.flatten.each { |arg| @original_io.puts "#{@prefix}: #{arg}" }
+      end
+    end
+  end
+end
+
+::Parallel::UndumpableException.prepend(::Switchman::Parallel::UndumpableException)

data/lib/switchman/r_spec_helper.rb

@@ -70,7 +70,10 @@ module Switchman
         Shard.default(reload: true)
         puts 'Done!'
 
+        main_pid = Process.pid
         at_exit do
+          next unless main_pid == Process.pid
+
           # preserve rspec's exit status
           status = $!.is_a?(::SystemExit) ? $!.status : nil
           puts 'Tearing down sharding for all specs'

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: @shard_host.pool.connection_klass&.current_role
         }
       end
       super name, payload

data/lib/switchman/standard_error.rb

@@ -6,13 +6,14 @@ module Switchman
       super
       # These seem to get themselves into a bad state if we try to lookup shards while processing
       return if is_a?(IO::EAGAINWaitReadable)
+
       return if Thread.current[:switchman_error_handler]
 
       begin
         Thread.current[:switchman_error_handler] = true
 
         begin
-          @active_shards = Shard.active_shards if defined?(Shard)
+          @active_shards ||= Shard.active_shards if defined?(Shard)
         rescue
           # If we hit an error really early in boot, activerecord may not be initialized yet
         end

data/lib/switchman/version.rb

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

data/lib/switchman.rb

@@ -1,7 +1,7 @@
 # frozen_string_literal: true
 
 require 'guard_rail'
-require 'switchman/open4'
+require 'switchman/parallel'
 require 'switchman/engine'
 
 module Switchman

data/lib/tasks/switchman.rake

@@ -86,10 +86,8 @@ 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

metadata

@@ -1,7 +1,7 @@
 --- !ruby/object:Gem::Specification
 name: switchman
 version: !ruby/object:Gem::Version
-  version: 3.0.14
+  version: 3.0.17
 platform: ruby
 authors:
 - Cody Cutrer
@@ -10,7 +10,7 @@ authors:
 autorequire:
 bindir: bin
 cert_chain: []
-date: 2022-03-23 00:00:00.000000000 Z
+date: 2022-04-05 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: activerecord
@@ -47,19 +47,19 @@ dependencies:
       - !ruby/object:Gem::Version
         version: 3.0.0
 - !ruby/object:Gem::Dependency
-  name: open4
+  name: parallel
   requirement: !ruby/object:Gem::Requirement
     requirements:
     - - "~>"
       - !ruby/object:Gem::Version
-        version: 1.3.0
+        version: '1.22'
   type: :runtime
   prerelease: false
   version_requirements: !ruby/object:Gem::Requirement
     requirements:
     - - "~>"
       - !ruby/object:Gem::Version
-        version: 1.3.0
+        version: '1.22'
 - !ruby/object:Gem::Dependency
   name: railties
   requirement: !ruby/object:Gem::Requirement
@@ -287,7 +287,7 @@ files:
 - lib/switchman/errors.rb
 - lib/switchman/guard_rail.rb
 - lib/switchman/guard_rail/relation.rb
-- lib/switchman/open4.rb
+- lib/switchman/parallel.rb
 - lib/switchman/r_spec_helper.rb
 - lib/switchman/rails.rb
 - lib/switchman/sharded_instrumenter.rb

data/lib/switchman/open4.rb

@@ -1,80 +0,0 @@
-# frozen_string_literal: true
-
-require 'open4'
-
-# This fixes a bug with exception handling,
-# see https://github.com/ahoward/open4/pull/30
-module Open4
-  def self.do_popen(b = nil, exception_propagation_at = nil, closefds=false, &cmd)
-    pw, pr, pe, ps = IO.pipe, IO.pipe, IO.pipe, IO.pipe
-
-    verbose = $VERBOSE
-    begin
-      $VERBOSE = nil
-
-      cid = fork {
-        if closefds
-          exlist = [0, 1, 2] | [pw,pr,pe,ps].map{|p| [p.first.fileno, p.last.fileno] }.flatten
-          ObjectSpace.each_object(IO){|io|
-            io.close if (not io.closed?) and (not exlist.include? io.fileno) rescue nil
-          }
-        end
-
-        pw.last.close
-        STDIN.reopen pw.first
-        pw.first.close
-
-        pr.first.close
-        STDOUT.reopen pr.last
-        pr.last.close
-
-        pe.first.close
-        STDERR.reopen pe.last
-        pe.last.close
-
-        STDOUT.sync = STDERR.sync = true
-
-        begin
-          cmd.call(ps)
-        rescue Exception => e
-          begin
-            Marshal.dump(e, ps.last)
-            ps.last.flush
-          rescue Errno::EPIPE
-            raise e
-          end
-        ensure
-          ps.last.close unless ps.last.closed?
-        end
-
-        exit!
-      }
-    ensure
-      $VERBOSE = verbose
-    end
-
-    [ pw.first, pr.last, pe.last, ps.last ].each { |fd| fd.close }
-
-    Open4.propagate_exception cid, ps.first if exception_propagation_at == :init
-
-    pw.last.sync = true
-
-    pi = [ pw.last, pr.first, pe.first ]
-
-    begin
-      return [cid, *pi] unless b
-
-      begin
-        b.call(cid, *pi)
-      ensure
-        pi.each { |fd| fd.close unless fd.closed? }
-      end
-
-      Open4.propagate_exception cid, ps.first if exception_propagation_at == :block
-
-      Process.waitpid2(cid).last
-    ensure
-      ps.first.close unless ps.first.closed?
-    end
-  end
-end