1# frozen_string_literal: true
2
3module Gitlab
4  module Database
5    module BackgroundMigration
6      class BatchedMigrationWrapper
7        extend Gitlab::Utils::StrongMemoize
8
9        # Wraps the execution of a batched_background_migration.
10        #
11        # Updates the job's tracking records with the status of the migration
12        # when starting and finishing execution, and optionally saves batch_metrics
13        # the migration provides, if any are given.
14        #
15        # The job's batch_metrics are serialized to JSON for storage.
16        def perform(batch_tracking_record)
17          start_tracking_execution(batch_tracking_record)
18
19          execute_batch(batch_tracking_record)
20
21          batch_tracking_record.status = :succeeded
22        rescue Exception # rubocop:disable Lint/RescueException
23          batch_tracking_record.status = :failed
24
25          raise
26        ensure
27          finish_tracking_execution(batch_tracking_record)
28          track_prometheus_metrics(batch_tracking_record)
29        end
30
31        private
32
33        def start_tracking_execution(tracking_record)
34          tracking_record.update!(attempts: tracking_record.attempts + 1, status: :running, started_at: Time.current, finished_at: nil, metrics: {})
35        end
36
37        def execute_batch(tracking_record)
38          job_instance = tracking_record.migration_job_class.new
39
40          job_instance.perform(
41            tracking_record.min_value,
42            tracking_record.max_value,
43            tracking_record.migration_table_name,
44            tracking_record.migration_column_name,
45            tracking_record.sub_batch_size,
46            tracking_record.pause_ms,
47            *tracking_record.migration_job_arguments)
48
49          if job_instance.respond_to?(:batch_metrics)
50            tracking_record.metrics = job_instance.batch_metrics
51          end
52        end
53
54        def finish_tracking_execution(tracking_record)
55          tracking_record.finished_at = Time.current
56          tracking_record.save!
57        end
58
59        def track_prometheus_metrics(tracking_record)
60          migration = tracking_record.batched_migration
61          base_labels = migration.prometheus_labels
62
63          metric_for(:gauge_batch_size).set(base_labels, tracking_record.batch_size)
64          metric_for(:gauge_sub_batch_size).set(base_labels, tracking_record.sub_batch_size)
65          metric_for(:gauge_interval).set(base_labels, tracking_record.batched_migration.interval)
66          metric_for(:gauge_job_duration).set(base_labels, (tracking_record.finished_at - tracking_record.started_at).to_i)
67          metric_for(:counter_updated_tuples).increment(base_labels, tracking_record.batch_size)
68          metric_for(:gauge_migrated_tuples).set(base_labels, tracking_record.batched_migration.migrated_tuple_count)
69          metric_for(:gauge_total_tuple_count).set(base_labels, tracking_record.batched_migration.total_tuple_count)
70          metric_for(:gauge_last_update_time).set(base_labels, Time.current.to_i)
71
72          if metrics = tracking_record.metrics
73            metrics['timings']&.each do |key, timings|
74              summary = metric_for(:histogram_timings)
75              labels = base_labels.merge(operation: key)
76
77              timings.each do |timing|
78                summary.observe(labels, timing)
79              end
80            end
81          end
82        end
83
84        def metric_for(name)
85          self.class.metrics[name]
86        end
87
88        def self.metrics
89          strong_memoize(:metrics) do
90            {
91              gauge_batch_size: Gitlab::Metrics.gauge(
92                :batched_migration_job_batch_size,
93                'Batch size for a batched migration job'
94              ),
95              gauge_sub_batch_size: Gitlab::Metrics.gauge(
96                :batched_migration_job_sub_batch_size,
97                'Sub-batch size for a batched migration job'
98              ),
99              gauge_interval: Gitlab::Metrics.gauge(
100                :batched_migration_job_interval_seconds,
101                'Interval for a batched migration job'
102              ),
103              gauge_job_duration: Gitlab::Metrics.gauge(
104                :batched_migration_job_duration_seconds,
105                'Duration for a batched migration job'
106              ),
107              counter_updated_tuples: Gitlab::Metrics.counter(
108                :batched_migration_job_updated_tuples_total,
109                'Number of tuples updated by batched migration job'
110              ),
111              gauge_migrated_tuples: Gitlab::Metrics.gauge(
112                :batched_migration_migrated_tuples_total,
113                'Total number of tuples migrated by a batched migration'
114              ),
115              histogram_timings: Gitlab::Metrics.histogram(
116                :batched_migration_job_query_duration_seconds,
117                'Query timings for a batched migration job',
118                {},
119                [0.1, 0.25, 0.5, 1, 5].freeze
120              ),
121              gauge_total_tuple_count: Gitlab::Metrics.gauge(
122                :batched_migration_total_tuple_count,
123                'Total tuple count the migration needs to touch'
124              ),
125              gauge_last_update_time: Gitlab::Metrics.gauge(
126                :batched_migration_last_update_time_seconds,
127                'Unix epoch time in seconds'
128              )
129            }
130          end
131        end
132      end
133    end
134  end
135end
136