From 4229e4486f8b885c5c081cdd5250588f12124ebe Mon Sep 17 00:00:00 2001
From: LovroColic <lovro.colic.lc@gmail.com>
Date: Wed, 9 Oct 2024 14:32:05 +0200
Subject: [PATCH] fix (wallet-race-condition): Retry wallet update upon race
 condition (#2664)

## Context

Recently, feature for updating wallet balance upon successful payment
was introduced.

## Description

If multiple paid credits are issues, there is a chance that after
successful payment processing we will receive all events from payment
provider in the same request.

Multiple events will schedule multiple jobs for increasing wallet
balance and each job is updating the same resource so there is a high
chance of race condition.

This PR introduced optimistic locking approach with 1 additional retry
if multiple processes are updating the same resource.
---
 app/models/wallet.rb                          |  1 +
 .../wallets/balance/increase_service.rb       | 47 +++++++++++++------
 ...41008080209_add_lock_version_to_wallets.rb |  7 +++
 db/schema.rb                                  |  3 +-
 4 files changed, 42 insertions(+), 16 deletions(-)
 create mode 100644 db/migrate/20241008080209_add_lock_version_to_wallets.rb

diff --git a/app/models/wallet.rb b/app/models/wallet.rb
index d6b45008dfd..376d0a4ecc8 100644
--- a/app/models/wallet.rb
+++ b/app/models/wallet.rb
@@ -57,6 +57,7 @@ def currency
 #  invoice_requires_successful_payment :boolean          default(FALSE), not null
 #  last_balance_sync_at                :datetime
 #  last_consumed_credit_at             :datetime
+#  lock_version                        :integer          default(0), not null
 #  name                                :string
 #  ongoing_balance_cents               :bigint           default(0), not null
 #  ongoing_usage_balance_cents         :bigint           default(0), not null
diff --git a/app/services/wallets/balance/increase_service.rb b/app/services/wallets/balance/increase_service.rb
index 87a1fdb81b5..86d135ecc16 100644
--- a/app/services/wallets/balance/increase_service.rb
+++ b/app/services/wallets/balance/increase_service.rb
@@ -3,34 +3,51 @@
 module Wallets
   module Balance
     class IncreaseService < BaseService
+      MAX_RETRIES = 5
+
       def initialize(wallet:, credits_amount:, reset_consumed_credits: false)
         super
 
         @wallet = wallet
         @credits_amount = credits_amount
         @reset_consumed_credits = reset_consumed_credits
+        @retries = 0
       end
 
       def call
-        currency = wallet.balance.currency
-        amount_cents = wallet.rate_amount * credits_amount * currency.subunit_to_unit
-
-        update_params = {
-          balance_cents: wallet.balance_cents + amount_cents,
-          credits_balance: wallet.credits_balance + credits_amount,
-          last_balance_sync_at: Time.current
-        }
-
-        if reset_consumed_credits
-          update_params[:consumed_credits] = [0.0, wallet.consumed_credits - credits_amount].max
-          update_params[:consumed_amount_cents] = [0, wallet.consumed_amount_cents - amount_cents].max
+        ActiveRecord::Base.transaction do
+          currency = wallet.balance.currency
+          amount_cents = wallet.rate_amount * credits_amount * currency.subunit_to_unit
+
+          update_params = {
+            balance_cents: wallet.balance_cents + amount_cents,
+            credits_balance: wallet.credits_balance + credits_amount,
+            last_balance_sync_at: Time.current
+          }
+
+          if reset_consumed_credits
+            update_params[:consumed_credits] = [0.0, wallet.consumed_credits - credits_amount].max
+            update_params[:consumed_amount_cents] = [0, wallet.consumed_amount_cents - amount_cents].max
+          end
+
+          wallet.update!(update_params)
+          Wallets::Balance::RefreshOngoingService.call(wallet:)
         end
 
-        wallet.update!(update_params)
-        Wallets::Balance::RefreshOngoingService.call(wallet:)
-
         result.wallet = wallet
         result
+      rescue ActiveRecord::StaleObjectError
+        @retries += 1
+
+        if @retries <= MAX_RETRIES
+          sleep(0.5)
+
+          wallet.reload
+
+          retry
+        end
+
+        result.service_failure!(code: 'race_condition_error', message: '')
       end
 
       private
diff --git a/db/migrate/20241008080209_add_lock_version_to_wallets.rb b/db/migrate/20241008080209_add_lock_version_to_wallets.rb
new file mode 100644
index 00000000000..368fbc84983
--- /dev/null
+++ b/db/migrate/20241008080209_add_lock_version_to_wallets.rb
@@ -0,0 +1,7 @@
+# frozen_string_literal: true
+
+class AddLockVersionToWallets < ActiveRecord::Migration[7.1]
+  def change
+    add_column :wallets, :lock_version, :integer, default: 0, null: false
+  end
+end
diff --git a/db/schema.rb b/db/schema.rb
index d8baa37660f..75244a8931a 100644
--- a/db/schema.rb
+++ b/db/schema.rb
@@ -10,7 +10,7 @@
 #
 # It's strongly recommended that you check this file into your version control system.
 
-ActiveRecord::Schema[7.1].define(version: 2024_10_01_112117) do
+ActiveRecord::Schema[7.1].define(version: 2024_10_08_080209) do
   # These are extensions that must be enabled in order to support this database
   enable_extension "pgcrypto"
   enable_extension "plpgsql"
@@ -1160,6 +1160,7 @@
     t.decimal "credits_ongoing_usage_balance", precision: 30, scale: 5, default: "0.0", null: false
     t.boolean "depleted_ongoing_balance", default: false, null: false
     t.boolean "invoice_requires_successful_payment", default: false, null: false
+    t.integer "lock_version", default: 0, null: false
     t.index ["customer_id"], name: "index_wallets_on_customer_id"
   end