activerecord-turntable 3.1.0 → 4.0.0

data/lib/active_record/turntable/base.rb

@@ -5,20 +5,33 @@ module ActiveRecord::Turntable
     extend ActiveSupport::Concern
 
     included do
-      class_attribute :turntable_connections, :turntable_clusters,
-                      :turntable_enabled, :turntable_sequencer_enabled
+      class_attribute :turntable_connections, :turntable_clusters, :turntable_sequencers,
+                      :turntable_enabled, :turntable_sequencer_enabled, :turntable_configuration
 
       self.turntable_connections = {}
       self.turntable_clusters = {}.with_indifferent_access
+      self.turntable_sequencers = {}.with_indifferent_access
       self.turntable_enabled = false
       self.turntable_sequencer_enabled = false
+
       class << self
         delegate :shards_transaction, :with_all, to: :connection
-      end
 
-      ActiveSupport.on_load(:turntable_config_loaded) do
-        self.initialize_clusters!
+        def reset_turntable_configuration(configuration, reset = true)
+          old = self.turntable_configuration
+          self.turntable_configuration = configuration
+
+          old.release! if old
+
+          if reset
+            # TODO: replace exitsting connection_pool when configurations reloaded
+            self.turntable_clusters = turntable_configuration.clusters
+            self.turntable_sequencers = turntable_configuration.sequencers
+            ActiveSupport.run_load_hooks(:turntable_configuration_loaded, ActiveRecord::Base)
+          end
+        end
       end
+
       include ClusterHelperMethods
     end
 
@@ -27,20 +40,18 @@ module ActiveRecord::Turntable
       # @param [Symbol] shard_key_name shard key attribute name
       # @param [Hash] options
       def turntable(cluster_name, shard_key_name, options = {})
-        class_attribute :turntable_shard_key,
-                        :turntable_cluster, :turntable_cluster_name
+        class_attribute :turntable_shard_key, :turntable_cluster_name
 
         self.turntable_enabled = true
         self.turntable_cluster_name = cluster_name
         self.turntable_shard_key = shard_key_name
-        self.turntable_cluster =
-          self.turntable_clusters[cluster_name] ||= Cluster.new(
-            turntable_config[:clusters][cluster_name],
-            options
-          )
         turntable_replace_connection_pool
       end
 
+      def turntable_cluster
+        turntable_clusters[turntable_cluster_name]
+      end
+
       def turntable_replace_connection_pool
         ch = connection_handler
         cp = ConnectionProxy.new(self, turntable_cluster)
@@ -50,31 +61,22 @@ module ActiveRecord::Turntable
         ch.send(:owner_to_pool)[connection_specification_name] = pp
       end
 
-      def initialize_clusters!
-        turntable_config[:clusters].each do |name, spec|
-          self.turntable_clusters[name] ||= Cluster.new(spec)
-        end
-      end
-
-      def spec_for(config)
-        begin
-          require "active_record/connection_adapters/#{config['adapter']}_adapter"
-        rescue LoadError => e
-          raise "Please install the #{config['adapter']} adapter: `gem install activerecord-#{config['adapter']}-adapter` (#{e})"
-        end
-        adapter_method = "#{config['adapter']}_connection"
-        ActiveRecord::ConnectionAdapters::ConnectionSpecification.new(config, adapter_method)
-      end
-
       def clear_all_connections!
         turntable_connections.values.each(&:disconnect!)
       end
 
-      def sequencer(sequence_name, *args)
-        class_attribute :turntable_sequencer
+      def sequencer(sequencer_name, *args)
+        class_attribute :turntable_sequencer_name
+        class << self
+          prepend ActiveRecordExt::Sequencer
+        end
 
         self.turntable_sequencer_enabled = true
-        self.turntable_sequencer = ActiveRecord::Turntable::Sequencer.build(self, sequence_name, *args)
+        self.turntable_sequencer_name = sequencer_name
+      end
+
+      def turntable_sequencer
+        turntable_sequencers[turntable_sequencer_name]
       end
 
       def turntable_enabled?
@@ -85,12 +87,8 @@ module ActiveRecord::Turntable
         turntable_sequencer_enabled
       end
 
-      def current_sequence
-        connection.current_sequence_value(self.sequence_name) if sequencer_enabled?
-      end
-
       def current_last_shard
-        turntable_cluster.select_shard(current_sequence) if sequencer_enabled?
+        turntable_cluster.select_shard(current_sequence_value) if sequencer_enabled?
       end
 
       def with_shard(any_shard)
@@ -104,12 +102,18 @@ module ActiveRecord::Turntable
                 end
         connection.with_shard(shard) { yield }
       end
-    end
 
-    def shards_transaction(options = {}, &block)
-      self.class.shards_transaction(options, &block)
+      def with_slave
+        connection.with_slave { yield }
+      end
+
+      def with_master
+        connection.with_master { yield }
+      end
     end
 
+    delegate :shards_transaction, :turntable_cluster, to: :class
+
     # @return [ActiveRecord::Turntable::Shard] current shard for self
     def turntable_shard
       turntable_cluster.shard_for(self.send(turntable_shard_key))

data/lib/active_record/turntable/cluster.rb

@@ -1,4 +1,5 @@
 require "active_support/core_ext/hash/indifferent_access"
+require "concurrent/atomic/thread_local_var"
 
 module ActiveRecord::Turntable
   class Cluster
@@ -7,40 +8,27 @@ module ActiveRecord::Turntable
       "algorithm" => "range",
     }.with_indifferent_access
 
-    def initialize(cluster_spec, options = {})
-      @config = DEFAULT_CONFIG.merge(cluster_spec)
-      @options = options.with_indifferent_access
-      @shards = {}.with_indifferent_access
+    attr_accessor :algorithm, :shard_registry, :sequencer_registry
 
-      # setup sequencer
-      seq = (@options[:seq] || @config[:seq])
-      if seq
-        if seq.values.size > 0 && seq.values.first[:seq_type] == "mysql"
-          @seq_shard = SeqShard.new(seq.values.first)
-        end
-      end
-
-      # setup shards
-      @config[:shards].each do |spec|
-        @shards[spec[:connection]] ||= Shard.new(spec)
-      end
-
-      # setup algorithm
-      alg_name = "ActiveRecord::Turntable::Algorithm::#{@config[:algorithm].camelize}Algorithm"
-      @algorithm = alg_name.constantize.new(@config)
+    def initialize
+      @slave_enabled = Concurrent::ThreadLocalVar.new(false)
     end
 
-    def seq
-      @seq_shard
+    def self.build(sequencer_registry)
+      self.new.tap do |instance|
+        instance.shard_registry = ShardRegistry.new(instance)
+        instance.sequencer_registry = sequencer_registry
+        yield instance
+      end
     end
 
-    attr_reader :shards
+    delegate :shards, :shard_maps, :release!, to: :shard_registry
 
     def shard_for(key)
-      @shards[@algorithm.calculate(key)]
+      algorithm.choose(shard_maps, key)
     rescue
       raise ActiveRecord::Turntable::CannotSpecifyShardError,
-            "cannot select_shard for key:#{key}"
+            "cannot select shard for key:#{key.inspect}"
     end
 
     def select_shard(key)
@@ -52,7 +40,7 @@ module ActiveRecord::Turntable
       unless in_recursion
         shards = Array.wrap(shards).dup
         if shards.blank?
-          shards = @shards.values.dup
+          shards = self.shards.dup
         end
       end
       shard = to_shard(shards.shift)
@@ -83,10 +71,24 @@ module ActiveRecord::Turntable
       end
     end
 
+    def slave_enabled?
+      @slave_enabled.value
+    end
+
+    def set_slave_enabled(enabled)
+      @slave_enabled.value = enabled
+    end
+
+    def sequencers
+      sequencer_registry.all
+    end
+
+    def sequencer(name)
+      sequencers[name]
+    end
+
     def weighted_shards(key = nil)
-      Hash[@algorithm.calculate_used_shards_with_weight(key).map do |k, v|
-        [@shards[k], v]
-      end]
+      @algorithm.shard_weights(shard_maps, key)
     end
   end
 end

data/lib/active_record/turntable/cluster_helper_methods.rb

@@ -3,7 +3,7 @@ module ActiveRecord::Turntable
     extend ActiveSupport::Concern
 
     included do
-      ActiveSupport.on_load(:turntable_config_loaded) do
+      ActiveSupport.on_load(:turntable_configuration_loaded) do
         turntable_clusters.each do |name, _cluster|
           turntable_define_cluster_methods(name)
         end
@@ -42,8 +42,18 @@ module ActiveRecord::Turntable
         end
       end
 
+      def spec_for(config)
+        begin
+          require "active_record/connection_adapters/#{config["adapter"]}_adapter"
+        rescue LoadError => e
+          raise "Please install the #{config["adapter"]} adapter: `gem install activerecord-#{config["adapter"]}-adapter` (#{e})"
+        end
+        adapter_method = "#{config["adapter"]}_connection"
+        ActiveRecord::ConnectionAdapters::ConnectionSpecification.new(config, adapter_method)
+      end
+
       def weighted_random_shard_with(*klasses, &block)
-        shards_weight = self.turntable_cluster.weighted_shards(self.current_sequence)
+        shards_weight = self.turntable_cluster.weighted_shards(self.current_sequence_value(sequence_name))
         sum = shards_weight.values.inject(&:+)
         idx = rand(sum)
         shard, _weight = shards_weight.find { |_k, v|

data/lib/active_record/turntable/cluster_registry.rb

@@ -0,0 +1,7 @@
+module ActiveRecord::Turntable
+  class ClusterRegistry < HashWithIndifferentAccess
+    def release!
+      values.each(&:release!)
+    end
+  end
+end

data/lib/active_record/turntable/configuration.rb

@@ -0,0 +1,50 @@
+module ActiveRecord::Turntable
+  class Configuration
+    extend ActiveSupport::Autoload
+    autoload :DSL
+    autoload :Loader
+
+    attr_reader :cluster_registry, :sequencer_registry
+    attr_accessor :raise_on_not_specified_shard_query,
+                  :raise_on_not_specified_shard_update
+    alias_method :configure, :instance_exec
+    alias_method :clusters, :cluster_registry
+
+    def initialize
+      @cluster_registry = ClusterRegistry.new
+      @sequencer_registry = SequencerRegistry.new
+    end
+
+    def cluster(name)
+      cluster_registry[name]
+    end
+
+    def sequencers
+      sequencer_registry.sequencers
+    end
+
+    def sequencer(name)
+      sequencer_registry[name]
+    end
+
+    def release!
+      cluster_registry.release!
+      sequencer_registry.release!
+    end
+
+    def self.configure(&block)
+      new.tap { |c| c.configure(&block) }
+    end
+
+    def self.load(path, env)
+      case File.extname(path)
+      when ".yml"
+        Loader::YAML.load(path, env)
+      when ".rb"
+        Loader::DSL.load(path)
+      else
+        raise InvalidConfigurationError, "Invalid configuration file path: #{path}"
+      end
+    end
+  end
+end

data/lib/active_record/turntable/configuration/dsl.rb

@@ -0,0 +1,79 @@
+require "active_record/turntable/configuration"
+
+module ActiveRecord::Turntable
+  class Configuration
+    class DSL
+      attr_reader :configuration
+
+      def initialize(configuration = Configuration.new)
+        @configuration = configuration
+      end
+
+      def cluster(name, &block)
+        cluster_dsl = ClusterDSL.new(configuration).tap { |dsl| dsl.instance_exec(&block) }
+        configuration.clusters[name] = cluster_dsl.cluster
+      end
+
+      GLOBAL_SETTINGS_KEYS = %w(
+        raise_on_not_specified_shard_query
+        raise_on_not_specified_shard_update
+      ).freeze
+
+      GLOBAL_SETTINGS_KEYS.each do |k|
+        define_method(k) do |value|
+          configuration.send("#{k}=", value)
+        end
+      end
+
+      def global_settings_keys
+        GLOBAL_SETTINGS_KEYS
+      end
+
+      class ClusterDSL < DSL
+        attr_reader :sequencers
+
+        def initialize(configuration)
+          super
+          @algorithm = Algorithm.class_for("range").new(nil)
+          @shard_settings = []
+          @sequencer_settings = []
+        end
+
+        def cluster
+          Cluster.build(configuration.sequencer_registry) do |c|
+            c.algorithm = @algorithm
+            @shard_settings.each do |setting|
+              c.shard_registry.add(setting)
+            end
+            @sequencer_settings.each do |name, type, options|
+              c.sequencer_registry.add(name, type, options, c)
+            end
+          end
+        end
+
+        def algorithm(type, options = {})
+          @algorithm = Algorithm.class_for(type.to_s).new(options)
+        end
+
+        def sequencer(sequencer_name, type, options = {})
+          @sequencer_settings << [sequencer_name.to_s, type.to_s, options]
+        end
+
+        ShardSetting = Struct.new(:name, :range, :slaves) do
+          def range
+            case self[:range]
+            when Integer
+              self[:range]..self[:range]
+            else
+              self[:range]
+            end
+          end
+        end
+
+        def shard(range,  slaves: [], to:)
+          @shard_settings << ShardSetting.new(to.to_s, range, slaves.map(&:to_s))
+        end
+      end
+    end
+  end
+end

data/lib/active_record/turntable/configuration/loader.rb

@@ -0,0 +1,10 @@
+module ActiveRecord::Turntable
+  class Configuration
+    module Loader
+      extend ActiveSupport::Autoload
+
+      autoload :YAML
+      autoload :DSL
+    end
+  end
+end

data/lib/active_record/turntable/configuration/loader/dsl.rb

@@ -0,0 +1,23 @@
+module ActiveRecord::Turntable
+  class Configuration
+    class Loader::DSL
+      attr_reader :path, :configuration, :dsl
+
+      def initialize(path, configuration = Configuration.new)
+        @path = path
+        @configuration = configuration
+        @dsl = DSL.new(@configuration)
+      end
+
+      def self.load(path, configuration = Configuration.new)
+        new(path, configuration).load
+      end
+
+      def load
+        @dsl.instance_eval(File.read(@path), @path, 1)
+
+        configuration
+      end
+    end
+  end
+end

data/lib/active_record/turntable/configuration/loader/yaml.rb

@@ -0,0 +1,62 @@
+module ActiveRecord::Turntable
+  class Configuration
+    class Loader::YAML
+      attr_reader :path, :configuration, :dsl
+
+      def initialize(path, configuration = Configuration.new)
+        @path = path
+        @configuration = configuration
+        @dsl = DSL.new(@configuration)
+      end
+
+      def self.load(path, env, configuration = Configuration.new)
+        new(path, configuration).load(env)
+      end
+
+      def load(env)
+        yaml = YAML.load(ERB.new(IO.read(path)).result).with_indifferent_access[env]
+        load_clusters(yaml[:clusters])
+        load_global_settings(yaml)
+
+        configuration
+      end
+
+      private
+
+        def load_clusters(clusters_config)
+          clusters_config.each do |cluster_name, conf|
+            @dsl.cluster(cluster_name) do
+              algorithm conf[:algorithm] if conf[:algorithm]
+
+              if conf[:seq]
+                conf[:seq].each do |sequence_name, sequence_conf|
+                  sequencer(sequence_name, (sequence_conf[:seq_type] || :mysql), sequence_conf)
+                end
+              end
+
+              if conf[:shards]
+                current_lower_limit = 1
+                conf[:shards].each do |shard_conf|
+                  upper_limit = if shard_conf.has_key?(:less_than)
+                                  shard_conf[:less_than] - 1
+                                else
+                                  current_lower_limit
+                                end
+                  shard current_lower_limit..upper_limit, to: shard_conf[:connection], slaves: Array.wrap(shard_conf[:slaves])
+                  current_lower_limit = upper_limit + 1
+                end
+              end
+            end
+          end
+        end
+
+        def load_global_settings(yaml)
+          yaml.each do |k, v|
+            if @dsl.global_settings_keys.include?(k)
+              @dsl.send(k, v)
+            end
+          end
+        end
+    end
+  end
+end