makara 0.3.8 → 0.5.0

data/lib/makara/pool.rb

@@ -1,7 +1,7 @@
 require 'active_support/core_ext/hash/keys'
+require 'makara/strategies/shard_aware'
 
 # Wraps a collection of similar connections and chooses which one to use
-# Uses the Makara::Context to determine if the connection needs rotation.
 # Provides convenience methods for accessing underlying connections
 
 module Makara
@@ -14,15 +14,22 @@ module Makara
     attr_reader :role
     attr_reader :connections
     attr_reader :strategy
+    attr_reader :shard_strategy_class
+    attr_reader :default_shard
 
     def initialize(role, proxy)
       @role             = role
       @proxy            = proxy
-      @context          = Makara::Context.get_current
       @connections      = []
       @blacklist_errors = []
       @disabled         = false
-      @strategy         = proxy.strategy_for(role)
+      if proxy.shard_aware_for(role)
+        @strategy = Makara::Strategies::ShardAware.new(self)
+        @shard_strategy_class = proxy.strategy_class_for(proxy.strategy_name_for(role))
+        @default_shard = proxy.default_shard_for(role)
+      else
+        @strategy = proxy.strategy_for(role)
+      end
     end
 
 
@@ -64,17 +71,14 @@ module Makara
       @connections.each do |con|
         next if con._makara_blacklisted?
         begin
-          if block
-            value = @proxy.error_handler.handle(con) do
+          ret = @proxy.error_handler.handle(con) do
+            if block
               yield con
+            else
+              con.send(method, *args)
             end
           end
 
-          if method
-            ret = con.send(method, *args)
-          else
-            ret = value
-          end
           one_worked = true
         rescue Makara::Errors::BlacklistConnection => e
           errors.insert(0, e)
@@ -96,33 +100,47 @@ module Makara
     # Provide a connection that is not blacklisted and connected. Handle any errors
     # that may occur within the block.
     def provide
-      provided_connection = self.next
+      attempt = 0
+      begin
+        provided_connection = self.next
 
-      # nil implies that it's blacklisted
-      if provided_connection
+        # nil implies that it's blacklisted
+        if provided_connection
 
-        value = @proxy.error_handler.handle(provided_connection) do
-          yield provided_connection
-        end
+          value = @proxy.error_handler.handle(provided_connection) do
+            yield provided_connection
+          end
 
-        @blacklist_errors = []
+          @blacklist_errors = []
 
-        value
+          value
 
-      # if we've made any connections within this pool, we should report the blackout.
-      elsif connection_made?
-        err = Makara::Errors::AllConnectionsBlacklisted.new(self, @blacklist_errors)
-        @blacklist_errors = []
-        raise err
-      else
-        raise Makara::Errors::NoConnectionsAvailable.new(@role) unless @disabled
-      end
+        # if we've made any connections within this pool, we should report the blackout.
+        elsif connection_made?
+          err = Makara::Errors::AllConnectionsBlacklisted.new(self, @blacklist_errors)
+          @blacklist_errors = []
+          raise err
+        else
+          raise Makara::Errors::NoConnectionsAvailable.new(@role) unless @disabled
+        end
 
-    # when a connection causes a blacklist error within the provided block, we blacklist it then retry
-    rescue Makara::Errors::BlacklistConnection => e
-      @blacklist_errors.insert(0, e)
-      provided_connection._makara_blacklist!
-      retry
+      # when a connection causes a blacklist error within the provided block, we blacklist it then retry
+      rescue Makara::Errors::BlacklistConnection => e
+        @blacklist_errors.insert(0, e)
+        in_transaction = self.role == "master" && provided_connection._makara_in_transaction?
+        provided_connection._makara_blacklist!
+        raise Makara::Errors::BlacklistedWhileInTransaction.new(@role) if in_transaction
+        attempt += 1
+        if attempt < @connections.length
+          retry
+        elsif connection_made?
+          err = Makara::Errors::AllConnectionsBlacklisted.new(self, @blacklist_errors)
+          @blacklist_errors = []
+          raise err
+        else
+          raise Makara::Errors::NoConnectionsAvailable.new(@role) unless @disabled
+        end
+      end
     end
 
 
@@ -140,16 +158,11 @@ module Makara
     # to be sticky, provide back the current connection assuming it is
     # not blacklisted.
     def next
-      if @proxy.sticky && Makara::Context.get_current == @context
-        con = @strategy.current
-        return con if con
-      end
-
-      con = @strategy.next
-      if con
-        @context = Makara::Context.get_current
+      if @proxy.sticky && (curr = @strategy.current)
+        curr
+      else
+        @strategy.next
       end
-      con
     end
   end
 end

data/lib/makara/proxy.rb

@@ -14,8 +14,9 @@ module Makara
 
     METHOD_MISSING_SKIP = [ :byebug, :puts ]
 
-    class_attribute :hijack_methods
+    class_attribute :hijack_methods, :control_methods
     self.hijack_methods = []
+    self.control_methods = []
 
     class << self
       def hijack_method(*method_names)
@@ -38,12 +39,24 @@ module Makara
           end
         end
       end
+
+      def control_method(*method_names)
+        self.control_methods = self.control_methods || []
+        self.control_methods |= method_names
+
+        method_names.each do |method_name|
+          define_method method_name do |*args, &block|
+            control&.send(method_name, *args, &block)
+          end
+        end
+      end
     end
 
 
     attr_reader :error_handler
     attr_reader :sticky
     attr_reader :config_parser
+    attr_reader :control
 
     def initialize(config)
       @config         = config.symbolize_keys
@@ -59,22 +72,21 @@ module Makara
     end
 
     def without_sticking
-      before_context = @master_context
-      @master_context = nil
       @skip_sticking = true
       yield
     ensure
       @skip_sticking = false
-      @master_context ||= before_context
     end
 
     def hijacked?
       @hijacked
     end
 
-    def stick_to_master!(write_to_cache = true)
-      @master_context = Makara::Context.get_current
-      Makara::Context.stick(@master_context, @id, @ttl) if write_to_cache
+    # If persist is true, we stick the proxy to master for subsequent requests
+    # up to master_ttl duration. Otherwise we just stick it for the current request
+    def stick_to_master!(persist = true)
+      stickiness_duration = persist ? @ttl : 0
+      Makara::Context.stick(@id, stickiness_duration)
     end
 
     def strategy_for(role)
@@ -85,6 +97,14 @@ module Makara
       @config_parser.makara_config["#{role}_strategy".to_sym]
     end
 
+    def shard_aware_for(role)
+      @config_parser.makara_config["#{role}_shard_aware".to_sym]
+    end
+
+    def default_shard_for(role)
+      @config_parser.makara_config["#{role}_default_shard".to_sym]
+    end
+
     def strategy_class_for(strategy_name)
       case strategy_name
       when 'round_robin', 'roundrobin', nil, ''
@@ -149,8 +169,14 @@ module Makara
     end
 
     def any_connection
-      @master_pool.provide do |con|
-        yield con
+      if @master_pool.disabled
+        @slave_pool.provide do |con|
+          yield con
+        end
+      else
+        @master_pool.provide do |con|
+          yield con
+        end
       end
     rescue ::Makara::Errors::AllConnectionsBlacklisted, ::Makara::Errors::NoConnectionsAvailable
       begin
@@ -182,7 +208,6 @@ module Makara
 
       # for testing purposes
       pool = _appropriate_pool(method_name, args)
-
       yield pool
 
     rescue ::Makara::Errors::AllConnectionsBlacklisted, ::Makara::Errors::NoConnectionsAvailable => e
@@ -202,15 +227,11 @@ module Makara
         stick_to_master(method_name, args)
         @master_pool
 
-      # in this context, we've already stuck to master
-      elsif Makara::Context.get_current == @master_context
-        @master_pool
-
-      elsif previously_stuck_to_master?
+      elsif stuck_to_master?
 
-        # we're only on master because of the previous context so
-        # behave like we're sticking to master but store the current context
-        stick_to_master(method_name, args, false)
+        # we're on master because we already stuck this proxy in this
+        # request or because we got stuck in previous requests and the
+        # stickiness is still valid
         @master_pool
 
       # all slaves are down (or empty)
@@ -248,28 +269,28 @@ module Makara
     end
 
 
-    def previously_stuck_to_master?
-      @sticky && Makara::Context.previously_stuck?(@id)
+    def stuck_to_master?
+      sticky? && Makara::Context.stuck?(@id)
     end
 
-
-    def stick_to_master(method_name, args, write_to_cache = true)
-      # if we're already stuck to master, don't bother doing it again
-      return if @master_context == Makara::Context.get_current
-
+    def stick_to_master(method_name, args)
       # check to see if we're configured, bypassed, or some custom implementation has input
       return unless should_stick?(method_name, args)
 
       # do the sticking
-      stick_to_master!(write_to_cache)
+      stick_to_master!
     end
 
-
-    # if we are configured to be sticky and we aren't bypassing stickiness
+    # For the generic proxy implementation, we stick if we are sticky,
+    # method and args don't matter
     def should_stick?(method_name, args)
-      @sticky && !@skip_sticking
+      sticky?
     end
 
+    # If we are configured to be sticky and we aren't bypassing stickiness,
+    def sticky?
+      @sticky && !@skip_sticking
+    end
 
     # use the config parser to generate a master and slave pool
     def instantiate_connections
@@ -292,7 +313,8 @@ module Makara
       yield
     rescue ::Makara::Errors::NoConnectionsAvailable => e
       if e.role == 'master'
-        Kernel.raise ::Makara::Errors::NoConnectionsAvailable.new('master and slave')
+        # this means slave connections are good.
+        return
       end
       @slave_pool.disabled = true
       yield

data/lib/makara/railtie.rb

@@ -5,11 +5,5 @@ module Makara
       app.middleware.use Makara::Middleware
     end
 
-    initializer "makara.initialize_logger" do |app|
-      ActiveRecord::LogSubscriber.log_subscribers.each do |subscriber|
-        subscriber.extend ::Makara::Logging::Subscriber
-      end
-    end
-
   end
 end

data/lib/makara/strategies/round_robin.rb

@@ -23,6 +23,8 @@ module Makara
       end
 
       def next
+        return safe_value(0, true) if single_one?
+
         idx = @current_idx
         begin
 
@@ -67,6 +69,10 @@ module Makara
       def should_shuffle?
         true
       end
+
+      def single_one?
+        false
+      end
     end
   end
 end

data/lib/makara/strategies/shard_aware.rb

@@ -0,0 +1,47 @@
+require 'makara/errors/invalid_shard'
+
+module Makara
+  module Strategies
+    class ShardAware < ::Makara::Strategies::Abstract
+
+      def init
+        @shards = {}
+        @default_shard = pool.default_shard
+      end
+
+      def connection_added(wrapper)
+        id = wrapper._makara_shard_id
+        shard_strategy(id).connection_added(wrapper)
+      end
+
+      def shard_strategy(shard_id)
+        id = shard_id
+        shard_strategy = @shards[id]
+        unless shard_strategy
+          shard_strategy = pool.shard_strategy_class.new(pool)
+          @shards[id] = shard_strategy
+        end
+        shard_strategy
+      end
+
+      def current
+        id = shard_id
+        raise Makara::Errors::InvalidShard.new(pool.role, id) unless id && @shards[id]
+
+        @shards[id].current
+      end
+
+      def next
+        id = shard_id
+        raise Makara::Errors::InvalidShard.new(pool.role, id) unless id && @shards[id]
+
+        @shards[id].next
+      end
+
+      def shard_id
+        Thread.current['makara_shard_id'] || pool.default_shard
+      end
+
+    end
+  end
+end

data/lib/makara/version.rb

@@ -2,8 +2,8 @@ module Makara
   module VERSION
 
     MAJOR = 0
-    MINOR = 3
-    PATCH = 8
+    MINOR = 5
+    PATCH = 0
     PRE = nil
 
     def self.to_s

data/makara.gemspec

@@ -6,7 +6,11 @@ Gem::Specification.new do |gem|
   gem.email         = ["mike@mikeonrails.com"]
   gem.description   = %q{Read-write split your DB yo}
   gem.summary       = %q{Read-write split your DB yo}
-  gem.homepage      = ""
+  gem.homepage      = "https://github.com/taskrabbit/makara"
+  gem.licenses      = ['MIT']
+  gem.metadata      = {
+                        "source_code_uri" => 'https://github.com/taskrabbit/makara'
+                      }
 
   gem.files         = `git ls-files`.split($\)
   gem.executables   = gem.files.grep(%r{^bin/}).map{ |f| File.basename(f) }

data/spec/active_record/connection_adapters/makara_abstract_adapter_spec.rb

@@ -27,7 +27,16 @@ describe ActiveRecord::ConnectionAdapters::MakaraAbstractAdapter do
     '    select * from users for update' => true,
     'select * from users lock in share mode' => true,
     'select * from users where name = "for update"' => false,
-    'select * from users where name = "lock in share mode"' => false
+    'select * from users where name = "lock in share mode"' => false,
+    'select nextval(\'users_id_seq\')' => true,
+    'select currval(\'users_id_seq\')' => true,
+    'select lastval()' => true,
+    'with fence as (select * from users) select * from fence' => false,
+    'with fence as (select * from felines) insert to cats' => true,
+    'select get_lock(\'foo\', 0)' => true,
+    'select release_lock(\'foo\')' => true,
+    'select pg_advisory_lock(12345)' => true,
+    'select pg_advisory_unlock(12345)' => true
   }.each do |sql, should_go_to_master|
 
     it "determines that \"#{sql}\" #{should_go_to_master ? 'requires' : 'does not require'} master" do
@@ -65,10 +74,6 @@ describe ActiveRecord::ConnectionAdapters::MakaraAbstractAdapter do
       end
 
       proxy.execute(sql)
-
-      if should_send_to_all_connections
-        expect(proxy.master_context).to be_nil
-      end
     end
 
   end