promiscuous 0.90.0 → 0.91.0

data/lib/promiscuous/publisher/operation/base.rb

@@ -1,41 +1,31 @@
 class Promiscuous::Publisher::Operation::Base
-  class TryAgain < RuntimeError; end
-  VERSION_FIELD = '_pv'
+  mattr_accessor :recovery_mechanisms
+  self.recovery_mechanisms = []
 
-  attr_accessor :operation, :operation_ext, :instance, :selector_keys
-
-  def initialize(options={})
-    # XXX instance is not always an instance, it can be a selector
-    # representation.
-    @instance      = options[:instance]
-    @operation     = options[:operation]
-    @operation_ext = options[:operation_ext]
-    @multi         = options[:multi]
+  def self.register_recovery_mechanism(method_name=nil, &block)
+    self.recovery_mechanisms << (block || method(method_name))
   end
 
-  def read?
-    operation == :read
+  def self.run_recovery_mechanisms
+    self.recovery_mechanisms.each(&:call)
   end
 
-  def write?
-    !read?
-  end
+  attr_accessor :operation
 
-  def multi?
-    !!@multi
+  def initialize(options={})
+    @operation = options[:operation]
   end
 
-  def single?
-    !@multi
+  def read?
+    @operation == :read
   end
 
-  def persists?
-    # TODO For writes in transactions, it should be false
-    write?
+  def write?
+    !read?
   end
 
-  def failed?
-    !!@exception
+  def recovering?
+    !!@recovering
   end
 
   def current_context
@@ -67,7 +57,7 @@ class Promiscuous::Publisher::Operation::Base
   end
 
   def publish_payload_in_rabbitmq_async
-    Promiscuous::AMQP.publish(:key => @amqp_key, :payload => @payload,
+    Promiscuous::AMQP.publish(:key => Promiscuous::Config.app, :payload => @payload,
                               :on_confirm => method(:on_rabbitmq_confirm))
   end
 
@@ -87,33 +77,37 @@ class Promiscuous::Publisher::Operation::Base
         node.zadd(rabbitmq_staging_set_key, Time.now.to_i, key)
         payload = node.get(key)
 
-        Promiscuous.info "[payload recovery] #{payload}"
-        new.instance_eval do
-          @payload_recovery_node = node
-          @payload_recovery_key = key
-          @amqp_key = MultiJson.load(payload)['__amqp__']
-          @payload = payload
-          publish_payload_in_rabbitmq_async
+        # It's possible that the payload is nil as the message could be
+        # recovered by another worker
+        if payload
+          Promiscuous.info "[payload recovery] #{payload}"
+          new.instance_eval do
+            @payload_recovery_node = node
+            @payload_recovery_key = key
+            @payload = payload
+            @recovery = true
+            publish_payload_in_rabbitmq_async
+          end
         end
       end
     end
   end
+  register_recovery_mechanism :recover_payloads_for_rabbitmq
 
   def publish_payload_in_redis
     # TODO Optimize and DRY this up
     r = @committed_read_deps
     w = @committed_write_deps
 
-    master_node = w.first.redis_node
-    operation_recovery_key = w.first.key(:pub).join('operation_recovery').to_s
-    # We identify a payload with a unique key (id:id_value:current_version) to
-    # avoid collisions with other updates on the same document.
+    # We identify a payload with a unique key (id:id_value:current_version:payload_recovery)
+    # to avoid collisions with other updates on the same document.
+    master_node = @op_lock.node
     @payload_recovery_node = master_node
-    @payload_recovery_key = w.first.key(:pub).join(w.first.version).to_s
+    @payload_recovery_key = Promiscuous::Key.new(:pub).join('payload_recovery', @op_lock.token).to_s
 
     # We need to be able to recover from a redis failure. By sending the
     # payload to the slave first, we ensure that we can replay the lost
-    # payloads if the primary came to fail.
+    # payloads if the master came to fail.
     # We still need to recover the lost operations. This can be done by doing a
     # version diff from what is stored in the database and the recovered redis slave.
     # XXX TODO
@@ -123,113 +117,108 @@ class Promiscuous::Publisher::Operation::Base
     # happen if we lost the lock without knowing about it.
     # The payload can be sent twice, which is okay since the subscribers
     # tolerate it.
+    operation_recovery_key = "#{@op_lock.key}:operation_recovery"
+    versions_recovery_key = "#{operation_recovery_key}:versions"
 
-    nodes = (w+r).map(&:redis_node).uniq
-    if nodes.size == 1
-      # We just have the master node. Since we are atomic, we don't need to do
-      # the 2pc dance.
-      master_node.multi do
-        master_node.del(operation_recovery_key)
-        master_node.set(@payload_recovery_key, @payload)
-        master_node.zadd(rabbitmq_staging_set_key, Time.now.to_i, @payload_recovery_key)
-      end
-    else
-      master_node.multi do
-        master_node.set(@payload_recovery_key, @payload)
-        master_node.zadd(rabbitmq_staging_set_key, Time.now.to_i, @payload_recovery_key)
-      end
-
-      # The payload is safe now. We can cleanup all the versions on the
-      # secondary. Note that we need to clear the master node at the end,
-      # as it acts as a lock on the other keys. This is important to avoid a
-      # race where we would delete data that doesn't belong to the current
-      # operation due to a lock loss.
-      nodes.reject { |node| node == master_node }
-            .each  { |node| node.del(operation_recovery_key) }
+    master_node.multi do
+      master_node.set(@payload_recovery_key, @payload)
+      master_node.zadd(rabbitmq_staging_set_key, Time.now.to_i, @payload_recovery_key)
       master_node.del(operation_recovery_key)
+      master_node.del(versions_recovery_key)
     end
+
+    # The payload is safe now. We can cleanup all the versions on the
+    # secondary. There are no harmful races that can happen since the
+    # secondary_operation_recovery_key is unique to the operation.
+    # XXX The caveat is that if we die here, the
+    # secondary_operation_recovery_key will never be cleaned up.
+    (w+r).map(&:redis_node).uniq
+      .reject { |node| node == master_node }
+      .each   { |node| node.del(versions_recovery_key) }
   end
 
-  def generate_payload_and_clear_operations
-    # TODO Transactions with multi writes
-    raise "We don't support multi writes yet" if previous_successful_operations.select(&:write?).size > 1
-    raise "The instance is gone, or there is a version mismatch" unless @instance
+  def payload_for(instance)
+    options = { :with_attributes => self.operation.in?([:create, :update]) }
+    instance.promiscuous.payload(options).tap do |payload|
+      payload[:operation] = self.operation
+    end
+  end
 
-    payload = @instance.promiscuous.payload(:with_attributes => operation.in?([:create, :update]))
+  def generate_payload
+    payload = {}
+    payload[:operations] = operation_payloads
     payload[:context] = current_context.name
+    payload[:app] = Promiscuous::Config.app
     payload[:timestamp] = @timestamp
-
-    # If the db operation has failed, so we publish a dummy operation on the
-    # failed instance. It's better than using the Dummy polisher class
-    # because a subscriber can choose not to receive any of these messages.
-    payload[:operation] = self.failed? ? :dummy : operation
-
-    # We need to consider the last write operation as an implicit read
-    # dependency. This is why we don't need to consider the read dependencies
-    # happening before a first write when publishing the second write in a
-    # context.
+    payload[:host] = Socket.gethostname
+    payload[:current_user_id] = Thread.current[:promiscuous_context].try(:current_user_id)
     payload[:dependencies] = {}
     payload[:dependencies][:read]  = @committed_read_deps if @committed_read_deps.present?
     payload[:dependencies][:write] = @committed_write_deps
 
-    current_context.last_write_dependency = @committed_write_deps.first
-    current_context.operations.clear
-
-    @amqp_key = payload[:__amqp__]
     @payload = MultiJson.dump(payload)
   end
 
+  def clear_previous_dependencies
+    current_context.read_operations.clear
+    current_context.extra_dependencies = [@committed_write_deps.first]
+  end
+
   def self.recover_operation_from_lock(lock)
     # We happen to have acquired a never released lock.
     # The database instance is thus still prestine.
-    # Three cases to consider:
-    # 1) the key is not an id dependency or the payload queue stage was passed
-    # 2) The write query was never executed, we must send a dummy operation
-    # 3) The write query was executed, but never passed the payload queue stage
 
     master_node = lock.node
-    recovery_data = master_node.hgetall("#{lock.key}:operation_recovery")
-    return nil unless recovery_data.present? # case 1)
+    recovery_data = master_node.get("#{lock.key}:operation_recovery")
+
+    unless recovery_data.present?
+      lock.unlock
+      return
+    end
 
     Promiscuous.info "[operation recovery] #{lock.key} -> #{recovery_data}"
 
-    collection, instance_id, operation,
-      document, read_dependencies, write_dependencies = *MultiJson.load(recovery_data['payload'])
+    op_klass, operation, read_dependencies,
+      write_dependencies, recovery_arguments = *MultiJson.load(recovery_data)
 
     operation = operation.to_sym
-    read_dependencies.map!  { |k| Promiscuous::Dependency.parse(k.to_s) }
-    write_dependencies.map! { |k| Promiscuous::Dependency.parse(k.to_s) }
+    read_dependencies.map!  { |k| Promiscuous::Dependency.parse(k.to_s, :type => :read) }
+    write_dependencies.map! { |k| Promiscuous::Dependency.parse(k.to_s, :type => :write) }
 
-    model = Promiscuous::Publisher::Model.publishers[collection]
-
-    if model.is_a? Promiscuous::Publisher::Model::Ephemeral
-      operation = :dummy
-    else
-      # TODO Abstract db operations.
-      # We need to query on the root model
-      model = model.collection.name.singularize.camelize.constantize
+    begin
+      op = op_klass.constantize.recover_operation(*recovery_arguments)
+    rescue NameError
+      raise "invalid recover operation class: #{op_klass}"
     end
 
-    op_klass = model.get_operation_class_for(operation)
-    op = op_klass.recover_operation(model, instance_id, document)
-    op.operation = operation
-
-    Promiscuous.context :operation_recovery, :detached_from_parent => true do
-      op.instance_eval do
-        @read_dependencies  = read_dependencies
-        @write_dependencies = write_dependencies
-        @locks = [lock]
-        execute_persistent_locked { recover_db_operation }
+    Thread.new do
+      # We run the recovery in another thread to ensure that we get a new
+      # database connection to avoid tempering with the current state of the
+      # connection, which can be in an open transaction.
+      # Thankfully, we are not in a fast path.
+      # Note that any exceptions will be passed through the thread join() method.
+      Promiscuous.context :operation_recovery do
+        op.instance_eval do
+          @operation = operation
+          @read_dependencies  = read_dependencies
+          @write_dependencies = write_dependencies
+          @op_lock = lock
+          @recovering = true
+
+          query = Promiscuous::Publisher::Operation::ProxyForQuery.new(self) { recover_db_operation }
+          execute_instrumented(query)
+          query.result
+        end
       end
-    end
+    end.join
 
-    lock.unlock
   rescue Exception => e
-    message = "cannot recover #{lock.key} -> #{recovery_data}"
+    message = "cannot recover #{lock.key}, failed to fetch recovery data"
+    message = "cannot recover #{lock.key}, recovery data: #{recovery_data}" if recovery_data
     raise Promiscuous::Error::Recovery.new(message, e)
   end
 
-  def increment_read_and_write_dependencies(read_dependencies, write_dependencies)
+  def increment_read_and_write_dependencies
     # We collapse all operations, ignoring the read/write interleaving.
     # It doesn't matter since all write operations are serialized, so the first
     # write in the transaction can have all the read dependencies.
@@ -241,92 +230,166 @@ class Promiscuous::Publisher::Operation::Base
     # r and w is empty) when it calculates the happens before relationships.
     r -= w
 
-    master_node = w.first.redis_node
-    operation_recovery_key = w.first
+    master_node = @op_lock.node
+    operation_recovery_key = "#{@op_lock.key}:operation_recovery"
 
     # We group all the dependencies by their respective shards
     # The master node will have the responsability to hold the recovery data.
     # We do the master node first. The seconaries can be done in parallel.
-    (w+r).group_by(&:redis_node).each do |node, deps|
-      r_deps = deps.select { |dep| dep.in? r }
-      w_deps = deps.select { |dep| dep.in? w }
+    @committed_read_deps  = []
+    @committed_write_deps = []
+
+    # We need to do the increments always in the same node order, otherwise.
+    # the subscriber can deadlock. But we must always put the recovery payload
+    # on the master before touching anything.
+    nodes_deps = (w+r).group_by(&:redis_node)
+                      .sort_by { |node, deps| -Promiscuous::Redis.master.nodes.index(node) }
+    if nodes_deps.first[0] != master_node
+      nodes_deps = [[master_node, []]] + nodes_deps
+    end
 
+    nodes_deps.each do |node, deps|
       argv = []
       argv << Promiscuous::Key.new(:pub) # key prefixes
-      argv << MultiJson.dump([r_deps, w_deps])
+      argv << operation_recovery_key
+
+      # The index of the first write is then used to pass to redis along with the
+      # dependencies. This is done because arguments to redis LUA scripts cannot
+      # accept complex data types.
+      argv << (deps.index(&:read?) || deps.length)
 
       # Each shard have their own recovery payload. The master recovery node
       # has the full operation recovery, and the others just have their versions.
-      argv << operation_recovery_key.as_json
-      if node == master_node
+      # Note that the operation_recovery_key on the secondaries have the current
+      # version of the instance appended to them. It's easier to cleanup when
+      # locks get lost.
+      if node == master_node && !self.recovering?
         # We are on the master node, which holds the recovery payload
-        document = serialize_document_for_create_recovery if operation == :create
-        argv << MultiJson.dump([@instance.class.promiscuous_collection_name,
-                                @instance.id, operation, document, r, w])
+        argv << MultiJson.dump([self.class.name, operation, r, w, self.recovery_payload])
       end
 
+      # FIXME If the lock is lost, we need to backoff
+
       # We are going to store all the versions in redis, to be able to recover.
       # We store all our increments in a transaction_id key in JSON format.
       # Note that the transaction_id is the id of the current instance.
       @@increment_script ||= Promiscuous::Redis::Script.new <<-SCRIPT
         local prefix = ARGV[1] .. ':'
-        local deps = cjson.decode(ARGV[2])
-        local read_deps = deps[1]
-        local write_deps = deps[2]
-        local operation_recovery_key = prefix .. ARGV[3] .. ':operation_recovery'
+        local operation_recovery_key = ARGV[2]
+        local versions_recovery_key = operation_recovery_key .. ':versions'
+        local first_read_index = tonumber(ARGV[3]) + 1
         local operation_recovery_payload = ARGV[4]
+        local deps = KEYS
 
-        local read_versions = {}
-        local write_versions = {}
+        local versions = {}
 
-        if redis.call('exists', operation_recovery_key) == 1 then
-          for i, dep in ipairs(read_deps) do
-            local key = prefix .. dep
-            read_versions[i] = redis.call('get', key .. ':w')
-          end
-          for i, dep in ipairs(write_deps) do
-            local key = prefix .. dep
-            write_versions[i] = redis.call('get', key .. ':w')
-          end
-        else
-          for i, dep in ipairs(read_deps) do
-            local key = prefix .. dep
-            redis.call('incr', key .. ':rw')
-            read_versions[i] = redis.call('get', key .. ':w')
-            redis.call('hset', operation_recovery_key, dep, read_versions[i])
+        if redis.call('exists', versions_recovery_key) == 1 then
+          first_read_index = tonumber(redis.call('hget', versions_recovery_key, 'read_index'))
+          if not first_read_index then
+            return redis.error_reply('Failed to read dependency index during recovery')
           end
 
-          for i, dep in ipairs(write_deps) do
-            local key = prefix .. dep
-            write_versions[i] = redis.call('incr', key .. ':rw')
-            redis.call('set', key .. ':w', write_versions[i])
-            redis.call('hset', operation_recovery_key, dep, write_versions[i])
+          for i, dep in ipairs(deps) do
+            versions[i] = tonumber(redis.call('hget', versions_recovery_key, dep))
+            if not versions[i] then
+              return redis.error_reply('Failed to read dependency ' .. dep .. ' during recovery')
+            end
           end
 
-          if operation_recovery_payload then
-            redis.call('hset', operation_recovery_key, 'payload', operation_recovery_payload)
+          return { first_read_index-1, versions }
+        end
+
+        if redis.call('exists', prefix .. 'bootstrap') == 1 then
+          first_read_index = #deps + 1
+        end
+
+        if #deps ~= 0 then
+          redis.call('hset', versions_recovery_key, 'read_index', first_read_index)
+        end
+
+        for i, dep in ipairs(deps) do
+          local key = prefix .. dep
+          local rw_version = redis.call('incr', key .. ':rw')
+          if i < first_read_index then
+            redis.call('set', key .. ':w', rw_version)
+            versions[i] = rw_version
+          else
+            versions[i] = tonumber(redis.call('get', key .. ':w')) or 0
           end
+          redis.call('hset', versions_recovery_key, dep, versions[i])
         end
 
-        return { read_versions, write_versions }
+        if operation_recovery_payload then
+          redis.call('set', operation_recovery_key, operation_recovery_payload)
+        end
+
+        return { first_read_index-1, versions }
       SCRIPT
-      read_versions, write_versions = @@increment_script.eval(node, :argv => argv)
 
-      r_deps.zip(read_versions).each  { |dep, version| dep.version = version.to_i }
-      w_deps.zip(write_versions).each { |dep, version| dep.version = version.to_i }
+      first_read_index, versions = @@increment_script.eval(node, :argv => argv, :keys => deps)
+
+      deps.zip(versions).each  { |dep, version| dep.version = version }
+
+      @committed_write_deps += deps[0...first_read_index]
+      @committed_read_deps  += deps[first_read_index..-1]
     end
 
-    @committed_read_deps  = r
-    @committed_write_deps = w
-    @instance_version = w.first.version
+    # The instance version must to be the first in the list to allow atomic
+    # subscribers to do their magic.
+    # TODO What happens with transactions with multiple operations?
+    instance_dep_index = @committed_write_deps.index(write_dependencies.first)
+    @committed_write_deps[0], @committed_write_deps[instance_dep_index] =
+      @committed_write_deps[instance_dep_index], @committed_write_deps[0]
   end
 
-  LOCK_OPTIONS = { :timeout => 10.seconds, # after 10 seconds, we give up
-                   :sleep   => 0.01,       # polling every 10ms.
-                   :expire  => 1.minute }  # after one minute, we are considered dead
-
   def self.lock_options
-    LOCK_OPTIONS.merge({ :lock_set => Promiscuous::Key.new(:pub).join('lock_set').to_s })
+    {
+      :timeout  => 10.seconds,   # after 10 seconds, we give up so we don't queue requests
+      :sleep    => 0.01.seconds, # polling every 10ms.
+      :expire   => 1.minute,     # after one minute, we are considered dead
+      :lock_set => Promiscuous::Key.new(:pub).join('lock_set').to_s
+    }
+  end
+  delegate :lock_options, :to => self
+
+  def dependency_for_op_lock
+    query_dependencies.first
+  end
+
+  def get_new_op_lock
+    dep = dependency_for_op_lock
+    Promiscuous::Redis::Mutex.new(dep.key(:pub).to_s, lock_options.merge(:node => dep.redis_node))
+  end
+
+  def self._acquire_lock(mutex)
+    loop do
+      case mutex.lock
+      # recover_operation_from_lock implicitely unlocks the lock.
+      when :recovered then recover_operation_from_lock(mutex)
+      when true       then return true
+      when false      then return false
+      end
+    end
+  end
+
+  def acquire_op_lock
+    @op_lock = get_new_op_lock
+
+    unless self.class._acquire_lock(@op_lock)
+      raise Promiscuous::Error::LockUnavailable.new(@op_lock.key)
+    end
+  end
+
+  def release_op_lock
+    @op_lock.unlock
+    @op_lock = nil
+  end
+
+  def ensure_op_still_locked
+    unless @op_lock.still_locked?
+      # We lost the lock, let the recovery mechanism do its thing.
+      raise Promiscuous::Error::LostLock.new(@op_lock.key)
+    end
   end
 
   def self.recover_locks
@@ -340,368 +403,110 @@ class Promiscuous::Publisher::Operation::Base
         break unless key && Time.now.to_i >= time.to_i + lock_options[:expire]
 
         mutex = Promiscuous::Redis::Mutex.new(key, lock_options.merge(:node => node))
-        case mutex.lock
-        when :recovered then recover_operation_from_lock(mutex)
-        when true       then mutex.unlock
-        when false      then ;
-        end
+        mutex.unlock if _acquire_lock(mutex)
       end
     end
   end
+  register_recovery_mechanism :recover_locks
 
-  def locks_from_write_dependencies
-    # XXX TODO Support multi row writes
-    instance_dep = write_dependencies.first
-    return [] unless instance_dep
-    options = self.class.lock_options.merge(:node => instance_dep.redis_node)
-    [Promiscuous::Redis::Mutex.new(instance_dep.key(:pub).to_s, options)]
-  end
-
-  def lock_write_dependencies
-    # returns true if we could get all the locks, false otherwise
-
-    start_at = Time.now
-    @recovered_locks = []
-
-    # We acquire all the locks in order, and unlock everything if one come
-    # to fail. lock/unlock return true/false when they succeed/fail
-    locks = locks_from_write_dependencies
-    locks.reduce(->{ @locks = locks; true }) do |chain, l|
-      lambda do
-        return false if Time.now - start_at > LOCK_OPTIONS[:timeout]
-        case l.lock
-          # Note that we do not unlock the recovered lock if the chain fails
-        when :recovered then @recovered_locks << l; chain.call
-        when true       then chain.call or (l.unlock; false)
-        when false      then @unavailable_lock = l; false
-        end
-      end
-    end.call
-  end
-
-  def unlock_write_dependencies
-    # returns true if we could unlock all the locks, false otherwise
-    return true if @locks.blank?
-    @locks.reduce(true) { |result, l| l.unlock && result }.tap { @locks = nil }
-  end
+  def dependencies_for(instance, options={})
+    return [] if instance.nil?
 
-  def _reload_instance_dependencies
     if read?
       # We want to use the smallest subset that we can depend on when doing
       # reads. tracked_dependencies comes sorted from the smallest subset to
       # the largest. For maximum performance on the subscriber side, we thus
       # pick the first one. In most cases, it should resolve to the id
       # dependency.
-      best_dependency = @instance.promiscuous.tracked_dependencies.first
-      unless best_dependency
-        raise Promiscuous::Error::Dependency.new(:operation => self)
-      end
-      [best_dependency]
+      # If we don't have any, the driver should track individual instances.
+      best_dependency = instance.promiscuous.tracked_dependencies(:allow_missing_attributes => true).first
+      [best_dependency].compact
     else
       # Note that tracked_dependencies will not return the id dependency if it
       # doesn't exist which can only happen for create operations and auto
-      # generated ids. Be aware that with auto generated id, create operation
-      # might not provide the id dependency.
-      @instance.promiscuous.tracked_dependencies
+      # generated ids.
+      instance.promiscuous.tracked_dependencies
     end
   end
 
-  def reload_instance_dependencies
-    # Returns true when the dependencies changed, false otherwise
-    @write_dependencies = nil
-    old = @instance_dependencies
-    @instance_dependencies = _reload_instance_dependencies
-    old != @instance_dependencies
-  end
-
-  def instance_dependencies
-    reload_instance_dependencies unless @instance_dependencies
-    @instance_dependencies
-  end
-
-  def previous_successful_operations
-    current_context.operations.reject(&:failed?)
-  end
-
   def read_dependencies
     # We memoize the read dependencies not just for performance, but also
     # because we store the versions once incremented in these.
     return @read_dependencies if @read_dependencies
-    read_dependencies = previous_successful_operations.select(&:read?)
-                             .map(&:instance_dependencies).flatten
+    read_dependencies = current_context.read_operations.map(&:query_dependencies).flatten
 
-    # We implicitly have a read dependency on the latest write.
-    if current_context.last_write_dependency
-      current_context.last_write_dependency.version = nil
-      read_dependencies << current_context.last_write_dependency
+    # We add extra_dependencies, which can contain the latest write, or user
+    # context, etc.
+    current_context.extra_dependencies.each do |dep|
+      dep.version = nil
+      read_dependencies << dep
     end
 
-    @read_dependencies = read_dependencies.uniq
+    @read_dependencies = read_dependencies.uniq.each { |d| d.type = :read }
   end
-  alias verify_read_dependencies read_dependencies
+  alias generate_read_dependencies read_dependencies
 
   def write_dependencies
-    # The cache is cleared when we call reload_instance_dependencies
-    @write_dependencies ||= previous_successful_operations.select(&:write?)
-                              .map(&:instance_dependencies).flatten.uniq
+    @write_dependencies ||= self.query_dependencies.uniq.each { |d| d.type = :write }
   end
 
-  def reload_instance
-    @instance = without_promiscuous { fetch_instance }
+  def should_instrument_query?
+    # current_context is later enforced for writes.
+    !Promiscuous.disabled? && (current_context || write?)
   end
 
-  def perform_db_operation_with_no_exceptions(&db_operation)
-    going_to_execute_db_operation
-    @result = db_operation.call(self)
-  rescue Exception => e
-    @exception = e
-  end
-
-  def lock_instance_for_execute_persistent
-    current_context.add_operation(self)
+  def execute(&query_config)
+    query = Promiscuous::Publisher::Operation::ProxyForQuery.new(self, &query_config)
 
-    # Note: At first, @instance can be a representation of a selector, to
-    # become a real model instance once we get to fetch it from the db with
-    # reload_instance to lock an instance that matches the selector.
-    # This is a good thing because we allow the underlying driver to hook from
-    # the model interface to the driver interface easily.
-    auto_unlock = true
-
-    begin
-      unless lock_write_dependencies
-        raise Promiscuous::Error::LockUnavailable.new(@unavailable_lock.key)
-      end
-
-      if @recovered_locks.present?
-        # When recovering locks, if we fail, we must not release the lock again
-        # to allow another one to do the recovery.
-        auto_unlock = false
-        @recovered_locks.each { |lock| self.class.recover_operation_from_lock(lock) }
-        auto_unlock = true
-        raise TryAgain
-      end
-
-      if operation != :create
-        # We need to lock and update all the dependencies before any other
-        # readers can see our write through any one of our tracked attributes.
-
-        # We want to reload the instance to make sure we have all the locked
-        # dependencies that we need. It's a query we cannot avoid when we have
-        # tracked dependencies. There is a bit of room for optimization.
-        # If the selector doesn't fetch any instance, the query has no effect
-        # so we can bypass it as if nothing happened.  If reload_instance
-        # raises an exception, it's okay to let it bubble up since we haven't
-        # touch anything yet except for the locks (which will be unlocked on
-        # the way out)
-        return false unless reload_instance
-
-        # If reload_instance changed the current instance because the selector,
-        # we need to unlock the old instance, lock this new instance, and
-        # retry. XXX What should we do if we are going in a live lock?
-        # Sleep with some jitter?
-        if reload_instance_dependencies
-          raise TryAgain
-        end
-      end
-    rescue TryAgain
-      unlock_write_dependencies if auto_unlock
-      retry
-    end
-
-    verify_read_dependencies
-    if write_dependencies.blank?
-      # TODO We don't like auto generated ids. A good solution is to do all
-      # writes in a transaction, so we can know the ids at commit time.
-      raise "We don't support auto generated id yet"
-    end
-
-    # We are now in the possession of an instance that matches the original
-    # selector, we can proceed.
-    auto_unlock = false
-    true
-  ensure
-    # In case of an exception was raised before we updated the version in
-    # redis, we can unlock because we don't need recovery.
-    unlock_write_dependencies if auto_unlock
-  end
-
-  def execute_persistent_locked(&db_operation)
-    # We are going to commit all the pending writes in the context if we are
-    # doing a transaction commit. We also commit the current write operation for
-    # atomic writes without transactions.  We enable the recovery mechanism by
-    # having someone expiring our lock if we die in the middle.
-
-    # All the versions are updated and a marked as pending for publish in Redis
-    # atomically in case we die before we could write the versions in the
-    # database. Once incremented, concurrent queries that are reading our
-    # instance will be serialized after our write, even through it may read our
-    # old instance. This is a race that we tolerate.
-    # XXX We also stash the document for create operations, so the recovery can
-    # redo the create to avoid races when instances are getting partitioned.
-    increment_read_and_write_dependencies(read_dependencies, write_dependencies)
-
-    # From this point, if we die, the one expiring our write locks must finish
-    # the publish, either by sending a dummy, or by sending the real instance.
-    # We could have die before or after the database query.
-
-    # We save the versions in the database, as it is our source of truth.
-    # This allow a reconstruction of redis in the face of failures.
-    # We would also need to send a special message to the subscribers to reset
-    # their read counters to the last write version since we would not be able
-    # to restore the read counters (and we don't want to store them because
-    # this would dramatically augment our footprint on the db).
-    #
-    # If we are doing a destroy operation, and redis dies right after, and
-    # we happen to lost contact with rabbitmq, recovery is going to be complex:
-    # we would need to do a diff from the dummy subscriber to see what
-    # documents are missing on our side to be able to resend the destroy
-    # message.
-
-    case operation
-    when :create
-      stash_version_in_write_query
-    when :update
-      stash_version_in_write_query
-      # We are now in the possession of an instance that matches the original
-      # selector. We need to make sure the db_operation will operate on it,
-      # instead of the original selector.
-      use_id_selector(:use_atomic_version_selector => true)
-      # We need to use an atomic versioned selector to make sure that
-      # if we lose the lock for a long period of time, we don't mess up
-      # with other people's updates. Also we make sure that the recovery
-      # mechanism is not racing with us.
-    when :destroy
-      use_id_selector(:use_atomic_version_selector => true)
-    end
-
-    # Perform the actual database query (single write or transaction commit).
-    # If successful, the result goes in @result, otherwise, @exception contains
-    # the thrown exception.
-    perform_db_operation_with_no_exceptions(&db_operation)
-
-    # We take a timestamp right after the write is performed because latency
-    # measurements are performed on the subscriber.
-    record_timestamp
-
-    if operation == :update && !failed?
-      # The underlying driver should implement some sort of find and modify
-      # operation in the previous write query to avoid this extra read query.
-      # If reload_instance raise an exception, we let it bubble up,
-      # and we'll trigger the recovery mechanism.
-      use_id_selector
-      reload_instance
-    end
-
-    unless @locks.first.still_locked?
-      # We lost the lock, let the recovery mechanism do its thing.
-      # This is a code optimization to avoid checking if the db operation
-      # succeeded or not because of the db operation race during recovery.
-      raise Promiscuous::Error::LostLock.new(@locks.first.key)
+    if should_instrument_query?
+      raise Promiscuous::Error::MissingContext if !current_context && write?
+      execute_instrumented(query)
+    else
+      query.call_and_remember_result(:non_instrumented)
     end
 
-    generate_payload_and_clear_operations
-
-    # As soon as we unlock the locks, the rescuer will not be able to assume
-    # that the database instance is still pristine, and so we need to stash the
-    # payload in redis. If redis dies, we don't care because it can be
-    # reconstructed. Subscribers can see "compressed" updates.
-    publish_payload_in_redis
-
-    # TODO Performance: merge these 3 redis operations to speed things up.
-    unlock_write_dependencies
-
-    # If we die from this point on, a recovery worker can republish our payload
-    # since we queued it in Redis.
-
-    # We don't care if we lost the lock and got recovered, subscribers are
-    # immune to duplicate messages.
-    publish_payload_in_rabbitmq_async
+    query.result
   end
 
-  # --- the following methods can be overridden by the driver  --- #
-
-  def execute_persistent(&db_operation)
-    return nil unless lock_instance_for_execute_persistent
-    execute_persistent_locked(&db_operation)
+  def query_dependencies
+    # Returns the list of dependencies that are involved in the database query.
+    # For an atomic write operation, the first one returned must be the one
+    # corresponding to the primary key.
+    raise
   end
 
-  def execute_non_persistent(&db_operation)
-    # We are getting here in the following cases:
-    # * read: we fetch the instance. It's the driver's job to cache the
-    #       raw instance and return it during db_operation.
-    # * multi read: nothing to do, we'll keep our current selector, sadly
-    # * write in a transaction: TODO
-
-    if single?
-      # If the query misses, we don't bother
-      return nil unless reload_instance
-      use_id_selector
-    end
-
-    # We don't do any reload_instance_dependencies at this point (and thus we
-    # won't raise an exception on a multi read that we cannot track).
-    # We'll wait until the commit, and hopefully with tainting, we'll be able to
-    # tell if we should depend the multi read operation in question.
-    perform_db_operation_with_no_exceptions(&db_operation)
-    # If the db_operation raises, we don't consider this failed operation when
-    # committing the next persistent write by omitting the operation in the
-    # context.
-    current_context.add_operation(self) unless failed?
-  end
-
-  def execute(&db_operation)
-    # execute returns the result of the db_operation to perform
-    db_operation ||= proc {}
-    return db_operation.call if Promiscuous.disabled
-
-    unless current_context
-      raise Promiscuous::Error::MissingContext if write?
-      return db_operation.call # Don't care for a read
-    end
-
-    self.persists? ? execute_persistent(&db_operation) :
-                     execute_non_persistent(&db_operation)
-
-    @exception ? (raise @exception) : @result
+  def execute_instrumented(db_operation)
+    # Implemented by subclasses
+    raise
   end
 
-  def fetch_instance
-    # This method is overridden to use the original query selector.
-    # Should return nil if the instance is not found.
-    @instance
+  def operation_payloads
+    # subclass can use payloads_for to generate the payload
+    raise
   end
 
-  def serialize_document_for_create_recovery
-    # Overridden to be able to redo the create during recovery.
-    nil
+  def recovery_payload
+    # Overridden to be able to recover the operation
+    []
   end
 
-  def self.recover_operation(model, instance_id, document)
-    # Overriden to reconstruct the operation. If the database is read, only the
-    # primary must be used.
-    new(:instance => model.new { |instance| instance.id = instance_id })
+  def self.recover_operation(*recovery_payload)
+    # Overridden to reconstruct the operation.
   end
 
   def recover_db_operation
-    # Overriden to reexecute the db operation during recovery (or make sure that
+    # Overridden to reexecute the db operation during recovery (or make sure that
     # it will never succeed).
   end
 
-  def use_id_selector(options={})
-    # Overridden to use the {:id => @instance.id} selector.
-    # if use_atomic_version_selector is passed, the driver must
-    # add the VERSION_FIELD selector if present in original instance.
-  end
-
-  def use_versioned_selector
-    # Overridden to use the {VERSION_FIELD => @instance[VERSION_FIELD]} selector.
-  end
-
-  def stash_version_in_write_query
-    # Overridden to update the query to set 'instance.VERSION_FIELD = @instance_version'
+  def trace_operation
+    if ENV['TRACE']
+      msg = self.explain_operation(70)
+      current_context.trace(msg, :color => self.read? ? '0;32' : '1;31')
+    end
   end
 
-  def going_to_execute_db_operation
-    # Test hook
+  def explain_operation(max_width)
+    "Unknown database operation"
   end
 end