feat(wallet): Improve how we're refreshing wallet's ongoing balance

app/jobs/clock/refresh_wallets_ongoing_balance_job.rb

@@ -11,7 +11,7 @@ class RefreshWalletsOngoingBalanceJob < ApplicationJob
     def perform
       return unless License.premium?
 
-      Wallet.active.find_each do |wallet|
+      Wallet.active.ready_to_be_refreshed.find_each do |wallet|
         Wallets::RefreshOngoingBalanceJob.perform_later(wallet)
       end
     end

app/jobs/wallets/refresh_ongoing_balance_job.rb

@@ -7,6 +7,8 @@ class RefreshOngoingBalanceJob < ApplicationJob
     unique :until_executed, on_conflict: :log
 
     def perform(wallet)
+      return unless wallet.ready_to_be_refreshed?
+
       Wallets::Balance::RefreshOngoingService.call(wallet:)
     end
   end

app/models/wallet.rb

@@ -28,6 +28,7 @@ def mark_as_terminated!(timestamp = Time.zone.now)
   end
 
   scope :expired, -> { where('wallets.expiration_at::timestamp(0) <= ?', Time.current) }
+  scope :ready_to_be_refreshed, -> { where(ready_to_be_refreshed: true) }
 
   def currency=(currency)
     self.balance_currency = currency
@@ -62,6 +63,7 @@ def currency
 #  ongoing_balance_cents               :bigint           default(0), not null
 #  ongoing_usage_balance_cents         :bigint           default(0), not null
 #  rate_amount                         :decimal(30, 5)   default(0.0), not null
+#  ready_to_be_refreshed               :boolean          default(FALSE), not null
 #  status                              :integer          not null
 #  terminated_at                       :datetime
 #  created_at                          :datetime         not null
@@ -70,7 +72,8 @@ def currency
 #
 # Indexes
 #
-#  index_wallets_on_customer_id  (customer_id)
+#  index_wallets_on_customer_id            (customer_id)
+#  index_wallets_on_ready_to_be_refreshed  (ready_to_be_refreshed) WHERE ready_to_be_refreshed
 #
 # Foreign Keys
 #

app/services/events/post_process_service.rb

@@ -25,6 +25,7 @@ def call
 
       expire_cached_charges(subscriptions)
       flag_refresh_from_subscription
+      flag_wallets_for_refresh
 
       handle_pay_in_advance
 
@@ -95,6 +96,13 @@ def flag_refresh_from_subscription
       subscriptions.select(&:active?).each { |s| LifetimeUsages::FlagRefreshFromSubscriptionService.new(subscription: s).call }
     end
 
+    def flag_wallets_for_refresh
+      return unless customer
+      return unless customer.wallets.active.any?
+
+      customer.wallets.active.update_all(ready_to_be_refreshed: true) # rubocop:disable Rails/SkipsModelValidations
+    end
+
     def handle_pay_in_advance
       return unless billable_metric
       return unless charges.any?

app/services/wallets/balance/update_ongoing_service.rb

@@ -34,7 +34,8 @@ def compute_update_params
           ongoing_usage_balance_cents: total_usage_amount_cents,
           credits_ongoing_usage_balance:,
           ongoing_balance_cents:,
-          credits_ongoing_balance:
+          credits_ongoing_balance:,
+          ready_to_be_refreshed: false
         }
 
         if !wallet.depleted_ongoing_balance? && ongoing_balance_cents <= 0

db/migrate/20241014093451_add_ready_to_be_refreshed_to_wallets.rb

@@ -0,0 +1,10 @@
+# frozen_string_literal: true
+
+class AddReadyToBeRefreshedToWallets < ActiveRecord::Migration[7.1]
+  disable_ddl_transaction!
+
+  def change
+    add_column :wallets, :ready_to_be_refreshed, :boolean, default: false, null: false
+    add_index :wallets, :ready_to_be_refreshed, where: "ready_to_be_refreshed", algorithm: :concurrently
+  end
+end

spec/jobs/clock/refresh_wallets_ongoing_balance_job_spec.rb

@@ -8,7 +8,7 @@
   describe '.perform' do
     let(:organization) { create(:organization) }
     let(:customer) { create(:customer, organization:) }
-    let(:wallet) { create(:wallet, customer:) }
+    let(:wallet) { create(:wallet, customer:, ready_to_be_refreshed: true) }
 
     before do
       wallet
@@ -30,8 +30,8 @@
         expect(Wallets::RefreshOngoingBalanceJob).to have_been_enqueued.with(wallet)
       end
 
-      context 'when not active' do
-        let(:wallet) { create(:wallet, :terminated) }
+      context 'when not ready to be refreshed' do
+        let(:wallet) { create(:wallet, customer:, ready_to_be_refreshed: false) }
 
         it 'does not call the refresh service' do
           described_class.perform_now

spec/jobs/wallets/refresh_ongoing_balance_job_spec.rb

@@ -3,19 +3,29 @@
 require 'rails_helper'
 
 RSpec.describe Wallets::RefreshOngoingBalanceJob, type: :job do
-  let(:wallet) { create(:wallet) }
+  let(:wallet) { create(:wallet, ready_to_be_refreshed: true) }
   let(:result) { BaseService::Result.new }
 
   let(:refresh_service) do
     instance_double(Wallets::Balance::RefreshOngoingService)
   end
 
-  it 'delegates to the RefreshOngoingBalance service' do
+  before do
     allow(Wallets::Balance::RefreshOngoingService).to receive(:new).with(wallet:).and_return(refresh_service)
     allow(refresh_service).to receive(:call).and_return(result)
+  end
 
-    described_class.perform_now(wallet)
+  context 'when wallet is ready to be refreshed' do
+    let(:wallet) { create(:wallet, ready_to_be_refreshed: false) }
 
+    it 'does not call the RefreshOngoingBalance service' do
+      described_class.perform_now(wallet)
+      expect(Wallets::Balance::RefreshOngoingService).not_to have_received(:new)
+    end
+  end
+
+  it 'delegates to the RefreshOngoingBalance service' do
+    described_class.perform_now(wallet)
     expect(Wallets::Balance::RefreshOngoingService).to have_received(:new)
     expect(refresh_service).to have_received(:call)
   end

spec/services/events/post_process_service_spec.rb

@@ -47,6 +47,12 @@
       expect(subscription.lifetime_usage.recalculate_current_usage).to be(true)
     end
 
+    it 'flags wallets for refresh' do
+      wallet = create(:wallet, customer:)
+
+      expect { process_service.call }.to change { wallet.reload.ready_to_be_refreshed }.from(false).to(true)
+    end
+
     context 'when event matches an pay_in_advance charge' do
       let(:charge) { create(:standard_charge, :pay_in_advance, plan:, billable_metric:, invoiceable: false) }
       let(:billable_metric) { create(:billable_metric, organization:, aggregation_type: 'sum_agg', field_name: 'item_id') }

spec/services/wallets/balance/update_ongoing_service_spec.rb

@@ -13,7 +13,8 @@
       ongoing_usage_balance_cents: 200,
       credits_balance: 10.0,
       credits_ongoing_balance: 8.0,
-      credits_ongoing_usage_balance: 2.0
+      credits_ongoing_usage_balance: 2.0,
+      ready_to_be_refreshed: true
     )
   end
   let(:total_usage_amount_cents) { 450 }
@@ -28,6 +29,7 @@
         .and change(wallet, :credits_ongoing_usage_balance).from(2.0).to(4.5)
         .and change(wallet, :ongoing_balance_cents).from(800).to(550)
         .and change(wallet, :credits_ongoing_balance).from(8.0).to(5.5)
+        .and change(wallet, :ready_to_be_refreshed).from(true).to(false)
 
       expect(wallet).not_to be_depleted_ongoing_balance
     end