add support for client timestamps and serial consistency in batch statements

Author: Bulat Shakirzyanov

lib/cassandra.rb

@@ -122,6 +122,12 @@ module Cassandra
   #   address resolver to use. Must be one of `:none` or
   #   `:ec2_multi_region`.
   #
+  # @option options [Boolean] :client_timestamps (false) whether the driver
+  #   should send timestamps for each execited statement. Enabling this setting
+  #   allows mitigating Cassandra cluster clock skew because the timestamp of
+  #   the client machine will be used. This does not help mitigate application
+  #   cluster clock skew.
+  #
   # @option options [Boolean] :synchronize_schema (true) whether the driver
   #   should automatically keep schema metadata synchronized. When enabled, the
   #   driver updates schema metadata after receiving schema change
@@ -217,7 +223,7 @@ def self.cluster_async(options = {})
         :connect_timeout, :futures_factory, :datacenter, :address_resolution,
         :address_resolution_policy, :idle_timeout, :heartbeat_interval, :timeout,
         :synchronize_schema, :schema_refresh_delay, :schema_refresh_timeout,
-        :shuffle_replicas
+        :shuffle_replicas, :client_timestamps
       ].include?(key)
     end
 
@@ -476,6 +482,10 @@ def self.cluster_async(options = {})
       options[:synchronize_schema] = !!options[:synchronize_schema]
     end
 
+    if options.has_key?(:client_timestamps)
+      options[:client_timestamps] = !!options[:client_timestamps]
+    end
+
     hosts = []
 
     Array(options.fetch(:hosts, '127.0.0.1')).each do |host|

lib/cassandra/cluster/client.rb

@@ -207,9 +207,11 @@ def keyspace_dropped(keyspace)
       def query(statement, options)
         return @futures.error(Errors::ClientError.new("Positional arguments are not supported by the current version of Apache Cassandra")) if !statement.params.empty? && @connection_options.protocol_version == 1
 
-        request = Protocol::QueryRequest.new(statement.cql, statement.params, statement.params_types, options.consistency, options.serial_consistency, options.page_size, options.paging_state, options.trace?, statement.params_names)
-        timeout = options.timeout
-        promise = @futures.promise
+        timestamp = nil
+        timestamp = Time.now if @connection_options.client_timestamps? && @connection_options.protocol_version > 2
+        request   = Protocol::QueryRequest.new(statement.cql, statement.params, statement.params_types, options.consistency, options.serial_consistency, options.page_size, options.paging_state, options.trace?, statement.params_names, timestamp)
+        timeout   = options.timeout
+        promise   = @futures.promise
 
         keyspace = @keyspace
         plan     = @load_balancing_policy.plan(keyspace, statement, options)
@@ -234,9 +236,11 @@ def prepare(cql, options)
       end
 
       def execute(statement, options)
+        timestamp       = nil
+        timestamp       = Time.now if @connection_options.client_timestamps? && @connection_options.protocol_version > 2
         timeout         = options.timeout
         result_metadata = statement.result_metadata
-        request         = Protocol::ExecuteRequest.new(nil, statement.params_types, statement.params, result_metadata.nil?, options.consistency, options.serial_consistency, options.page_size, options.paging_state, options.trace?)
+        request         = Protocol::ExecuteRequest.new(nil, statement.params_types, statement.params, result_metadata.nil?, options.consistency, options.serial_consistency, options.page_size, options.paging_state, options.trace?, timestamp)
         promise         = @futures.promise
 
         keyspace = @keyspace
@@ -250,12 +254,15 @@ def execute(statement, options)
       def batch(statement, options)
         return @futures.error(Errors::ClientError.new("Batch statements are not supported by the current version of Apache Cassandra")) if @connection_options.protocol_version < 2
 
-        timeout  = options.timeout
-        keyspace = @keyspace
-        plan     = @load_balancing_policy.plan(keyspace, statement, options)
-        promise  = @futures.promise
+        timestamp = nil
+        timestamp = Time.now if @connection_options.client_timestamps? && @connection_options.protocol_version > 2
+        timeout   = options.timeout
+        request   = Protocol::BatchRequest.new(BATCH_TYPES[statement.type], options.consistency, options.trace?, options.serial_consistency, timestamp)
+        keyspace  = @keyspace
+        plan      = @load_balancing_policy.plan(keyspace, statement, options)
+        promise   = @futures.promise
 
-        batch_by_plan(promise, keyspace, statement, options, plan, timeout)
+        batch_by_plan(promise, keyspace, statement, options, request, plan, timeout)
 
         promise.future
       end
@@ -509,7 +516,7 @@ def prepare_and_send_request_by_plan(host, connection, promise, keyspace, statem
         end
       end
 
-      def batch_by_plan(promise, keyspace, statement, options, plan, timeout, errors = nil, hosts = [])
+      def batch_by_plan(promise, keyspace, statement, options, request, plan, timeout, errors = nil, hosts = [])
         unless plan.has_next?
           promise.break(Errors::NoHostsAvailable.new(errors))
           return
@@ -522,7 +529,7 @@ def batch_by_plan(promise, keyspace, statement, options, plan, timeout, errors =
         unless pool
           errors ||= {}
           errors[host] = NOT_CONNECTED
-          return batch_by_plan(promise, keyspace, statement, options, plan, timeout, errors, hosts)
+          return batch_by_plan(promise, keyspace, statement, options, request, plan, timeout, errors, hosts)
         end
 
         connection = pool.random_connection
@@ -531,31 +538,31 @@ def batch_by_plan(promise, keyspace, statement, options, plan, timeout, errors =
           switch = switch_keyspace(connection, keyspace, timeout)
           switch.on_complete do |s|
             if s.resolved?
-              batch_and_send_request_by_plan(host, connection, promise, keyspace, statement, options, plan, timeout, errors, hosts)
+              batch_and_send_request_by_plan(host, connection, promise, keyspace, statement, request, options, plan, timeout, errors, hosts)
             else
               s.on_failure do |e|
                 case e
                 when Errors::HostError
                   errors ||= {}
                   errors[host] = e
-                  batch_by_plan(promise, keyspace, statement, options, plan, timeout, errors, hosts)
+                  batch_by_plan(promise, keyspace, statement, options, request, plan, timeout, errors, hosts)
                 else
                   promise.break(e)
                 end
               end
             end
           end
         else
-          batch_and_send_request_by_plan(host, connection, promise, keyspace, statement, options, plan, timeout, errors, hosts)
+          batch_and_send_request_by_plan(host, connection, promise, keyspace, statement, request, options, plan, timeout, errors, hosts)
         end
       rescue => e
         errors ||= {}
         errors[host] = e
-        batch_by_plan(promise, keyspace, statement, options, plan, timeout, errors, hosts)
+        batch_by_plan(promise, keyspace, statement, options, request, plan, timeout, errors, hosts)
       end
 
-      def batch_and_send_request_by_plan(host, connection, promise, keyspace, statement, options, plan, timeout, errors, hosts)
-        request    = Protocol::BatchRequest.new(BATCH_TYPES[statement.type], options.consistency, options.trace?)
+      def batch_and_send_request_by_plan(host, connection, promise, keyspace, statement, request, options, plan, timeout, errors, hosts)
+        request.clear
         unprepared = Hash.new {|hash, cql| hash[cql] = []}
 
         statement.statements.each do |statement|
@@ -598,7 +605,7 @@ def batch_and_send_request_by_plan(host, connection, promise, keyspace, statemen
                 when Errors::HostError
                   errors ||= {}
                   errors[host] = e
-                  batch_by_plan(promise, keyspace, statement, options, plan, timeout, errors, hosts)
+                  batch_by_plan(promise, keyspace, statement, options, request, plan, timeout, errors, hosts)
                 else
                   promise.break(e)
                 end
@@ -711,7 +718,7 @@ def do_send_request_by_plan(host, connection, promise, keyspace, statement, opti
                   when Protocol::ExecuteRequest
                     execute_by_plan(promise, keyspace, statement, options, request, plan, timeout, errors, hosts)
                   when Protocol::BatchRequest
-                    batch_by_plan(promise, keyspace, statement, options, plan, timeout, errors, hosts)
+                    batch_by_plan(promise, keyspace, statement, options, request, plan, timeout, errors, hosts)
                   end
                 else
                   promise.break(error)
@@ -774,7 +781,7 @@ def do_send_request_by_plan(host, connection, promise, keyspace, statement, opti
               when Protocol::ExecuteRequest
                 execute_by_plan(promise, keyspace, statement, options, request, plan, timeout, errors, hosts)
               when Protocol::BatchRequest
-                batch_by_plan(promise, keyspace, statement, options, plan, timeout, errors, hosts)
+                batch_by_plan(promise, keyspace, statement, options, request, plan, timeout, errors, hosts)
               else
                 promise.break(e)
               end

lib/cassandra/cluster/options.rb

@@ -25,7 +25,7 @@ class Options
                   :schema_refresh_delay, :schema_refresh_timeout
       attr_accessor :protocol_version
 
-      def initialize(protocol_version, credentials, auth_provider, compressor, port, connect_timeout, ssl, connections_per_local_node, connections_per_remote_node, heartbeat_interval, idle_timeout, synchronize_schema, schema_refresh_delay, schema_refresh_timeout)
+      def initialize(protocol_version, credentials, auth_provider, compressor, port, connect_timeout, ssl, connections_per_local_node, connections_per_remote_node, heartbeat_interval, idle_timeout, synchronize_schema, schema_refresh_delay, schema_refresh_timeout, client_timestamps)
         @protocol_version       = protocol_version
         @credentials            = credentials
         @auth_provider          = auth_provider
@@ -38,6 +38,7 @@ def initialize(protocol_version, credentials, auth_provider, compressor, port, c
         @synchronize_schema     = synchronize_schema
         @schema_refresh_delay   = schema_refresh_delay
         @schema_refresh_timeout = schema_refresh_timeout
+        @client_timestamps      = client_timestamps
 
         @connections_per_local_node  = connections_per_local_node
         @connections_per_remote_node = connections_per_remote_node
@@ -47,6 +48,10 @@ def synchronize_schema?
         @synchronize_schema
       end
 
+      def client_timestamps?
+        @client_timestamps
+      end
+
       def compression
         @compressor && @compressor.algorithm
       end

lib/cassandra/driver.rb

@@ -86,7 +86,8 @@ def self.let(name, &block)
         idle_timeout,
         synchronize_schema,
         schema_refresh_delay,
-        schema_refresh_timeout
+        schema_refresh_timeout,
+        client_timestamps
       )
     end
 
@@ -114,6 +115,7 @@ def self.let(name, &block)
     let(:schema_refresh_timeout)    { 10 }
     let(:thread_pool_size)          { 4 }
     let(:shuffle_replicas)          { true }
+    let(:client_timestamps)         { false }
 
     let(:connections_per_local_node)  { 2 }
     let(:connections_per_remote_node) { 1 }

lib/cassandra/protocol/cql_byte_buffer.rb

@@ -258,6 +258,10 @@ def append_string_map(map)
         self
       end
 
+      def append_timestamp(timestamp)
+        append_long(timestamp.tv_sec * 1000000 + timestamp.tv_usec)
+      end
+
       def append_long(n)
         top = n >> 32
         bottom = n & 0xffffffff

lib/cassandra/protocol/requests/batch_request.rb

@@ -23,14 +23,21 @@ class BatchRequest < Request
       UNLOGGED_TYPE = 1
       COUNTER_TYPE = 2
 
-      attr_reader :type
+      attr_reader :type, :timestamp
       attr_accessor :consistency, :retries
 
-      def initialize(type, consistency, trace=false)
+      def initialize(type, consistency, trace=false, serial_consistency = nil, timestamp = nil)
         super(0x0D, trace)
         @type  = type
         @parts = []
         @consistency = consistency
+        @serial_consistency = serial_consistency
+        @timestamp = timestamp
+      end
+
+      def clear
+        @parts.clear
+        nil
       end
 
       def add_query(cql, values, types)
@@ -54,9 +61,13 @@ def write(buffer, protocol_version, encoder)
         buffer.append_consistency(@consistency)
 
         if protocol_version > 2
-          flags = 0
+          flags  = 0
+          flags |= 0x10 if @serial_consistency
+          flags |= 0x20 if @timestamp
 
           buffer.append(flags.chr)
+          buffer.append_consistency(@serial_consistency) if @serial_consistency
+          buffer.append_timestamp(@timestamp) if @timestamp
         end
 
         buffer

lib/cassandra/protocol/requests/execute_request.rb

@@ -19,10 +19,10 @@
 module Cassandra
   module Protocol
     class ExecuteRequest < Request
-      attr_reader :metadata, :values, :request_metadata, :serial_consistency, :page_size, :paging_state
+      attr_reader :metadata, :values, :request_metadata, :serial_consistency, :page_size, :paging_state, :timestamp
       attr_accessor :consistency, :retries, :id
 
-      def initialize(id, metadata, values, request_metadata, consistency, serial_consistency=nil, page_size=nil, paging_state=nil, trace=false)
+      def initialize(id, metadata, values, request_metadata, consistency, serial_consistency=nil, page_size=nil, paging_state=nil, trace=false, timestamp = nil)
         raise ArgumentError, "Metadata for #{metadata.size} columns, but #{values.size} values given" if metadata.size != values.size
         raise ArgumentError, %(No such consistency: #{consistency.inspect}) if consistency.nil? || !CONSISTENCIES.include?(consistency)
         raise ArgumentError, %(No such consistency: #{serial_consistency.inspect}) unless serial_consistency.nil? || CONSISTENCIES.include?(serial_consistency)
@@ -36,6 +36,7 @@ def initialize(id, metadata, values, request_metadata, consistency, serial_consi
         @serial_consistency = serial_consistency
         @page_size = page_size
         @paging_state = paging_state
+        @timestamp = timestamp
       end
 
       def write(buffer, protocol_version, encoder)
@@ -48,13 +49,19 @@ def write(buffer, protocol_version, encoder)
           flags |= 0x04 if @page_size
           flags |= 0x08 if @paging_state
           flags |= 0x10 if @serial_consistency
+          if protocol_version > 2
+            flags |= 0x20 if @timestamp
+          end
           buffer.append(flags.chr)
           if @values.size > 0
             encoder.write_parameters(buffer, @values, @metadata)
           end
           buffer.append_int(@page_size) if @page_size
           buffer.append_bytes(@paging_state) if @paging_state
           buffer.append_consistency(@serial_consistency) if @serial_consistency
+          if protocol_version > 2
+            buffer.append_timestamp(@timestamp) if @timestamp
+          end
         else
           encoder.write_parameters(buffer, @values, @metadata)
           buffer.append_consistency(@consistency)

lib/cassandra/protocol/requests/query_request.rb

@@ -19,10 +19,10 @@
 module Cassandra
   module Protocol
     class QueryRequest < Request
-      attr_reader :cql, :values, :type_hints, :serial_consistency, :page_size, :paging_state
+      attr_reader :cql, :values, :type_hints, :serial_consistency, :page_size, :paging_state, :timestamp
       attr_accessor :consistency, :retries
 
-      def initialize(cql, values, type_hints, consistency, serial_consistency = nil, page_size = nil, paging_state = nil, trace = false, names = EMPTY_LIST)
+      def initialize(cql, values, type_hints, consistency, serial_consistency = nil, page_size = nil, paging_state = nil, trace = false, names = EMPTY_LIST, timestamp = nil)
         super(7, trace)
         @cql = cql
         @values = values
@@ -32,6 +32,7 @@ def initialize(cql, values, type_hints, consistency, serial_consistency = nil, p
         @page_size = page_size
         @paging_state = paging_state
         @names = names
+        @timestamp = timestamp
       end
 
       def write(buffer, protocol_version, encoder)
@@ -44,7 +45,10 @@ def write(buffer, protocol_version, encoder)
           flags |= 0x10 if @serial_consistency
           if @values && @values.size > 0
             flags |= 0x01
-            flags |= 0x40 unless @names.empty?
+            if protocol_version > 2
+              flags |= 0x20 if @timestamp
+              flags |= 0x40 unless @names.empty?
+            end
             buffer.append(flags.chr)
             encoder.write_parameters(buffer, @values, @type_hints, @names)
           else
@@ -53,6 +57,9 @@ def write(buffer, protocol_version, encoder)
           buffer.append_int(@page_size) if @page_size
           buffer.append_bytes(@paging_state) if @paging_state
           buffer.append_consistency(@serial_consistency) if @serial_consistency
+          if protocol_version > 2
+            buffer.append_timestamp(@timestamp) if @timestamp
+          end
         end
         buffer
       end

spec/cassandra/cluster/client_spec.rb

@@ -770,7 +770,8 @@ class Cluster
 
           client.connect.value
 
-          expect(Cassandra::Protocol::BatchRequest).to receive(:new).once.with(0, :one, false).and_return(batch_request)
+          expect(Cassandra::Protocol::BatchRequest).to receive(:new).once.with(0, :one, false, nil, nil).and_return(batch_request)
+          allow(batch_request).to receive(:clear)
           expect(batch_request).to receive(:add_query).once.with('INSERT INTO songs (id, title, album, artist, tags) VALUES (?, ?, ?, ?, ?)', [1, 2, 3, 4, 5], [Cassandra::Types.bigint, Cassandra::Types.bigint, Cassandra::Types.bigint, Cassandra::Types.bigint, Cassandra::Types.bigint])
           expect(batch_request).to receive(:retries=).once.with(0)
           client.batch(batch, Execution::Options.new(:consistency => :one, :trace => false)).get
@@ -810,7 +811,8 @@ class Cluster
 
           batch.add(statement, [Cassandra::Uuid.new(1), 'some title', 'some album', 'some artist', Set['cool', 'stuff']])
 
-          expect(Cassandra::Protocol::BatchRequest).to receive(:new).once.with(0, :one, false).and_return(batch_request)
+          expect(Cassandra::Protocol::BatchRequest).to receive(:new).once.with(0, :one, false, nil, nil).and_return(batch_request)
+          allow(batch_request).to receive(:clear)
           expect(batch_request).to receive(:add_prepared).once.with(123, [Cassandra::Uuid.new(1), 'some title', 'some album', 'some artist', Set['cool', 'stuff']], params_metadata.map(&:last))
           expect(batch_request).to receive(:retries=).once.with(0)
           client.batch(batch, Execution::Options.new(:consistency => :one, :trace => false)).get
@@ -852,7 +854,8 @@ class Cluster
 
           batch.add(statement, [Cassandra::Uuid.new(1), 'some title', 'some album', 'some artist', Set['cool', 'stuff']])
 
-          expect(Cassandra::Protocol::BatchRequest).to receive(:new).once.with(0, :one, false).and_return(batch_request)
+          expect(Cassandra::Protocol::BatchRequest).to receive(:new).once.with(0, :one, false, nil, nil).and_return(batch_request)
+          allow(batch_request).to receive(:clear)
           expect(batch_request).to receive(:add_prepared).once.with(123, [Cassandra::Uuid.new(1), 'some title', 'some album', 'some artist', Set['cool', 'stuff']], params_metadata.map(&:last))
           expect(batch_request).to receive(:retries=).once.with(0)