CCOL-2039: Post process valid and invalid records from batch consumption

lib/deimos/active_record_consume/batch_consumption.rb

@@ -28,15 +28,16 @@ def consume_batch(payloads, metadata)
           zip(metadata[:keys]).
           map { |p, k| Deimos::Message.new(p, nil, key: k) }
 
-        @tags = %W(topic:#{metadata[:topic]})
+        Deimos.config.tracer.active_span.set_tag('topic', metadata[:topic])
 
-        Deimos.instrument('ar_consumer.consume_batch', @tags) do
+        Deimos.instrument('ar_consumer.consume_batch',
+                          Deimos.config.tracer.active_span.get_tag('topic')) do
           if @compacted || self.class.config[:no_keys]
             update_database(compact_messages(messages))
           else
             uncompacted_update(messages)
           end
-        end
+                          end
       end
 
     protected
@@ -152,20 +153,18 @@ def update_database(messages)
       def upsert_records(messages)
         record_list = build_records(messages)
         invalid = filter_records(record_list)
-        unless invalid.blank?
+        if invalid.any?
           ActiveSupport::Notifications.instrument('batch_consumption.invalid_records', {
                                                     records: invalid,
                                                     consumer: self.class
                                                   })
-
         end
         return if record_list.empty?
 
         key_col_proc = self.method(:key_columns).to_proc
         col_proc = self.method(:columns).to_proc
 
         updater = MassUpdater.new(@klass,
-                                  @tags,
                                   key_col_proc: key_col_proc,
                                   col_proc: col_proc,
                                   replace_associations: self.class.replace_associations,
@@ -221,9 +220,11 @@ def build_records(messages)
       # deleted records.
       # @return [void]
       def remove_records(messages)
-        clause = deleted_query(messages)
+        Deimos::Utils::DeadlockRetry.wrap(Deimos.config.tracer.active_span.get_tag('topic')) do
+          clause = deleted_query(messages)
 
-        clause.delete_all
+          clause.delete_all
+        end
       end
     end
   end

lib/deimos/active_record_consume/mass_updater.rb

@@ -19,10 +19,9 @@ def default_cols(klass)
       # @param key_col_proc [Proc<Class < ActiveRecord::Base>]
       # @param col_proc [Proc<Class < ActiveRecord::Base>]
       # @param replace_associations [Boolean]
-      def initialize(klass, tags, key_col_proc: nil, col_proc: nil,
+      def initialize(klass, key_col_proc: nil, col_proc: nil,
                      replace_associations: true, bulk_import_id_generator: nil)
         @klass = klass
-        @tags = tags
         @replace_associations = replace_associations
         @bulk_import_id_generator = bulk_import_id_generator
 
@@ -90,7 +89,7 @@ def mass_update(record_list)
         # The entire batch should be treated as one transaction so that if
         # any message fails, the whole thing is rolled back or retried
         # if there is deadlock
-        Deimos::Utils::DeadlockRetry.wrap(@tags) do
+        Deimos::Utils::DeadlockRetry.wrap(Deimos.config.tracer.active_span.get_tag('topic')) do
           save_records_to_database(record_list)
           import_associations(record_list) if record_list.associations.any?
         end

lib/deimos/test_helpers.rb

@@ -28,6 +28,7 @@ def unit_test!
           deimos_config.kafka.seed_brokers ||= ['test_broker']
           deimos_config.schema.backend = Deimos.schema_backend_class.mock_backend
           deimos_config.producers.backend = :test
+          deimos_config.tracer = Deimos::Tracing::Mock.new
         end
       end
 

lib/deimos/tracing/mock.rb

@@ -10,6 +10,7 @@ class Mock < Tracing::Provider
       def initialize(logger=nil)
         @logger = logger || Logger.new(STDOUT)
         @logger.info('MockTracingProvider initialized')
+        @active_span = MockSpan.new
       end
 
       # @param span_name [String]
@@ -32,12 +33,22 @@ def finish(span)
 
       # :nodoc:
       def active_span
-        nil
+        @active_span ||= MockSpan.new
       end
 
       # :nodoc:
       def set_tag(tag, value, span=nil)
-        nil
+        if span
+          span.set_tag(tag, value)
+        else
+          active_span.set_tag(tag, value)
+        end
+      end
+
+      # Get a tag from a span with the specified tag.
+      # @param tag [String]
+      def get_tag(tag)
+        @span.get_tag(tag)
       end
 
       # :nodoc:
@@ -47,5 +58,23 @@ def set_error(span, exception)
         @logger.info("Mock span '#{name}' set an error: #{exception}")
       end
     end
+
+    # Mock Span class
+    class MockSpan
+      # :nodoc:
+      def initialize
+        @span = {}
+      end
+
+      # :nodoc:
+      def set_tag(tag, value)
+        @span[tag] = value
+      end
+
+      # :nodoc:
+      def get_tag(tag)
+        @span[tag]
+      end
+    end
   end
 end

lib/deimos/tracing/provider.rb

@@ -42,6 +42,12 @@ def set_tag(tag, value, span=nil)
         raise NotImplementedError
       end
 
+      # Get a tag from a span with the specified tag.
+      # @param tag [String]
+      def get_tag(tag)
+        raise NotImplementedError
+      end
+
     end
   end
 end

spec/active_record_batch_consumer_spec.rb

@@ -206,6 +206,48 @@ def publish_batch(messages)
                 ]
               )
           end
+
+          it 'should handle deletes with deadlock retries' do
+            allow(Deimos::Utils::DeadlockRetry).to receive(:sleep)
+            allow(instance_double(ActiveRecord::Relation)).to receive(:delete_all).and_raise(
+              ActiveRecord::Deadlocked.new('Lock wait timeout exceeded')
+            ).twice.ordered
+
+            Widget.create!(id: 1, test_id: 'abc', some_int: 2)
+
+            publish_batch(
+              [
+                { key: 1,
+                  payload: nil },
+                { key: 1,
+                  payload: nil }
+              ]
+            )
+
+            expect(all_widgets).to be_empty
+          end
+
+          it 'should not delete after multiple deadlock retries' do
+            allow(Deimos::Utils::DeadlockRetry).to receive(:sleep)
+            allow(instance_double(ActiveRecord::Relation)).to receive(:delete_all).and_raise(
+              ActiveRecord::Deadlocked.new('Lock wait timeout exceeded')
+            ).exactly(3).times
+
+            Widget.create!(id: 1, test_id: 'abc', some_int: 2)
+
+            publish_batch(
+              [
+                { key: 1,
+                  payload: nil },
+                { key: 1,
+                  payload: nil }
+              ]
+            )
+
+            expect(Widget.count).to eq(0)
+
+          end
+
         end
       end
     end

spec/active_record_consume/mass_updater_spec.rb

@@ -69,8 +69,7 @@
 
       it 'should mass update the batch' do
         allow(SecureRandom).to receive(:uuid).and_return('1', '2')
-        results = described_class.new(Widget, %w(topic:my-topic),
-                                      bulk_import_id_generator: bulk_id_generator).mass_update(batch)
+        results = described_class.new(Widget, bulk_import_id_generator: bulk_id_generator).mass_update(batch)
         expect(results.count).to eq(2)
         expect(results.map(&:test_id)).to match(%w(id1 id2))
         expect(Widget.count).to eq(2)
@@ -80,5 +79,40 @@
         expect(Widget.last.detail).not_to be_nil
       end
 
+      context 'with deadlock retries' do
+        before(:each) do
+          allow(Deimos::Utils::DeadlockRetry).to receive(:sleep)
+        end
+
+        it 'should upsert rows after deadlocks' do
+          allow(Widget).to receive(:import!).and_raise(
+            ActiveRecord::Deadlocked.new('Lock wait timeout exceeded')
+          ).twice.ordered
+          allow(Widget).to receive(:import!).and_raise(
+            ActiveRecord::Deadlocked.new('Lock wait timeout exceeded')
+          ).once.and_call_original
+
+          results = described_class.new(Widget, bulk_import_id_generator: bulk_id_generator).mass_update(batch)
+          expect(results.count).to eq(2)
+          expect(results.map(&:test_id)).to match(%w(id1 id2))
+          expect(Widget.count).to eq(2)
+          expect(Detail.count).to eq(2)
+          expect(Widget.first.detail).not_to be_nil
+          expect(Widget.last.detail).not_to be_nil
+        end
+
+        it 'should not upsert after encountering multiple deadlocks' do
+          allow(Widget).to receive(:import!).and_raise(
+            ActiveRecord::Deadlocked.new('Lock wait timeout exceeded')
+          ).exactly(3).times
+          expect {
+            described_class.new(Widget, bulk_import_id_generator: bulk_id_generator).mass_update(batch)
+          }.to raise_error(ActiveRecord::Deadlocked)
+          expect(Widget.count).to eq(0)
+          expect(Detail.count).to eq(0)
+        end
+
+      end
+
     end
 end