makara 0.3.8 → 0.5.0

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
 
@@ -127,7 +128,7 @@ module Makara
     # NOTE: Does not use ENV['DATABASE_URL']
     def self.merge_and_resolve_default_url_config(config)
       if ENV['DATABASE_URL']
-        Logging::Logger.log "Please rename DATABASE_URL to use in the database.yml", :warn
+        Makara::Logging::Logger.log "Please rename DATABASE_URL to use in the database.yml", :warn
       end
       return config unless config.key?(:url)
       url = config[:url]
@@ -144,7 +145,7 @@ module Makara
       @config = config.symbolize_keys
       @makara_config = DEFAULTS.merge(@config[:makara] || {})
       @makara_config = @makara_config.symbolize_keys
-      @id = @makara_config[:id]
+      @id = sanitize_id(@makara_config[:id])
     end
 
 
@@ -197,5 +198,15 @@ module Makara
 
     end
 
+
+    def sanitize_id(id)
+      return if id.nil? || id.empty?
+
+      id.gsub(/[\|:]/, '').tap do |sanitized_id|
+        if sanitized_id.size != id.size
+          Makara::Logging::Logger.log "Proxy id '#{id}' changed to '#{sanitized_id}'", :warn
+        end
+      end
+    end
   end
 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
@@ -146,10 +154,10 @@ module Makara
       # The new definition should allow for the proxy to intercept the invocation if required.
       @proxy.class.hijack_methods.each do |meth|
         extension << %Q{
-          def #{meth}(*args)
+          def #{meth}(*args, &block)
             _makara_hijack do |proxy|
               if proxy
-                proxy.#{meth}(*args)
+                proxy.#{meth}(*args, &block)
               else
                 super
               end
@@ -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/context.rb

@@ -1,63 +1,134 @@
 require 'digest/md5'
 
-# Keeps track of the current and previous context (hexdigests)
-# If a new context is needed it can be generated via Makara::Context.generate
-
+# Keeps track of the current stickiness state for different Makara proxies
 module Makara
   class Context
-    class << self
+    attr_accessor :stored_data, :staged_data
+    attr_reader :current_timestamp
+
+    def initialize(context_data)
+      @stored_data = context_data
+      @staged_data = {}
+      @dirty = @was_dirty = false
+
+      freeze_time
+    end
 
-      def generate(seed = nil)
-        seed ||= "#{Time.now.to_i}#{Thread.current.object_id}#{rand(99999)}"
-        Digest::MD5.hexdigest(seed)
+    def stage(proxy_id, ttl)
+      staged_data[proxy_id] = [staged_data[proxy_id].to_f, ttl.to_f].max
+    end
+
+    def stuck?(proxy_id)
+      stored_data[proxy_id] && !expired?(stored_data[proxy_id])
+    end
+
+    def staged?(proxy_id)
+      staged_data.key?(proxy_id)
+    end
+
+    def release(proxy_id)
+      @dirty ||= !!stored_data.delete(proxy_id)
+      staged_data.delete(proxy_id)
+    end
+
+    def release_all
+      if self.stored_data.any?
+        self.stored_data = {}
+        # We need to track a change made to the current stored data
+        # so we can commit it later
+        @dirty = true
       end
+      self.staged_data = {}
+    end
+
+    # Stores the staged data with an expiration time based on the current time,
+    # and clears any expired entries. Returns true if any changes were made to
+    # the current store
+    def commit
+      freeze_time
+      release_expired
+      store_staged_data
+      clean
 
-      def get_previous
-        fetch(:makara_context_previous) { generate }
+      was_dirty?
+    end
+
+    private
+
+    def freeze_time
+      @current_timestamp = Time.now.to_f
+    end
+
+    # Indicates whether there have been changes to the context that need
+    # to be persisted when the request finishes
+    def dirty?
+      @dirty
+    end
+
+    def was_dirty?
+      @was_dirty
+    end
+
+    def expired?(timestamp)
+      timestamp <= current_timestamp
+    end
+
+    def release_expired
+      previous_size = stored_data.size
+      stored_data.delete_if { |_, timestamp| expired?(timestamp) }
+      @dirty ||= previous_size != stored_data.size
+    end
+
+    def store_staged_data
+      staged_data.each do |proxy_id, ttl|
+        if ttl > 0 && self.stored_data[proxy_id].to_f < current_timestamp + ttl
+          self.stored_data[proxy_id] = current_timestamp + ttl
+          @dirty = true
+        end
       end
+    end
 
-      def set_previous(context)
-        previously_sticky.clear
-        set(:makara_context_previous,context)
+    def clean
+      @was_dirty = dirty?
+      @dirty = false
+      @staged_data = {}
+    end
+
+    class << self
+      def set_current(context_data)
+        set(:makara_current_context, new(context_data))
       end
 
-      def get_current
-        fetch(:makara_context_current) { generate }
+      # Called by `Proxy#stick_to_master!` to use master in subsequent requests
+      def stick(proxy_id, ttl)
+        current.stage(proxy_id, ttl)
       end
 
-      def set_current(context)
-        set(:makara_context_current,context)
+      def stuck?(proxy_id)
+        current.staged?(proxy_id) || current.stuck?(proxy_id)
       end
 
-      def previously_stuck?(config_id)
-        previously_sticky.fetch(config_id) do
-          stuck?(Makara::Context.get_previous, config_id)
+      def next
+        if current.commit
+          current.stored_data
         end
       end
 
-      # Called by `Proxy#stick_to_master!` to stick subsequent requests to
-      # master. They'll see the current context as their previous context
-      # when they're asking whether they should be stuck to master.
-      def stick(context, config_id, ttl)
-        Makara::Cache.write(cache_key_for(context, config_id), '1', ttl)
+      def release(proxy_id)
+        current.release(proxy_id)
       end
 
-      def stuck?(context, config_id)
-        !!Makara::Cache.read(cache_key_for(context, config_id))
+      def release_all
+        current.release_all
       end
 
       protected
-
-      def previously_sticky
-        fetch(:makara_previously_sticky) { Hash.new }
-      end
-
-      def cache_key_for(context, config_id)
-        "makara::#{context}-#{config_id}"
+      def current
+        fetch(:makara_current_context) { new({}) }
       end
 
       def fetch(key)
-        get(key) || set(key,yield)
+        get(key) || set(key, yield)
       end
 
       if Thread.current.respond_to?(:thread_variable_get)
@@ -65,19 +136,18 @@ module Makara
           Thread.current.thread_variable_get(key)
         end
 
-        def set(key,value)
-          Thread.current.thread_variable_set(key,value)
+        def set(key, value)
+          Thread.current.thread_variable_set(key, value)
         end
       else
         def get(key)
           Thread.current[key]
         end
 
-        def set(key,value)
+        def set(key, value)
           Thread.current[key]=value
         end
       end
-
     end
   end
 end

data/lib/makara/cookie.rb

@@ -0,0 +1,52 @@
+module Makara
+  module Cookie
+    extend self
+
+    IDENTIFIER = '_mkra_stck'.freeze
+
+    DEFAULT_OPTIONS = {
+      path: "/",
+      http_only: true
+    }
+
+    MAX_AGE_BUFFER = 5
+
+    def fetch(request)
+      parse(request.cookies[IDENTIFIER].to_s)
+    end
+
+    def store(context_data, headers, options = {})
+      unless context_data.nil?
+        Rack::Utils.set_cookie_header! headers, IDENTIFIER, build_cookie(context_data, options)
+      end
+    end
+
+    private
+    # Pairs of {proxy_id}:{timestamp}, separated by "|"
+    # proxy_id1:1518270031.3132212|proxy_id2:1518270030.313232 ..
+    def parse(cookie_string)
+      return {} if cookie_string.empty?
+
+      states = cookie_string.split("|")
+      states.each_with_object({}) do |state, context_data|
+        proxy_id, timestamp = state.split(":")
+        context_data[proxy_id] = timestamp.to_f if proxy_id && timestamp
+      end
+    end
+
+    def build_cookie(context_data, options)
+      cookie = DEFAULT_OPTIONS.merge(options)
+      now = Time.now
+
+      cookie[:max_age] = if context_data.any?
+        (context_data.values.max - now.to_f).ceil + MAX_AGE_BUFFER
+      else
+        0
+      end
+      cookie[:expires] = now + cookie[:max_age]
+      cookie[:value] = context_data.collect { |proxy_id, ttl| "#{proxy_id}:#{ttl}" }.join('|')
+
+      cookie
+    end
+  end
+end

data/lib/makara/error_handler.rb

@@ -25,13 +25,13 @@ module Makara
 
     def gracefully(connection, e)
       err = Makara::Errors::BlacklistConnection.new(connection, e)
-      ::Makara::Logging::Logger.log("[Makara] Gracefully handling: #{err}")
+      ::Makara::Logging::Logger.log("Gracefully handling: #{err}")
       raise err
     end
 
 
     def harshly(e)
-      ::Makara::Logging::Logger.log("[Makara] Harshly handling: #{e}\n#{e.backtrace.join("\n\t")}")
+      ::Makara::Logging::Logger.log("Harshly handling: #{e}\n#{e.backtrace.join("\n\t")}")
       raise e
     end
 

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/logging/logger.rb

@@ -5,7 +5,7 @@ module Makara
       class << self
 
         def log(msg, format = :error)
-          logger.send(format, msg) if logger
+          logger.send(format, "[Makara] #{msg}") if logger
         end
 
         def logger

data/lib/makara/middleware.rb

@@ -1,37 +1,20 @@
 require 'rack'
 
-# Persists the Makara::Context across requests ensuring the same master pool is used on the subsequent request.
-# Simply sets the cookie with the current context and the status code of this request. The next request then sets
-# the Makara::Context's previous context based on the the previous request. If a redirect is encountered the middleware
-# will defer generation of a new context until a non-redirect request occurs.
-
+# Persists the Makara::Context across requests ensuring the same master pool is used on subsequent requests.
 module Makara
   class Middleware
 
-    IDENTIFIER = '_mkra_ctxt'
-
-    DEFAULT_COOKIE = {
-      :path => '/',
-      :http_only => true,
-      :max_age => '5'
-    }
-
     def initialize(app, cookie_options = {})
       @app = app
-      @cookie = DEFAULT_COOKIE.merge(cookie_options)
+      @cookie_options = cookie_options
     end
 
-
     def call(env)
-
       return @app.call(env) if ignore_request?(env)
-
-      Makara::Context.set_previous previous_context(env)
-      Makara::Context.set_current new_context(env)
+      set_current_context(env)
 
       status, headers, body = @app.call(env)
-
-      store_context(status, headers)
+      store_new_context(headers)
 
       [status, headers, body]
     end
@@ -39,6 +22,14 @@ module Makara
 
     protected
 
+    def set_current_context(env)
+      context_data = Makara::Cookie.fetch(Rack::Request.new(env))
+      Makara::Context.set_current(context_data)
+    end
+
+    def store_new_context(headers)
+      Makara::Cookie.store(Makara::Context.next, headers, @cookie_options)
+    end
 
     # ignore asset paths
     # consider allowing a filter proc to be provided in an initializer
@@ -50,59 +41,5 @@ module Makara
       end
       false
     end
-
-
-    # generate a new context based on the request
-    # if the previous request was a redirect, we keep the same context
-    def new_context(env)
-
-      makara_context, makara_status = makara_values(env)
-
-      context = nil
-
-      # if the previous request was a redirect, let's keep that context
-      if makara_status.to_s =~ /^3/ # 300+ redirect
-        context = makara_context
-      end
-
-      context ||= Makara::Context.get_current if env['rack.test']
-      context ||= Makara::Context.generate(env["action_dispatch.request_id"])
-      context
-    end
-
-
-    # pulls the previous context out of the request
-    def previous_context(env)
-      context = makara_values(env).first
-      context ||= Makara::Context.get_previous if env['rack.test']
-      context ||= Makara::Context.generate
-      context
-    end
-
-
-    # retrieve the stored content from the cookie or query
-    # The value contains the hexdigest and status code of the previous
-    # response in the format: $digest--$status
-    def makara_values(env)
-      regex = /#{IDENTIFIER}=([\-a-z0-9A-Z]+)/
-
-      env['HTTP_COOKIE'].to_s =~ regex
-      return $1.split('--') if $1
-
-      env['QUERY_STRING'].to_s =~ regex
-      return $1.split('--') if $1
-
-      [nil, nil]
-    end
-
-
-    # push the current context into the cookie
-    # it should always be for the same path, only
-    # accessible via http and live for a short amount
-    # of time
-    def store_context(status, headers)
-      cookie = @cookie.merge(:value => "#{Makara::Context.get_current}--#{status}")
-      Rack::Utils.set_cookie_header! headers, IDENTIFIER, cookie
-    end
   end
 end