Preserve rollup invalidations across concurrent refreshes

Co-authored-by: Amp <amp@ampcode.com>

Author: skyfallwastaken

ARCHITECTURE.md

@@ -162,9 +162,13 @@ disposable.
 [`DashboardRollupRefreshService`](app/services/dashboard_rollup_refresh_service.rb)
 rebuilds totals, dimensions, weekly projects, project details, filter options,
 activity graph and today's stats from the user's non-archived heartbeats. It
-atomically replaces all of one user's rows in a transaction. The refresh job
-marks the user dirty before enqueue, coalesces scheduling with a cache key, and
-uses a per-user GoodJob concurrency limit. Heartbeat commits, soft-delete/
+reads and replaces one user's rows in a repeatable-read transaction (callers
+with an existing transaction own its isolation level). Invalidations increment
+the user's durable `dashboard_rollup_generation`; the total row records the
+generation it read. A newer invalidation cannot be acknowledged by an older
+refresh. The job coalesces enqueueing with a disposable cache key and schedules
+a follow-up if generations still differ after completion. GoodJob serialises
+execution per user while allowing a pending follow-up. Heartbeat commits, soft-delete/
 restore, timezone changes, and project archive changes schedule refreshes.
 
 [`ProfileStatsService`](app/services/profile_stats_service.rb) is a thin

app/jobs/dashboard_rollup_refresh_job.rb

@@ -4,14 +4,20 @@ class DashboardRollupRefreshJob < ApplicationJob
   include GoodJob::ActiveJobExtensions::Concurrency
 
   good_job_control_concurrency_with(
-    total_limit: 1, key: -> { "dashboard_rollup_refresh_job_#{arguments.first}" }
+    perform_limit: 1, key: -> { "dashboard_rollup_refresh_job_#{arguments.first}" }
   )
+  retry_on ActiveRecord::SerializationFailure, ActiveRecord::Deadlocked,
+    GoodJob::ActiveJobExtensions::Concurrency::ConcurrencyExceededError, wait: 5.seconds, attempts: 10
 
   DEFAULT_WAIT = 2.minutes
   ENQUEUE_CACHE_KEY_PREFIX = "dashboard_rollup_refresh_enqueued".freeze
 
   def self.schedule_for(user_id, wait: DEFAULT_WAIT)
     DashboardRollup.mark_dirty(user_id)
+    enqueue_for(user_id, wait:)
+  end
+
+  def self.enqueue_for(user_id, wait: DEFAULT_WAIT)
     return unless Rails.cache.write(enqueue_cache_key(user_id), true, expires_in: wait + 1.minute, unless_exist: true)
     set(wait: wait).perform_later(user_id)
   end
@@ -22,7 +28,9 @@ def perform(user_id)
     user = User.find_by(id: user_id)
     return unless user
     DashboardRollupRefreshService.new(user:).call
+    refreshed = true
   ensure
     Rails.cache.delete(self.class.enqueue_cache_key(user_id))
+    self.class.enqueue_for(user_id) if refreshed && DashboardRollup.dirty?(user_id)
   end
 end

app/models/dashboard_rollup.rb

@@ -6,7 +6,6 @@ class DashboardRollup < ApplicationRecord
   TODAY_STATS_DIMENSION = "today_stats".freeze
   FILTER_OPTIONS_DIMENSION = "filter_options".freeze
   CODING_RHYTHM_DIMENSION = "coding_rhythm".freeze
-  DIRTY_CACHE_KEY_PREFIX = "dashboard_rollup_dirty".freeze
 
   belongs_to :user
 
@@ -20,8 +19,17 @@ class DashboardRollup < ApplicationRecord
   def total_dimension? = dimension == TOTAL_DIMENSION
   def bucket = bucket_value_present ? bucket_value : nil
 
-  def self.dirty_cache_key(user_id) = "#{DIRTY_CACHE_KEY_PREFIX}_#{user_id}"
-  def self.mark_dirty(user_id) = Rails.cache.write(dirty_cache_key(user_id), true, expires_in: 1.day, unless_exist: true)
-  def self.clear_dirty(user_id) = Rails.cache.delete(dirty_cache_key(user_id))
-  def self.dirty?(user_id) = Rails.cache.exist?(dirty_cache_key(user_id))
+  def self.generation(user_id) = User.where(id: user_id).pick(:dashboard_rollup_generation)
+
+  def self.mark_dirty(user_id)
+    User.where(id: user_id).update_all("dashboard_rollup_generation = dashboard_rollup_generation + 1")
+  end
+
+  def self.dirty?(user_id)
+    current_generation = generation(user_id)
+    return false unless current_generation
+
+    payload = find_by(user_id: user_id, dimension: TOTAL_DIMENSION)&.payload
+    payload&.fetch("source_generation", nil) != current_generation
+  end
 end

app/services/dashboard_rollup_refresh_service.rb

@@ -4,14 +4,27 @@ class DashboardRollupRefreshService < ApplicationService
 
   def initialize(user:)
     @user = user
-    @scope = user.heartbeats_excluding_archived_projects
   end
 
   def call
+    # When called inside a transaction, its owner also owns the isolation level.
+    isolation = :repeatable_read unless DashboardRollup.connection.transaction_open?
+    DashboardRollup.transaction(isolation:) do
+      @scope = @user.heartbeats_excluding_archived_projects
+      records = build_records(DashboardRollup.generation(@user.id))
+      DashboardRollup.where(user_id: @user.id).delete_all
+      DashboardRollup.insert_all!(records)
+    end
+  end
+
+  private
+
+  def build_records(generation)
     now = Time.current
     records = [
       build_record(dimension: DashboardRollup::TOTAL_DIMENSION, bucket: nil,
                    total_seconds: @scope.duration_seconds, now:,
+                   payload: { source_generation: generation },
                    source_heartbeats_count: @scope.count,
                    source_max_heartbeat_time: @scope.maximum(:time)),
       build_record(dimension: DashboardRollup::FILTER_OPTIONS_DIMENSION, bucket: nil,
@@ -52,15 +65,9 @@ def call
       )
     end
 
-    DashboardRollup.transaction do
-      DashboardRollup.where(user_id: @user.id).delete_all
-      DashboardRollup.insert_all!(records)
-    end
-    DashboardRollup.clear_dirty(@user.id)
+    records
   end
 
-  private
-
   def build_record(
     dimension:,
     bucket:,

app/services/dashboard_stats.rb

@@ -266,7 +266,7 @@ def aggregate_rollup_stale?(total_row)
 
   def schedule_rollup_refresh(wait:)
     return if @rollup_refresh_scheduled
-    DashboardRollupRefreshJob.schedule_for(user.id, wait: wait)
+    DashboardRollupRefreshJob.enqueue_for(user.id, wait: wait)
     @rollup_refresh_scheduled = true
   end
 

db/migrate/20260906124526_add_dashboard_rollup_generation_to_users.rb

@@ -0,0 +1,5 @@
+class AddDashboardRollupGenerationToUsers < ActiveRecord::Migration[8.1]
+  def change
+    add_column :users, :dashboard_rollup_generation, :bigint, null: false, default: 0
+  end
+end

db/schema.rb

@@ -0,0 +1,57 @@
+require "test_helper"
+
+class DashboardRollupConcurrencyTest < ActiveSupport::TestCase
+  self.use_transactional_tests = false
+  include ActiveJob::TestHelper
+
+  test "a correction during refresh leaves a coherent snapshot and a pending refresh" do
+    original_cache = Rails.cache
+    original_adapter = ActiveJob::Base.queue_adapter
+    Rails.cache = ActiveSupport::Cache::MemoryStore.new
+    ActiveJob::Base.queue_adapter = :test
+    user = create(:user)
+    create(:heartbeat, user: user, project: "api", language: "Ruby", time: 1_700_000_000.0)
+    heartbeat = create(:heartbeat, user: user, project: "api", language: "Ruby", time: 1_700_000_060.0)
+    DashboardRollupRefreshService.new(user: user).call
+    clear_enqueued_jobs
+    Rails.cache.clear
+    DashboardRollupRefreshJob.schedule_for(user.id)
+    main_thread = Thread.current
+    corrected = false
+    subscriber = lambda do |*args|
+      payload = args.last
+      if Thread.current == main_thread && !corrected && payload[:sql].include?('SELECT COUNT(*) FROM "heartbeats"')
+        corrected = true
+        Thread.new do
+          ActiveRecord::Base.connection_pool.with_connection do
+            Heartbeat.find(heartbeat.id).update!(language: "Python")
+          end
+        end.value
+      end
+    end
+
+    ActiveSupport::Notifications.subscribed(subscriber, "sql.active_record") do
+      DashboardRollupRefreshJob.perform_now(user.id)
+    end
+
+    assert corrected, "the concurrent correction must run during the aggregate reads"
+    assert DashboardRollup.dirty?(user.id), "the newer invalidation must survive refresh completion"
+    buckets = DashboardRollup.where(user: user, dimension: "language").to_h { |row| [ row.bucket, row.total_seconds ] }
+    assert_equal({ "Ruby" => 60 }, buckets)
+    assert_enqueued_jobs 2, only: DashboardRollupRefreshJob
+
+    DashboardRollupRefreshJob.perform_now(user.id)
+    assert_not DashboardRollup.dirty?(user.id)
+    assert_equal 60, DashboardRollup.find_by!(user: user, dimension: "language", bucket_value: "Python").total_seconds
+    DashboardRollup.mark_dirty(user.id)
+    Rails.cache.clear
+    assert DashboardRollup.dirty?(user.id), "invalidations must survive cache loss"
+  ensure
+    DashboardRollup.where(user: user).delete_all if user
+    Heartbeat.with_deleted.where(user: user).delete_all if user
+    user&.destroy!
+    clear_enqueued_jobs
+    Rails.cache = original_cache
+    ActiveJob::Base.queue_adapter = original_adapter
+  end
+end

test/services/dashboard_rollup_concurrency_test.rb

@@ -359,7 +359,7 @@ def stats.grouped_durations_snapshot(_scope) = raise("expected rollup-backed das
         create_heartbeat(user, project: "beta", language: "javascript", editor: "zed", operating_system: "linux", category: "coding")
       end
 
-      DashboardRollup.clear_dirty(user.id)
+      total_row.update!(payload: { source_generation: DashboardRollup.generation(user.id) })
       Rails.cache.delete(DashboardRollupRefreshJob.enqueue_cache_key(user.id))
 
       stats = build_stats(user)