makara 0.4.1 → 0.5.0

data/lib/makara.rb

@@ -14,10 +14,12 @@ module Makara
   autoload :Proxy,              'makara/proxy'
 
   module Errors
-    autoload :MakaraError,                'makara/errors/makara_error'
-    autoload :AllConnectionsBlacklisted,  'makara/errors/all_connections_blacklisted'
-    autoload :BlacklistConnection,        'makara/errors/blacklist_connection'
-    autoload :NoConnectionsAvailable,     'makara/errors/no_connections_available'
+    autoload :MakaraError,                   'makara/errors/makara_error'
+    autoload :AllConnectionsBlacklisted,     'makara/errors/all_connections_blacklisted'
+    autoload :BlacklistConnection,           'makara/errors/blacklist_connection'
+    autoload :NoConnectionsAvailable,        'makara/errors/no_connections_available'
+    autoload :BlacklistedWhileInTransaction, 'makara/errors/blacklisted_while_in_transaction'
+    autoload :InvalidShard,                  'makara/errors/invalid_shard'
   end
 
   module Logging
@@ -29,6 +31,7 @@ module Makara
     autoload :Abstract,         'makara/strategies/abstract'
     autoload :RoundRobin,       'makara/strategies/round_robin'
     autoload :PriorityFailover, 'makara/strategies/priority_failover'
+    autoload :ShardAware,       'makara/strategies/shard_aware'
   end
 
 end

data/lib/makara/config_parser.rb

@@ -1,6 +1,7 @@
 require 'digest/md5'
 require 'active_support/core_ext/hash/keys'
 require 'active_support/core_ext/hash/except'
+require 'cgi'
 
 # Convenience methods to grab subconfigs out of the primary configuration.
 # Provides a way to generate a consistent ID based on a unique config.
@@ -63,7 +64,7 @@ module Makara
       # Converts the given URL to a full connection hash.
       def to_hash
         config = raw_config.reject { |_,value| value.blank? }
-        config.map { |key,value| config[key] = URI.unescape(value) if value.is_a? String }
+        config.map { |key,value| config[key] = CGI.unescape(value) if value.is_a? String }
         config
       end
 

data/lib/makara/connection_wrapper.rb

@@ -35,11 +35,19 @@ module Makara
       @config[:name]
     end
 
+    def _makara_shard_id
+      @config[:shard_id]
+    end
+
     # has this node been blacklisted?
     def _makara_blacklisted?
       @blacklisted_until.present? && @blacklisted_until.to_i > Time.now.to_i
     end
 
+    def _makara_in_transaction?
+      @connection && @connection.open_transactions > 0 ? true : false
+    end
+
     # blacklist this node for @config[:blacklist_duration] seconds
     def _makara_blacklist!
       @connection.disconnect! if @connection
@@ -158,6 +166,22 @@ module Makara
         }
       end
 
+      # Control methods must always be passed to the
+      # Makara::Proxy control object for handling (typically
+      # related to ActiveRecord connection pool management)
+      @proxy.class.control_methods.each do |meth|
+        extension << %Q{
+          def #{meth}(*args, &block)
+            proxy = _makara
+            if proxy
+              proxy.control.#{meth}(*args=args, block)
+            else
+              super # Only if we are not wrapped any longer
+            end
+          end
+        }
+      end
+
       # extend the instance
       con.instance_eval(extension)
       # set the makara context

data/lib/makara/errors/blacklisted_while_in_transaction.rb

@@ -0,0 +1,14 @@
+module Makara
+  module Errors
+    class BlacklistedWhileInTransaction < MakaraError
+
+      attr_reader :role
+
+      def initialize(role)
+        @role = role
+        super "[Makara] Blacklisted while in transaction in the #{role} pool"
+      end
+
+    end
+  end
+end

data/lib/makara/errors/invalid_shard.rb

@@ -0,0 +1,16 @@
+module Makara
+  module Errors
+    class InvalidShard < MakaraError
+
+      attr_reader :role
+      attr_reader :shard_id
+
+      def initialize(role, shard_id)
+        @role = role
+        @shard_id = shard_id
+        super "[Makara] Invalid shard_id #{shard_id} for the #{role} pool"
+      end
+
+    end
+  end
+end

data/lib/makara/pool.rb

@@ -1,4 +1,5 @@
 require 'active_support/core_ext/hash/keys'
+require 'makara/strategies/shard_aware'
 
 # Wraps a collection of similar connections and chooses which one to use
 # Provides convenience methods for accessing underlying connections
@@ -13,6 +14,8 @@ 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
@@ -20,7 +23,13 @@ module Makara
       @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
 
 
@@ -91,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
 
 
@@ -135,8 +158,8 @@ module Makara
     # to be sticky, provide back the current connection assuming it is
     # not blacklisted.
     def next
-      if @proxy.sticky && @strategy.current
-        @strategy.current
+      if @proxy.sticky && (curr = @strategy.current)
+        curr
       else
         @strategy.next
       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
@@ -84,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, ''
@@ -148,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
@@ -286,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/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 = 4
-    PATCH = 1
+    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_postgresql_adapter_spec.rb

@@ -1,5 +1,6 @@
 require 'spec_helper'
 require 'active_record/connection_adapters/postgresql_adapter'
+require 'active_record/errors'
 
 describe 'MakaraPostgreSQLAdapter' do
 
@@ -191,4 +192,60 @@ describe 'MakaraPostgreSQLAdapter' do
       it_behaves_like 'a transaction supporter'
     end
   end
+
+  context 'with two activerecord connection pools' do
+
+    before :each do
+      class Model1 < ActiveRecord::Base
+      end
+
+      class Model2 < ActiveRecord::Base
+      end
+
+      Model1.establish_connection(config)
+      Model2.establish_connection(config)
+
+    end
+
+    it 'should not leak raw connection into activerecord pool' do
+      # checkout a connection from Model1 pool and remove from the pool
+      conn = Model1.connection_pool.checkout
+      Model1.connection_pool.remove(conn)
+
+      # assign the connection to Model2 pool
+      conn.pool=Model2.connection_pool
+
+      # now close the connection to return it back to the pool
+      conn.close
+
+      # checkout the connection and make sure it is still a makara proxy
+      expect(Model2.connection).to eq(conn)
+    end
+
+    it 'should be able to steal the connection from a different thread' do
+      conn = Model1.connection_pool.checkout
+      conn.steal!
+      expect(conn.owner).to eq(Thread.current)
+      # steal! is not thread safe. it should be done while holding connection pool's mutex
+      t = Thread.new { conn.steal! }
+      t.join
+      expect(conn.owner).to eq(t)
+    end
+
+    it 'should not be able to expire the connection from same thread' do
+      conn = Model2.connection_pool.checkout
+      # expire is not thread safe. it should be done while holding connection pool's mutex
+      expect {
+        t = Thread.new { conn.expire }
+        t.join
+      }.to raise_error(ActiveRecord::ActiveRecordError)
+    end
+
+    it 'should be able to checkin connection back into activerecord pool' do
+      conn = Model1.connection_pool.checkout
+      Model1.connection_pool.checkin(conn)
+      # checkout the connection again and make sure it is same connection
+      expect(Model1.connection).to eq(conn)
+    end
+  end
 end

data/spec/cookie_spec.rb

@@ -53,20 +53,20 @@ describe Makara::Cookie do
       Makara::Cookie.store(context_data, headers)
 
       expect(headers['Set-Cookie']).to include("#{cookie_key}=mysql%3A#{(now + 5).to_f}%7Credis%3A#{(now + 5).to_f};")
-      expect(headers['Set-Cookie']).to include("path=/; max-age=10; expires=#{(Time.now + 10).gmtime.rfc2822}; HttpOnly")
+      expect(headers['Set-Cookie']).to include("path=/; max-age=10; expires=#{(Time.now + 10).httpdate}; HttpOnly")
     end
 
     it 'expires the cookie if the next context is empty' do
       Makara::Cookie.store({}, headers)
 
-      expect(headers['Set-Cookie']).to eq("#{cookie_key}=; path=/; max-age=0; expires=#{Time.now.gmtime.rfc2822}; HttpOnly")
+      expect(headers['Set-Cookie']).to eq("#{cookie_key}=; path=/; max-age=0; expires=#{Time.now.httpdate}; HttpOnly")
     end
 
     it 'allows custom cookie options to be provided' do
       Makara::Cookie.store(context_data, headers, { :secure => true })
 
       expect(headers['Set-Cookie']).to include("#{cookie_key}=mysql%3A#{(now + 5).to_f}%7Credis%3A#{(now + 5).to_f};")
-      expect(headers['Set-Cookie']).to include("path=/; max-age=10; expires=#{(Time.now + 10).gmtime.rfc2822}; secure; HttpOnly")
+      expect(headers['Set-Cookie']).to include("path=/; max-age=10; expires=#{(Time.now + 10).httpdate}; secure; HttpOnly")
     end
   end
 end