ar-octopus 0.8.1 → 0.8.2

data/lib/octopus/collection_association.rb

@@ -0,0 +1,9 @@
+module Octopus::CollectionAssociation
+  def self.included(base)
+    base.sharded_methods :reader, :writer, :ids_reader, :ids_writer, :create, :create!,
+                         :build, :any?, :count, :empty?, :first, :include?, :last, :length,
+                         :load_target, :many?, :reload, :size, :select, :uniq
+  end
+end
+
+ActiveRecord::Associations::CollectionAssociation.send(:include, Octopus::CollectionAssociation)

data/lib/octopus/collection_proxy.rb

@@ -0,0 +1,14 @@
+module Octopus::CollectionProxy
+  def self.included(base)
+    base.send(:include, Octopus::ShardTracking::Dynamic)
+    base.sharded_methods :any?, :build, :count, :create, :create!, :concat, :delete, :delete_all,
+                         :destroy, :destroy_all, :empty?, :find, :first, :include?, :last, :length,
+                         :many?, :pluck, :replace, :select, :size, :sum, :to_a, :uniq
+  end
+
+  def current_shard
+    @association.owner.current_shard
+  end
+end
+
+ActiveRecord::Associations::CollectionProxy.send(:include, Octopus::CollectionProxy)

data/lib/octopus/has_and_belongs_to_many_association.rb

@@ -1,17 +1,7 @@
 module Octopus::HasAndBelongsToManyAssociation
   def self.included(base)
-    base.instance_eval do
-      alias_method_chain :insert_record, :octopus
-    end
-  end
-
-  def insert_record_with_octopus(record, force = true, validate = true)
-    if should_wrap_the_connection?
-      Octopus.using(@owner.current_shard) { insert_record_without_octopus(record, force, validate) }
-    else
-      insert_record_without_octopus(record, force, validate)
-    end
+    base.sharded_methods :insert_record
   end
 end
 
-ActiveRecord::Associations::HasAndBelongsToManyAssociation.send(:include, Octopus::HasAndBelongsToManyAssociation)
+ActiveRecord::Associations::HasAndBelongsToManyAssociation.send(:include, Octopus::HasAndBelongsToManyAssociation)

data/lib/octopus/load_balancing.rb

@@ -0,0 +1,3 @@
+module Octopus::LoadBalancing
+
+end

data/lib/octopus/load_balancing/round_robin.rb

@@ -0,0 +1,15 @@
+require 'octopus/load_balancing'
+
+# The round-robin load balancing of slaves belonging to the same shard.
+# It is a pool that contains slaves which queries are distributed to.
+class Octopus::LoadBalancing::RoundRobin
+  def initialize(slaves_list)
+    @slaves_list = slaves_list
+    @slave_index = 0
+  end
+
+  # Returns the next available slave in the pool
+  def next
+    @slaves_list[@slave_index = (@slave_index + 1) % @slaves_list.length]
+  end
+end

data/lib/octopus/model.rb

@@ -2,10 +2,9 @@ require 'active_support/deprecation'
 
 module Octopus::Model
   def self.extended(base)
+    base.send(:include, Octopus::ShardTracking::Attribute)
     base.send(:include, InstanceMethods)
     base.extend(ClassMethods)
-    base.hijack_connection()
-    base.hijack_initializer()
   end
 
   module SharedMethods
@@ -28,82 +27,6 @@ module Octopus::Model
         self
       end
     end
-
-    def hijack_initializer()
-      attr_accessor :current_shard
-      around_save :run_on_shard
-
-      def set_current_shard
-        return unless Octopus.enabled?
-
-        if new_record? || self.class.connection_proxy.block
-          self.current_shard = self.class.connection_proxy.current_shard
-        else
-          self.current_shard = self.class.connection_proxy.last_current_shard || self.class.connection_proxy.current_shard
-        end
-      end
-
-      after_initialize :set_current_shard
-    end
-
-    def hijack_connection()
-      def self.should_use_normal_connection?
-        !Octopus.enabled? || self.custom_octopus_connection
-      end
-
-      def self.connection_proxy
-        @@connection_proxy ||= Octopus::Proxy.new
-      end
-
-      def self.connection_with_octopus
-        if should_use_normal_connection?
-          connection_without_octopus
-        else
-          self.connection_proxy.current_model = self
-          self.connection_proxy
-        end
-      end
-
-      def self.connection_pool_with_octopus
-        if should_use_normal_connection?
-          connection_pool_without_octopus
-        else
-          connection_proxy.connection_pool
-        end
-      end
-
-      def self.clear_active_connections_with_octopus!
-        if should_use_normal_connection?
-          clear_active_connections_without_octopus!
-        else
-          connection_proxy.clear_active_connections!
-        end
-      end
-
-      def self.clear_all_connections_with_octopus!
-        if should_use_normal_connection?
-          clear_all_connections_without_octopus!
-        else
-          connection_proxy.clear_all_connections!
-        end
-      end
-
-      def self.connected_with_octopus?
-        if should_use_normal_connection?
-          connected_without_octopus?
-        else
-          connection_proxy.connected?
-        end
-      end
-
-      class << self
-        alias_method_chain :connection, :octopus
-        alias_method_chain :connection_pool, :octopus
-        alias_method_chain :clear_all_connections!, :octopus
-        alias_method_chain :clear_active_connections!, :octopus
-        alias_method_chain :connected?, :octopus
-      end
-    end
   end
 
   module InstanceMethods
@@ -116,18 +39,20 @@ module Octopus::Model
       base.send(:alias_method_chain, :perform_validations, :octopus)
     end
 
-    def should_set_current_shard?
-      self.respond_to?(:current_shard) && !self.current_shard.nil?
-    end
+    def set_current_shard
+      return unless Octopus.enabled?
 
-    def run_on_shard(&block)
-      if self.current_shard
-        self.class.connection_proxy.run_queries_on_shard(self.current_shard, &block)
+      if new_record? || self.class.connection_proxy.block
+        self.current_shard = self.class.connection_proxy.current_shard
       else
-        yield
+        self.current_shard = self.class.connection_proxy.last_current_shard || self.class.connection_proxy.current_shard
       end
     end
 
+    def should_set_current_shard?
+      self.respond_to?(:current_shard) && !self.current_shard.nil?
+    end
+
     def equality_with_octopus(comparison_object)
       equality_without_octopus(comparison_object) && comparison_object.current_shard == current_shard
     end
@@ -161,10 +86,19 @@ module Octopus::Model
     end
 
     def hijack_methods
+      around_save :run_on_shard
+      after_initialize :set_current_shard
+
       class << self
         attr_accessor :custom_octopus_connection
         attr_accessor :custom_octopus_table_name
 
+        alias_method_chain :connection, :octopus
+        alias_method_chain :connection_pool, :octopus
+        alias_method_chain :clear_all_connections!, :octopus
+        alias_method_chain :clear_active_connections!, :octopus
+        alias_method_chain :connected?, :octopus
+
         if Octopus.rails3?
           alias_method_chain(:set_table_name, :octopus)
         end
@@ -176,6 +110,61 @@ module Octopus::Model
       end
     end
 
+    def connection_proxy
+      cached = ActiveRecord::Base.class_variable_get :@@connection_proxy rescue nil
+      cached ||
+        begin
+          p = Octopus::Proxy.new
+          ActiveRecord::Base.class_variable_set :@@connection_proxy, p
+          p
+        end
+    end
+
+    def should_use_normal_connection?
+      !Octopus.enabled? || custom_octopus_connection
+    end
+
+    def connection_with_octopus
+      if should_use_normal_connection?
+        connection_without_octopus
+      else
+        connection_proxy.current_model = self
+        connection_proxy
+      end
+    end
+
+    def connection_pool_with_octopus
+      if should_use_normal_connection?
+        connection_pool_without_octopus
+      else
+        connection_proxy.connection_pool
+      end
+    end
+
+    def clear_active_connections_with_octopus!
+      if should_use_normal_connection?
+        clear_active_connections_without_octopus!
+      else
+        connection_proxy.clear_active_connections!
+      end
+    end
+
+    def clear_all_connections_with_octopus!
+      if should_use_normal_connection?
+        clear_all_connections_without_octopus!
+      else
+        connection_proxy.clear_all_connections!
+      end
+    end
+
+    def connected_with_octopus?
+      if should_use_normal_connection?
+        connected_without_octopus?
+      else
+        connection_proxy.connected?
+      end
+    end
+
     def set_table_name_with_octopus(value = nil, &block)
       self.custom_octopus_table_name = true
       set_table_name_without_octopus(value, &block)

data/lib/octopus/proxy.rb

@@ -1,7 +1,9 @@
 require "set"
+require 'octopus/slave_group'
+require 'octopus/load_balancing/round_robin'
 
 class Octopus::Proxy
-  attr_accessor :config
+  attr_accessor :config, :sharded
 
   def initialize(config = Octopus.config)
     initialize_shards(config)
@@ -10,18 +12,20 @@ class Octopus::Proxy
 
   def initialize_shards(config)
     @shards = HashWithIndifferentAccess.new
+    @shards_slave_groups = HashWithIndifferentAccess.new
+    @slave_groups = HashWithIndifferentAccess.new
     @groups = {}
     @adapters = Set.new
     @config = ActiveRecord::Base.connection_pool_without_octopus.connection.instance_variable_get(:@config)
 
     if !config.nil?
       @entire_sharded = config['entire_sharded']
-      shards_config = config[Octopus.rails_env()]
+      @shards_config = config[Octopus.rails_env()]
     end
 
-    shards_config ||= []
+    @shards_config ||= []
 
-    shards_config.each do |key, value|
+    @shards_config.each do |key, value|
       if value.is_a?(String)
         value = resolve_string_connection(value).merge(:octopus_shard => key)
         initialize_adapter(value['adapter'])
@@ -30,6 +34,24 @@ class Octopus::Proxy
         value.merge!(:octopus_shard => key)
         initialize_adapter(value['adapter'])
         @shards[key.to_sym] = connection_pool_for(value, "#{value['adapter']}_connection")
+
+        slave_group_configs = value.select do |k,v|
+          structurally_slave_group? v
+        end
+
+        if slave_group_configs.present?
+          slave_groups = HashWithIndifferentAccess.new
+          slave_group_configs.each do |slave_group_name, slave_configs|
+            slaves = HashWithIndifferentAccess.new
+            slave_configs.each do |slave_name, slave_config|
+              @shards[slave_name.to_sym] = connection_pool_for(slave_config, "#{value['adapter']}_connection")
+              slaves[slave_name.to_sym] = slave_name.to_sym
+            end
+            slave_groups[slave_group_name.to_sym] = Octopus::SlaveGroup.new(slaves)
+          end
+          @shards_slave_groups[key.to_sym] = slave_groups
+          @sharded = true
+        end
       elsif value.is_a?(Hash)
         @groups[key.to_s] = []
 
@@ -42,6 +64,11 @@ class Octopus::Proxy
           @shards[k.to_sym] = connection_pool_for(config_with_octopus_shard, "#{v['adapter']}_connection")
           @groups[key.to_s] << k.to_sym
         end
+
+        if structurally_slave_group? value
+          slaves = Hash[@groups[key.to_s].map { |v| [v, v ] }]
+          @slave_groups[key.to_sym] = Octopus::SlaveGroup.new(slaves)
+        end
       end
     end
 
@@ -58,7 +85,7 @@ class Octopus::Proxy
 
     @slaves_list = @shards.keys.map {|sym| sym.to_s}.sort
     @slaves_list.delete('master')
-    @slave_index = 0
+    @slaves_load_balancer = Octopus::LoadBalancing::RoundRobin.new(@slaves_list)
   end
 
   def current_model
@@ -74,8 +101,29 @@ class Octopus::Proxy
   end
 
   def current_shard=(shard_symbol)
+    self.current_slave_group = nil
     if shard_symbol.is_a?(Array)
       shard_symbol.each {|symbol| raise "Nonexistent Shard Name: #{symbol}" if @shards[symbol].nil? }
+    elsif shard_symbol.is_a?(Hash)
+      hash = shard_symbol
+      shard_symbol = hash[:shard]
+      slave_group_symbol = hash[:slave_group]
+
+      if shard_symbol.nil? && slave_group_symbol.nil?
+        raise "Neither shard or slave group must be specified"
+      end
+
+      if shard_symbol.present?
+        raise "Nonexistent Shard Name: #{shard_symbol}" if @shards[shard_symbol].nil?
+      end
+
+      if slave_group_symbol.present?
+        if (@shards_slave_groups.try(:[], shard_symbol).present? && @shards_slave_groups[shard_symbol][slave_group_symbol].nil?) ||
+            (@shards_slave_groups.try(:[], shard_symbol).nil? && @slave_groups[slave_group_symbol].nil?)
+          raise "Nonexistent Slave Group Name: #{slave_group_symbol} in shards config: #{@shards_config.inspect}"
+        end
+        self.current_slave_group = slave_group_symbol
+      end
     else
       raise "Nonexistent Shard Name: #{shard_symbol}" if @shards[shard_symbol].nil?
     end
@@ -96,6 +144,14 @@ class Octopus::Proxy
     Thread.current["octopus.current_group"] = group_symbol
   end
 
+  def current_slave_group
+    Thread.current["octopus.current_slave_group"]
+  end
+
+  def current_slave_group=(slave_group_symbol)
+    Thread.current["octopus.current_slave_group"] = slave_group_symbol
+  end
+
   def block
     Thread.current["octopus.block"]
   end
@@ -112,6 +168,10 @@ class Octopus::Proxy
     Thread.current["octopus.last_current_shard"] = last_current_shard
   end
 
+  def fully_replicated?
+    @fully_replicated || Thread.current["octopus.fully_replicated"]
+  end
+
   # Public: Whether or not a group exists with the given name converted to a
   # string.
   #
@@ -157,16 +217,10 @@ class Octopus::Proxy
   end
 
   def run_queries_on_shard(shard, &block)
-    older_shard = self.current_shard
-    last_block = self.block
-
-    begin
-      self.block = true
-      self.current_shard = shard
-      yield
-    ensure
-      self.block = last_block || false
-      self.current_shard = older_shard
+    keeping_connection_proxy do
+      using_shard(shard) do
+        yield
+      end
     end
   end
 
@@ -176,7 +230,7 @@ class Octopus::Proxy
     end
   end
 
-  def clean_proxy()
+  def clean_connection_proxy()
     self.current_shard = :master
     self.current_group = nil
     self.block = false
@@ -189,7 +243,8 @@ class Octopus::Proxy
   end
 
   def transaction(options = {}, &block)
-    if @replicated && (current_model.replicated || @fully_replicated)
+    replicated = @replicated && (current_model.replicated || fully_replicated?)
+    if !sharded && replicated
       self.run_queries_on_shard(:master) do
         select_connection.transaction(options, &block)
       end
@@ -199,11 +254,15 @@ class Octopus::Proxy
   end
 
   def method_missing(method, *args, &block)
-    if should_clean_connection?(method)
+    if should_clean_connection_proxy?(method)
       conn = select_connection()
       self.last_current_shard = self.current_shard
-      clean_proxy()
+      clean_connection_proxy()
       conn.send(method, *args, &block)
+    elsif should_send_queries_to_shard_slave_group?(method)
+      send_queries_to_shard_slave_group(method, *args, &block)
+    elsif should_send_queries_to_slave_group?(method)
+      send_queries_to_slave_group(method, *args, &block)
     elsif should_send_queries_to_replicated_databases?(method)
       send_queries_to_selected_slave(method, *args, &block)
     else
@@ -244,6 +303,22 @@ class Octopus::Proxy
     @shards.any? { |k, v| v.connected? }
   end
 
+  def should_send_queries_to_shard_slave_group?(method)
+    should_use_slaves_for_method?(method) && @shards_slave_groups.try(:[], current_shard).try(:[], current_slave_group).present?
+  end
+
+  def send_queries_to_shard_slave_group(method, *args, &block)
+    send_queries_to_balancer(@shards_slave_groups[current_shard][current_slave_group], method, *args, &block)
+  end
+
+  def should_send_queries_to_slave_group?(method)
+    should_use_slaves_for_method?(method) && @slave_groups.try(:[], current_slave_group).present?
+  end
+
+  def send_queries_to_slave_group(method, *args, &block)
+    send_queries_to_balancer(@slave_groups[current_slave_group], method, *args, &block)
+  end
+
   protected
 
   def connection_pool_for(adapter, config)
@@ -274,27 +349,89 @@ class Octopus::Proxy
     resolver.spec.config.stringify_keys
   end
 
-  def should_clean_connection?(method)
+  def should_clean_connection_proxy?(method)
     method.to_s =~ /insert|select|execute/ && !@replicated && !self.block
   end
 
+  # Try to use slaves if and only if `replicated: true` is specified in `shards.yml` and no slaves groups are defined
   def should_send_queries_to_replicated_databases?(method)
-    @replicated && method.to_s =~ /select/ && !self.block
+    @replicated && method.to_s =~ /select/ && !self.block && !slaves_grouped?
   end
 
   def send_queries_to_selected_slave(method, *args, &block)
-    old_shard = self.current_shard
+    if current_model.replicated || fully_replicated?
+      selected_slave = @slaves_load_balancer.next
+    else
+      selected_slave = :master
+    end
 
-    begin
-      if current_model.replicated || @fully_replicated
-        self.current_shard = @slaves_list[@slave_index = (@slave_index + 1) % @slaves_list.length]
-      else
-        self.current_shard = :master
-      end
+    send_queries_to_slave(selected_slave, method, *args, &block)
+  end
+
+  # We should use slaves if and only if its safe to do so.
+  #
+  # We can safely use slaves when:
+  # (1) `replicated: true` is specified in `shards.yml`
+  # (2) The current model is `replicated()`, or `fully_replicated: true` is specified in `shards.yml` which means that
+  #     all the model is `replicated()`
+  # (3) It's a SELECT query
+  # while ensuring that we revert `current_shard` from the selected slave to the (shard's) master
+  # not to make queries other than SELECT leak to the slave.
+  def should_use_slaves_for_method?(method)
+    @replicated && (current_model.replicated || fully_replicated?) && method.to_s =~ /select/
+  end
+
+  def slaves_grouped?
+    @slave_groups.present?
+  end
 
+  # Temporarily switch `current_shard` to the next slave in a slave group and send queries to it
+  # while preserving `current_shard`
+  def send_queries_to_balancer(balancer, method, *args, &block)
+    send_queries_to_slave(balancer.next, method, *args, &block)
+  end
+
+  # Temporarily switch `current_shard` to the specified slave and send queries to it
+  # while preserving `current_shard`
+  def send_queries_to_slave(slave, method, *args, &block)
+    using_shard(slave) do
       select_connection.send(method, *args, &block)
+    end
+  end
+
+  # Temporarily block cleaning connection proxy and run the block
+  #
+  # @see Octopus::Proxy#should_clean_connection?
+  # @see Octopus::Proxy#clean_connection_proxy
+  def keeping_connection_proxy(&block)
+    last_block = self.block
+
+    begin
+      self.block = true
+      yield
+    ensure
+      self.block = last_block || false
+    end
+  end
+
+  # Temporarily switch `current_shard` and run the block
+  def using_shard(shard, &block)
+    older_shard = self.current_shard
+
+    begin
+      self.current_shard = shard
+      yield
     ensure
-      self.current_shard = old_shard
+      self.current_shard = older_shard
     end
   end
+
+  def structurally_slave?(config)
+    config.is_a?(Hash) && config.key?("adapter")
+  end
+
+  def structurally_slave_group?(config)
+    config.is_a?(Hash) && config.values.any? {|v| structurally_slave? v }
+  end
+
 end