diff --git a/enterprise/app/models/enterprise/account/plan_usage_and_limits.rb b/enterprise/app/models/enterprise/account/plan_usage_and_limits.rb index 705097030..0937c7b82 100644 --- a/enterprise/app/models/enterprise/account/plan_usage_and_limits.rb +++ b/enterprise/app/models/enterprise/account/plan_usage_and_limits.rb @@ -16,20 +16,15 @@ module Enterprise::Account::PlanUsageAndLimits # rubocop:disable Metrics/ModuleL 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_custom_attribute(CAPTAIN_RESPONSES_USAGE) end def reset_response_usage - custom_attributes[CAPTAIN_RESPONSES_USAGE] = 0 - save + update_custom_attribute(CAPTAIN_RESPONSES_USAGE, 0) end def update_document_usage - # this will ensure that the document count is always accurate - custom_attributes[CAPTAIN_DOCUMENTS_USAGE] = captain_documents.count - save + update_custom_attribute(CAPTAIN_DOCUMENTS_USAGE, captain_documents.count) end def email_transcript_enabled? @@ -130,6 +125,27 @@ module Enterprise::Account::PlanUsageAndLimits # rubocop:disable Metrics/ModuleL ChatwootApp.max_limit end + # Atomic jsonb_set to avoid clobbering concurrent writes to other custom_attributes keys. + # Goes through Account relation (rather than raw connection) so shard routing is respected. + # rubocop:disable Rails/SkipsModelValidations + def update_custom_attribute(key, value) + Account.where(id: id).update_all([ + "custom_attributes = jsonb_set(COALESCE(custom_attributes, '{}'), ARRAY[:key], :value::jsonb)", + { key: key, value: value.to_json } + ]) + custom_attributes[key] = value + end + + def increment_custom_attribute(key) + Account.where(id: id).update_all([ + "custom_attributes = jsonb_set(COALESCE(custom_attributes, '{}'), ARRAY[:key], " \ + '(COALESCE((custom_attributes ->> :key)::int, 0) + 1)::text::jsonb)', + { key: key } + ]) + custom_attributes[key] = custom_attributes[key].to_i + 1 + end + # rubocop:enable Rails/SkipsModelValidations + def validate_limit_keys errors.add(:limits, ': Invalid data') unless self[:limits].is_a? Hash self[:limits] = {} if self[:limits].blank? diff --git a/spec/enterprise/services/enterprise/billing/handle_stripe_event_service_spec.rb b/spec/enterprise/services/enterprise/billing/handle_stripe_event_service_spec.rb index 0f8ce1494..f9b550ef8 100644 --- a/spec/enterprise/services/enterprise/billing/handle_stripe_event_service_spec.rb +++ b/spec/enterprise/services/enterprise/billing/handle_stripe_event_service_spec.rb @@ -83,6 +83,37 @@ describe Enterprise::Billing::HandleStripeEventService do end end + describe 'subscription quantity update' do + before do + allow(subscription).to receive(:[]).with('plan') + .and_return({ 'id' => 'price_startups', 'product' => 'plan_id_startups', 'name' => 'Startups' }) + end + + it 'updates subscribed_quantity' do + allow(subscription).to receive(:[]).with('quantity').and_return(6) + + stripe_event_service.new.perform(event: event) + + expect(account.reload.custom_attributes['subscribed_quantity']).to eq(6) + end + + it 'persists quantity even when increment_response_usage runs concurrently' do + allow(subscription).to receive(:[]).with('quantity').and_return(6) + account.update!(custom_attributes: account.custom_attributes.merge('captain_responses_usage' => 100)) + + # Simulate: webhook updates quantity, then a concurrent increment_response_usage writes usage + stripe_event_service.new.perform(event: event) + account.reload + + # Simulate concurrent increment_response_usage (atomic jsonb_set, not full hash overwrite) + account.increment_response_usage + + # Quantity must survive the concurrent usage update + expect(account.reload.custom_attributes['subscribed_quantity']).to eq(6) + expect(account.reload.custom_attributes['captain_responses_usage']).to eq(101) + end + end + describe 'subscription deletion handling' do it 'calls CreateStripeCustomerService on subscription deletion' do allow(event).to receive(:type).and_return('customer.subscription.deleted')