move concerns to the enterprise concerns

enterprise/app/models/concerns/account_limit_validation.rb

@@ -0,0 +1,22 @@
+module Concerns::AccountLimitValidation
+  extend ActiveSupport::Concern
+
+  LIMIT_SCHEMA = {
+    'type' => 'object',
+    'properties' => {
+      'inboxes' => { 'type': 'number' },
+      'agents' => { 'type': 'number' },
+      'captain_responses' => { 'type': 'number' },
+      'captain_documents' => { 'type': 'number' }
+    },
+    'required' => [],
+    'additionalProperties' => false
+  }.freeze
+
+  def validate_limit_keys
+    errors.add(:limits, ': Invalid data') unless self[:limits].is_a? Hash
+    self[:limits] = {} if self[:limits].blank?
+
+    errors.add(:limits, ': Invalid data') unless JSONSchemer.schema(LIMIT_SCHEMA).valid?(self[:limits])
+  end
+end

enterprise/app/models/enterprise/account/plan_usage_and_limits.rb

@@ -1,4 +1,6 @@
 module Enterprise::Account::PlanUsageAndLimits
+  include Concerns::AccountLimitValidation
+
   CAPTAIN_RESPONSES = 'captain_responses'.freeze
   CAPTAIN_DOCUMENTS = 'captain_documents'.freeze
   CAPTAIN_RESPONSES_USAGE = 'captain_responses_usage'.freeze
@@ -16,9 +18,17 @@ module Enterprise::Account::PlanUsageAndLimits
   end
 
   def increment_response_usage
-    current_usage = custom_attributes[CAPTAIN_RESPONSES_USAGE].to_i || 0
-    custom_attributes[CAPTAIN_RESPONSES_USAGE] = current_usage + 1
-    save
+    increment_sql = <<~SQL.squish
+      custom_attributes = jsonb_set(
+        custom_attributes,
+        '{#{CAPTAIN_RESPONSES_USAGE}}',
+        to_jsonb(COALESCE((custom_attributes->>'#{CAPTAIN_RESPONSES_USAGE}')::int, 0) + 1),
+        true
+      )
+    SQL
+    # Skipping validations is safe: no validations on custom_attributes,
+    updated = self.class.where(id: id).update_all(increment_sql) # rubocop:disable Rails/SkipsModelValidations
+    reload if updated.positive?
   end
 
   def reset_response_usage
@@ -108,23 +118,4 @@ module Enterprise::Account::PlanUsageAndLimits
 
     ChatwootApp.max_limit
   end
-
-  def validate_limit_keys
-    errors.add(:limits, ': Invalid data') unless self[:limits].is_a? Hash
-    self[:limits] = {} if self[:limits].blank?
-
-    limit_schema = {
-      'type' => 'object',
-      'properties' => {
-        'inboxes' => { 'type': 'number' },
-        'agents' => { 'type': 'number' },
-        'captain_responses' => { 'type': 'number' },
-        'captain_documents' => { 'type': 'number' }
-      },
-      'required' => [],
-      'additionalProperties' => false
-    }
-
-    errors.add(:limits, ': Invalid data') unless JSONSchemer.schema(limit_schema).valid?(self[:limits])
-  end
 end

spec/enterprise/models/account_spec.rb

@@ -70,6 +70,7 @@ RSpec.describe Account, type: :model do
       ## Responses
       it 'incrementing responses updates usage_limits' do
         account.increment_response_usage
+        account.reload
 
         responses_limits = account.usage_limits[:captain][:responses]
 
@@ -78,6 +79,34 @@ RSpec.describe Account, type: :model do
         expect(responses_limits[:current_available]).to eq captain_limits[:startups][:responses] - 1
       end
 
+      it 'handles concurrent increments without losing updates' do
+        concurrent_calls = 50
+
+        ready = Concurrent::CountDownLatch.new(concurrent_calls)
+        start = Concurrent::CountDownLatch.new(1)
+
+        threads = Array.new(concurrent_calls) do
+          Thread.new do
+            thread_count = described_class.find(account.id)
+            ready.count_down
+            start.wait
+
+            thread_count.increment_response_usage
+          end
+        end
+        ready.wait
+
+        start.count_down
+        threads.each(&:join)
+
+        account.reload
+
+        expect(account.custom_attributes['captain_responses_usage']).to eq concurrent_calls
+        responses_limits = account.usage_limits[:captain][:responses]
+        expect(responses_limits[:consumed]).to eq concurrent_calls
+        expect(responses_limits[:current_available]).to eq captain_limits[:startups][:responses] - concurrent_calls
+      end
+
       it 'reseting responses limits updates usage_limits' do
         account.custom_attributes['captain_responses_usage'] = 30
         account.save!

spec/enterprise/services/enterprise/billing/handle_stripe_event_service_spec.rb

@@ -66,6 +66,7 @@ describe Enterprise::Billing::HandleStripeEventService do
     it 'resets captain usage on billing period renewal' do
       # Prime the account with some usage
       5.times { account.increment_response_usage }
+      account.reload
       expect(account.custom_attributes['captain_responses_usage']).to eq(5)
 
       # Setup for any plan