switchman 3.0.1 → 4.2.5

data/lib/switchman/environment.rb

@@ -1,10 +1,10 @@
 # frozen_string_literal: true
 
-require 'etc'
+require "etc"
 
 module Switchman
   class Environment
-    def self.cpu_count(nproc_bin = 'nproc')
+    def self.cpu_count(nproc_bin = "nproc")
       return Etc.nprocessors if Etc.respond_to?(:nprocessors)
 
       `#{nproc_bin}`.to_i

data/lib/switchman/errors.rb

@@ -1,7 +1,22 @@
 # frozen_string_literal: true
 
 module Switchman
-  class NonExistentShardError < RuntimeError; end
+  module Errors
+    class ManuallyCreatedShadowRecordError < RuntimeError
+      DEFAULT_MSG = "It looks like you're trying to manually create a shadow record. " \
+                    "Please use Switchman::ActiveRecord::Base#save_shadow_record instead."
 
-  class ParallelShardExecError < RuntimeError; end
+      def initialize(msg = DEFAULT_MSG)
+        super
+      end
+    end
+
+    class NonExistentShardError < RuntimeError; end
+
+    class ParallelShardExecError < RuntimeError; end
+
+    class ShadowRecordError < RuntimeError; end
+
+    class UnshardedTableError < RuntimeError; end
+  end
 end

data/lib/switchman/guard_rail/relation.rb

@@ -5,21 +5,18 @@ module Switchman
     module Relation
       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 = Shard.current(connection_class_for_self).database_server
+          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
+          def #{method}(*args, **kwargs)
+            db = Shard.current(connection_class_for_self).database_server
+            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 TransformingIO
+      delegate_missing_to :@original_io
+
+      def initialize(transformer, original_io)
+        @transformer = transformer
+        @original_io = original_io
+      end
+
+      def puts(*args)
+        args.flatten.each { |arg| @original_io.puts @transformer.call(arg) }
+      end
+    end
+  end
+end
+
+Parallel::UndumpableException.prepend(Switchman::Parallel::UndumpableException)

data/lib/switchman/r_spec_helper.rb

@@ -1,6 +1,6 @@
 # frozen_string_literal: true
 
-require 'switchman/test_helper'
+require "switchman/test_helper"
 
 module Switchman
   # including this module in your specs will give you several shards to
@@ -34,9 +34,9 @@ module Switchman
         groups = group.class.descendant_filtered_examples.map(&:example_group).uniq
         next unless groups.any? { |descendant_group| RSpecHelper.included_in?(descendant_group) }
 
-        puts 'Setting up sharding for all specs...'
+        puts "Setting up sharding for all specs..."
         Shard.delete_all
-        Switchman.cache.delete('default_shard')
+        Switchman.cache.delete("default_shard")
 
         @@shard1, @@shard2 = TestHelper.recreate_persistent_test_shards
         @@default_shard = Shard.default
@@ -48,7 +48,7 @@ module Switchman
             @@shard1 = @@shard1.create_new_shard
             @@shard2 = @@shard2.create_new_shard
           rescue => e
-            warn 'Sharding setup FAILED!:'
+            warn "Sharding setup FAILED!:"
             while e
               warn "\n#{e}\n"
               warn e.backtrace
@@ -66,14 +66,17 @@ module Switchman
         # we'll re-persist in the group's `before :all`; we don't want them to exist
         # in the db before then
         Shard.delete_all
-        Switchman.cache.delete('default_shard')
+        Switchman.cache.delete("default_shard")
         Shard.default(reload: true)
-        puts 'Done!'
+        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'
+          puts "Tearing down sharding for all specs"
           @@shard1.database_server.destroy unless @@shard1.database_server == Shard.default.database_server
           unless @@keep_the_shards
             @@shard1.drop_database
@@ -92,7 +95,7 @@ module Switchman
         dup = @@default_shard.dup
         dup.id = @@default_shard.id
         dup.save!
-        Switchman.cache.delete('default_shard')
+        Switchman.cache.delete("default_shard")
         Shard.default(reload: true)
         dup = @@shard1.dup
         dup.id = @@shard1.id
@@ -104,47 +107,33 @@ module Switchman
       end
 
       klass.before do
-        raise 'Sharding did not set up correctly' if @@sharding_failed
+        raise "Sharding did not set up correctly" if @@sharding_failed
 
         Shard.clear_cache
         if use_transactional_tests
           Shard.default(reload: true)
           @shard1 = Shard.find(@shard1.id)
           @shard2 = Shard.find(@shard2.id)
-          shards = [@shard2]
-          shards << @shard1 unless @shard1.database_server == Shard.default.database_server
-          shards.each do |shard|
-            shard.activate do
-              ::ActiveRecord::Base.connection.begin_transaction joinable: false
-            end
-          end
         end
       end
 
       klass.after do
         next if @@sharding_failed
 
-        if use_transactional_tests
-          shards = [@shard2]
-          shards << @shard1 unless @shard1.database_server == Shard.default.database_server
-          shards.each do |shard|
-            shard.activate do
-              ::ActiveRecord::Base.connection.rollback_transaction if ::ActiveRecord::Base.connection.transaction_open?
-            end
-          end
-        end
         # clean up after specs
-        DatabaseServer.all.each do |ds|
+        DatabaseServer.each do |ds|
           if ds.fake? && ds != @shard2.database_server
             ds.shards.delete_all unless use_transactional_tests
             ds.destroy
           end
+          ds.remove_instance_variable(:@primary_shard_id) if ds.instance_variable_defined?(:@primary_shard_id)
         end
       end
 
       klass.after(:all) do
-        Shard.connection.update("TRUNCATE #{Shard.quoted_table_name} CASCADE")
-        Switchman.cache.delete('default_shard')
+        # Don't truncate because that can create some fun cross-connection lock contention
+        Shard.delete_all
+        Switchman.cache.delete("default_shard")
         Shard.default(reload: true)
       end
     end

data/lib/switchman/rails.rb

@@ -4,10 +4,7 @@ module Switchman
   module Rails
     module ClassMethods
       def self.prepended(klass)
-        # in Rails 4+, the Rails.cache= method was used during bootstrap to set
-        # Rails.cache(_without_sharding) to the value from the config file. but now
-        # that that's done (the bootstrap happened before this module is included
-        # into Rails), we want to make sure no one tries to assign to Rails.cache,
+        # we want to make sure no one tries to assign to Rails.cache,
         # because it would be wrong w.r.t. sharding.
         klass.send(:remove_method, :cache=)
       end