LG-15560: Enhance A/B tests to support configurable max participants

app/jobs/reports/ab_tests_report.rb

@@ -15,40 +15,39 @@ def initialize(report_date = nil, *args, **kwargs)
     def perform(report_date)
       @report_date = report_date
 
-      report_configs.each do |config|
-        tables_report(config).deliver_now
+      reported_ab_tests.each do |ab_test|
+        tables_report(ab_test).deliver_now
       end
     end
 
-    def tables_report(config)
-      experiment_name = config.experiment_name
+    def tables_report(ab_test)
+      experiment_name = ab_test.experiment_name
       subject = "A/B Tests Report - #{experiment_name} - #{report_date}"
-      reports = ab_tests_report(config).as_emailable_reports
+      report = ab_tests_report(ab_test)
 
       ReportMailer.tables_report(
-        email: config.email,
+        email: ab_test.report.email,
         subject:,
-        message: subject,
-        reports:,
+        message: [subject, report.participants_message].compact,
+        reports: report.as_emailable_reports,
         attachment_format: :csv,
       )
     end
 
-    def ab_tests_report(config)
+    def ab_tests_report(ab_test)
       Reporting::AbTestsReport.new(
-        queries: config.queries,
+        ab_test:,
         time_range: report_date.yesterday..report_date,
       )
     end
 
     private
 
-    def report_configs
+    def reported_ab_tests
       AbTests
         .all
         .values
         .select { |ab_test| ab_test.report&.email&.present? && ab_test.active? }
-        .map(&:report)
     end
   end
 end

app/mailers/report_mailer.rb

@@ -26,6 +26,7 @@ def warn_error(email:, error:, env: Rails.env)
   # @param [String] email
   # @param [String] subject
   # @param [String] env name of current deploy environment
+  # @param [String, String[]] message Message(s) to use as email prelude
   # @param [:csv,:xlsx] attachment_format
   # @param [Array<EmailableReport>] reports
   #   an array of tables (which are arrays of rows (arrays of strings))

app/models/ab_test_assignment.rb

@@ -0,0 +1,9 @@
+# frozen_string_literal: true
+
+class AbTestAssignment < ApplicationRecord
+  class << self
+    def bucket(**)
+      find_by(**)&.bucket&.to_sym
+    end
+  end
+end

app/services/analytics.rb

@@ -88,6 +88,7 @@ def ab_test_attributes(event)
         session:,
         user:,
         user_session:,
+        persisted_read_only: true,
       )
       if !bucket.blank?
         obj[test_id.downcase] = {

app/views/report_mailer/tables_report.html.erb

@@ -11,7 +11,9 @@
       ) %>
 </header>
 <hr class="height-05 margin-bottom-4 border-transparent bg-secondary">
-<%= @message %><br>
+<% [*@message].each do |message| %>
+  <p><%= message %></p>
+<% end %>
 <% @reports.each do |report| %>
   <% header, *rows = report.table %>
   <%# Allow nil title if a subtitle is set %>

db/primary_migrate/20250207144037_create_ab_test_assignments.rb

@@ -0,0 +1,10 @@
+class CreateAbTestAssignments < ActiveRecord::Migration[8.0]
+  def change
+    create_table :ab_test_assignments do |t|
+      t.string :experiment, null: false, comment: 'sensitive=false'
+      t.string :discriminator, null: false, comment: 'sensitive=false'
+      t.string :bucket, null: false, comment: 'sensitive=false'
+      t.index [:experiment, :discriminator], unique: true, using: :btree
+    end
+  end
+end

db/schema.rb

@@ -10,12 +10,19 @@
 #
 # It's strongly recommended that you check this file into your version control system.
 
-ActiveRecord::Schema[8.0].define(version: 2025_01_06_232958) do
+ActiveRecord::Schema[8.0].define(version: 2025_02_07_144037) do
   # These are extensions that must be enabled in order to support this database
   enable_extension "citext"
   enable_extension "pg_catalog.plpgsql"
   enable_extension "pg_stat_statements"
 
+  create_table "ab_test_assignments", force: :cascade do |t|
+    t.string "experiment", null: false, comment: "sensitive=false"
+    t.string "discriminator", null: false, comment: "sensitive=false"
+    t.string "bucket", null: false, comment: "sensitive=false"
+    t.index ["experiment", "discriminator"], name: "index_ab_test_assignments_on_experiment_and_discriminator", unique: true
+  end
+
   create_table "account_reset_requests", force: :cascade do |t|
     t.integer "user_id", null: false, comment: "sensitive=false"
     t.datetime "requested_at", precision: nil, comment: "sensitive=false"

lib/ab_test.rb

@@ -3,13 +3,22 @@
 class AbTest
   include ::NewRelic::Agent::MethodTracer
 
-  attr_reader :buckets, :experiment_name, :default_bucket, :should_log, :report
+  attr_reader :buckets,
+              :experiment_name,
+              :default_bucket,
+              :should_log,
+              :report,
+              :persist,
+              :max_participants
+
+  alias_method :experiment, :experiment_name
+  alias_method :persist?, :persist
 
   MAX_SHA = (16 ** 64) - 1
 
   ReportQueryConfig = Struct.new(:title, :query, :row_labels, keyword_init: true).freeze
 
-  ReportConfig = Struct.new(:experiment_name, :email, :queries, keyword_init: true) do
+  ReportConfig = Struct.new(:email, :queries, keyword_init: true) do
     def initialize(queries: [], **)
       super
       self.queries.map!(&ReportQueryConfig.method(:new))
@@ -19,6 +28,9 @@ def initialize(queries: [], **)
   # @param [Regexp,#include?,nil] should_log A list of analytics event names for which the A/B test
   #   bucket assignment should be logged, or a regular expression pattern which is tested against an
   #   analytics event name when an event is being logged.
+  # @param [Hash] report Report mailer configuration.
+  # @param [Boolean] persist Whether the test assignment should be persisted to the database.
+  # @param [Integer] max_participants The maximum number of participants allowed in the test.
   # @yieldparam [ActionDispatch::Request] request
   # @yieldparam [String,nil] service_provider Issuer string for the service provider associated with
   #                                           the current session.
@@ -30,14 +42,19 @@ def initialize(
     should_log: nil,
     default_bucket: :default,
     report: nil,
+    persist: false,
+    max_participants: Float::INFINITY,
     &discriminator
   )
     @buckets = buckets
     @discriminator = discriminator
     @experiment_name = experiment_name
     @default_bucket = default_bucket
     @should_log = should_log
-    @report = ReportConfig.new(experiment_name:, **report.to_h) if report
+    @report = ReportConfig.new(**report.to_h) if report
+    @persist = persist
+    raise 'max_participants requires persist to be true' if max_participants.finite? && !persist?
+    @max_participants = max_participants
     raise 'invalid bucket data structure' unless valid_bucket_data_structure?
     ensure_numeric_percentages
     raise 'bucket percentages exceed 100' unless within_100_percent?
@@ -50,7 +67,16 @@ def initialize(
   # @param [Hash] session
   # @param [User] user
   # @param [Hash] user_session
-  def bucket(request:, service_provider:, session:, user:, user_session:)
+  # @param [Boolean] persisted_read_only Avoid new bucket assignment if test is configured to be
+  #   persisted but there is no persisted value.
+  def bucket(
+    request:,
+    service_provider:,
+    session:,
+    user:,
+    user_session:,
+    persisted_read_only: false
+  )
     return nil if !active?
 
     discriminator = resolve_discriminator(
@@ -59,16 +85,28 @@ def bucket(request:, service_provider:, session:, user:, user_session:)
     )
     return nil if discriminator.blank?
 
+    persisted_value = AbTestAssignment.bucket(experiment:, discriminator:) if persist?
+    return persisted_value if persisted_value || (persist? && persisted_read_only)
+
+    return nil if maxed?
+
     user_value = percent(discriminator)
 
+    bucket = @default_bucket
+
     min = 0
     buckets.keys.each do |key|
       max = min + buckets[key]
-      return key if user_value > min && user_value <= max
+      if user_value > min && user_value <= max
+        bucket = key
+        break
+      end
       min = max
     end
 
-    @default_bucket
+    AbTestAssignment.create(experiment:, discriminator:, bucket:) if persist?
+
+    bucket
   end
 
   def include_in_analytics_event?(event_name)
@@ -87,8 +125,18 @@ def active?
     @active
   end
 
+  def participants_count
+    return if !persist?
+    AbTestAssignment.where(experiment:).count
+  end
+
   private
 
+  def maxed?
+    return false if !persist? || !max_participants.finite?
+    participants_count >= max_participants
+  end
+
   def resolve_discriminator(user:, **)
     if @discriminator
       @discriminator.call(user:, **)

lib/reporting/ab_tests_report.rb

@@ -5,15 +5,15 @@
 
 module Reporting
   class AbTestsReport
-    attr_reader :queries, :time_range
+    attr_reader :ab_test, :time_range
 
-    # @param [Array<AbTest::ReportQueryConfig>] queries
+    # @param [AbTest] ab_test
     # @param [Range<Time>] time_range
     def initialize(
-      queries:,
+      ab_test:,
       time_range:
     )
-      @queries = queries
+      @ab_test = ab_test
       @time_range = time_range
     end
 
@@ -30,8 +30,18 @@ def as_emailable_reports
       end
     end
 
+    def participants_message
+      return unless ab_test.persist?
+      message = "Total participants: #{participants_count.to_fs(:delimited)}"
+      message += " (of #{max_participants.to_fs(:delimited)} maximum)" if max_participants.finite?
+      message
+    end
+
     private
 
+    delegate :participants_count, :max_participants, :report, to: :ab_test
+    delegate :queries, to: :report
+
     def table_for_query(query)
       query_data = fetch_results(query: query.query)
       headers = column_labels(query_data.first)

lib/tasks/ab_tests.rake

@@ -0,0 +1,42 @@
+# frozen_string_literal: true
+
+namespace :ab_tests do
+  desc 'Deletes persisted A/B test assignments for outdated tests'
+  task :delete_outdated, [:confirm] => :environment do |_t, args|
+    args.with_defaults(confirm: false)
+
+    configured_experiments = AbTests.all.values.map(&:experiment)
+    outdated_assignments = AbTestAssignment.where.not(experiment: configured_experiments)
+
+    names = outdated_assignments.distinct.pluck(:experiment).sort
+    if names.empty?
+      warn 'No outdated test assignments to delete!'
+      next
+    end
+
+    puts 'Found outdated test assignments with experiment names:'
+    names.each { |name| puts "  - #{name}" }
+    puts ''
+
+    if args.confirm
+      deleted = log_to_stdout { outdated_assignments.in_batches.delete_all }
+      puts "\nSuccessfully deleted #{deleted} #{'record'.pluralize(deleted)}."
+    else
+      puts <<~STR
+        Re-run command with `confirm` arg to delete:
+
+          rake "ab_tests:delete_outdated[confirm]"
+
+      STR
+    end
+  end
+end
+
+def log_to_stdout
+  original_logger = ActiveRecord::Base.logger
+  stdout_logger = Logger.new(STDOUT)
+  ActiveRecord::Base.logger = ActiveSupport::BroadcastLogger.new(original_logger, stdout_logger)
+  yield
+ensure
+  ActiveRecord::Base.logger = original_logger
+end

spec/factories/ab_test_assignment.rb

@@ -0,0 +1,9 @@
+FactoryBot.define do
+  factory :ab_test_assignment do
+    ab_test = AbTests.all.values.sample
+
+    experiment { ab_test.experiment }
+    discriminator { Random.uuid }
+    bucket { [*ab_test.buckets.keys, ab_test.default_bucket].sample }
+  end
+end

spec/jobs/reports/ab_tests_report_spec.rb

@@ -22,9 +22,15 @@
     {
       RECAPTCHA_SIGN_IN: AbTest.new(
         experiment_name: 'reCAPTCHA at Sign-In',
+        persist: true,
+        max_participants: 10_000,
         buckets: { sign_in_recaptcha: tested_percent },
         report: { email:, queries: },
       ),
+      UNREPORTED_EXPERIMENT: AbTest.new(
+        experiment_name: 'Unreported Experiment',
+        buckets: { group_a: 1 },
+      ),
     }
   end
 
@@ -47,12 +53,16 @@
     end
 
     let(:report) do
-      double(Reporting::AbTestsReport, as_emailable_reports: emailable_reports)
+      double(
+        Reporting::AbTestsReport,
+        as_emailable_reports: emailable_reports,
+        participants_message: 'Total participants: 0 (of 10,000 maximum)',
+      )
     end
 
     before do
       allow(Reporting::AbTestsReport).to receive(:new).with(
-        queries:,
+        ab_test: ab_tests[:RECAPTCHA_SIGN_IN],
         time_range: report_date.yesterday..report_date,
       ).and_return(report)
 
@@ -63,7 +73,10 @@
       expect(ReportMailer).to receive(:tables_report).with(
         email:,
         subject: "A/B Tests Report - reCAPTCHA at Sign-In - #{report_date}",
-        message: "A/B Tests Report - reCAPTCHA at Sign-In - #{report_date}",
+        message: [
+          "A/B Tests Report - reCAPTCHA at Sign-In - #{report_date}",
+          'Total participants: 0 (of 10,000 maximum)',
+        ],
         reports: emailable_reports,
         attachment_format: :csv,
       )