switchman 1.6.1 → 2.0.9

data/app/models/switchman/shard.rb

@@ -1,11 +1,746 @@
-# ActiveRecord::Base _must_ be loaded prior to the file that defines Shard
-# Because Switchman hooks into ActiveSupport.on_load(:active_record) and
-# reference it again (technically it gets referenced from inside
-# ConnectionHandler#establish_connection, when AR's on_load(:active_record)
-# hook gets called
-ActiveRecord::Base
-# in case it was referenced _before_ the on_load hook is added, then
-# we need to make sure we define the class someone wanted (in which case
-# AR's on_load hook won't be called either yet, and establish_connection
-# will be safe)
-require 'switchman/shard_internal'
+# frozen_string_literal: true
+
+require 'switchman/database_server'
+require 'switchman/default_shard'
+require 'switchman/environment'
+require 'switchman/errors'
+
+module Switchman
+  class Shard < ::ActiveRecord::Base
+    # 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? }
+
+    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
+      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
+        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
+          # is up and running everywhere else).  This includes for looking up the
+          # default shard itself. This also needs to be a local so that this method
+          # can be re-entrant
+          default = DefaultShard.instance
+
+          # 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
+
+          # 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
+          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
+
+          # 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
+          end
+        end
+        @default
+      end
+
+      def current(category = :primary)
+        active_shards[category] || Shard.default
+      end
+
+      def activate(shards)
+        old_shards = activate!(shards)
+        yield
+      ensure
+        active_shards.merge!(old_shards) if old_shards
+      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
+        end
+        old_shards
+      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
+
+        unless cached_shards.has_key?(id)
+          cached_shards[id] = Shard.default.activate do
+            find_cached(['shard', id]) { find_by(id: id) }
+          end
+        end
+        cached_shards[id]
+      end
+
+      def clear_cache
+        cached_shards.clear
+      end
+
+      # ==== 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
+      #                error), or a proc
+      def with_each_shard(*args)
+        raise ArgumentError, "wrong number of arguments (#{args.length} for 0...3)" if args.length > 3
+
+        unless default.is_a?(Shard)
+          return Array.wrap(yield)
+        end
+
+        options = args.extract_options!
+        if args.length == 1
+          if Array === args.first && args.first.first.is_a?(Symbol)
+            categories = args.first
+          else
+            scope = args.first
+          end
+        else
+          scope, categories = args
+        end
+
+        parallel = case options[:parallel]
+                     when true
+                       1
+                     when false, nil
+                       0
+                     else
+                       options[:parallel]
+                   end
+        options.delete(:parallel)
+
+        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
+
+        if parallel > 0
+          max_procs = determine_max_procs(options.delete(: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
+            # 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]
+          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 = []
+          pids = []
+          out_fds = []
+          err_fds = []
+          pid_to_name_map = {}
+          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)
+            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, categories, options) { yield }
+          end
+
+          # 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
+              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 => 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
+              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(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
+              end
+            end
+          end
+
+          while out_fds.any? || err_fds.any?
+            wait_for_output.call(out_fds, err_fds, fd_to_name_map)
+          end
+          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|
+            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
+          end
+
+          unless errors.empty?
+            raise ParallelShardExecError.new("The following subprocesses did not exit cleanly: #{errors.sort.join(", ")}")
+          end
+          return
+        end
+
+        categories ||= []
+
+        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
+        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
+            end
+          end
+          previous_shard = shard
+        end
+        close_connections_if_needed.call(Shard.current)
+        raise exception if exception
+        result
+      end
+
+      def partition_by_shard(array, partition_proc = nil)
+        shard_arrays = {}
+        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
+              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
+          end
+          shard ||= Shard.current
+          shard_arrays[shard] ||= []
+          shard_arrays[shard] << object
+        end
+        # 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) }
+        end
+      end
+
+      # it's tedious to hold onto this same
+      # kind of sign state and transform the
+      # result in multiple places, so
+      # here we can operate on the absolute value
+      # in a provided block and trust the sign will
+      # 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
+        output = yield input_id.abs
+        output * sign
+      end
+
+      # 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"
+          any_id = any_id.value.value_before_type_cast
+        end
+
+        case any_id
+        when ::ActiveRecord::Base
+          any_id.id
+        when /^(\d+)~(-?\d+)$/
+          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
+
+      # takes an id-ish, and returns a local id and the shard it's
+      # local to. [nil, nil] if it can't be interpreted. [id, nil]
+      # if it's already a local ID. [nil, nil] if it's a well formed
+      # id, but the shard it refers to does not exist
+      NIL_NIL_ID = [nil, nil].freeze
+      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)
+            abs_id % IDS_PER_SHARD
+          else
+            return NIL_NIL_ID
+          end
+        end
+        [local_id, return_shard]
+      end
+
+      # takes an id-ish, and returns an integral id relative to
+      # target_shard. returns nil if it can't be interpreted,
+      # or the integral version of the id if it refers to a shard
+      # that does not exist
+      def relative_id_for(any_id, source_shard, target_shard)
+        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
+
+      # takes an id-ish, and returns a shortened global
+      # string id if global, and itself if local.
+      # returns any_id itself if it can't be interpreted
+      def short_id_for(any_id)
+        local_id, shard = local_id_for(any_id)
+        return any_id unless local_id
+        return local_id unless shard
+        "#{shard.id}~#{local_id}"
+      end
+
+      # takes an id-ish, and returns an integral global id.
+      # returns nil if it can't be interpreted
+      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
+          else
+            source_shard ||= Shard.current
+            source_shard.global_id_for(abs_id)
+          end
+        end
+      end
+
+      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
+          end
+        end
+
+        return max_procs
+      end
+
+      private
+      # in-process caching
+      def cached_shards
+        @cached_shards ||= {}.compare_by_identity
+      end
+
+      def add_to_cache(shard)
+        cached_shards[shard.id] = shard
+      end
+
+      def remove_from_cache(shard)
+        cached_shards.delete(shard.id)
+      end
+
+      def find_cached(key)
+        # can't simply cache the AR object since Shard has a custom serializer
+        # that calls this method
+        attributes = Switchman.cache.fetch(key) { yield&.attributes }
+        return nil unless attributes
+
+        shard = Shard.new
+        attributes.each do |attr, value|
+          shard.send(:"#{attr}=", value) if shard.respond_to?(:"#{attr}=")
+        end
+        shard.clear_changes_information
+        shard.instance_variable_set(:@new_record, false)
+        # connection info doesn't exist in database.yml;
+        # pretend the shard doesn't exist either
+        shard = nil unless shard.database_server
+        shard
+      end
+
+      def active_shards
+        Thread.current[:active_shards] ||= {}.compare_by_identity
+      end
+    end
+
+    def name
+      unless instance_variable_defined?(:@name)
+        # protect against re-entrancy
+        @name = nil
+        @name = read_attribute(:name) || default_name
+      end
+      @name
+    end
+
+    def name=(name)
+      write_attribute(:name, @name = name)
+      remove_instance_variable(:@name) if name == nil
+    end
+
+    def database_server
+      @database_server ||= DatabaseServer.find(self.database_server_id)
+    end
+
+    def database_server=(database_server)
+      self.database_server_id = database_server.id
+      @database_server = database_server
+    end
+
+    def primary?
+      self == database_server.primary_shard
+    end
+
+    def description
+      [database_server.id, name].compact.join(':')
+    end
+
+    # Shards are always on the default shard
+    def shard
+      Shard.default
+    end
+
+    def activate(*categories)
+      shards = hashify_categories(categories)
+      Shard.activate(shards) do
+        yield
+      end
+    end
+
+    # for use from console ONLY
+    def activate!(*categories)
+      shards = hashify_categories(categories)
+      Shard.activate!(shards)
+      nil
+    end
+
+    # custom serialization, since shard is self-referential
+    def _dump(depth)
+      self.id.to_s
+    end
+
+    def self._load(str)
+      lookup(str.to_i)
+    end
+
+    def drop_database
+      raise("Cannot drop the database of the default shard") if self.default?
+      return unless read_attribute(:name)
+
+      begin
+        adapter = self.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) }
+        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
+              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
+              end
+            end
+        end
+      rescue
+        logger.info "Drop failed: #{$!}"
+      end
+    end
+
+    # 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
+      end
+    end
+
+    # skip global_id.hash
+    def hash
+      id.hash
+    end
+
+    def destroy
+      raise("Cannot destroy the default shard") if self.default?
+      super
+    end
+
+    private
+
+    def clear_cache
+      Shard.default.activate do
+        Switchman.cache.delete(['shard', id].join('/'))
+        Switchman.cache.delete("default_shard") if default?
+      end
+      self.class.clear_cache
+    end
+
+    def default_name
+      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 }
+      else
+        categories.inject({}) { |h, category| h[category] = self; h }
+      end
+    end
+
+  end
+end