使用数据库触发函数解决竞赛条件

78 阅读6分钟

在这篇文章中,我们将使用两个ActiveRecord模型,像这样。

ERD

举例来说,我们有一个预算模型,其中分配给预算的金额是allocation_cents,它可以用来购买预算中建议的项目。

一个预算可以有很多项目,我们需要确保这些项目的价格(price_cents)之和不超过预算的分配额(allocation_cents)。

这里是budget.rb模型文件:

class Budget < ApplicationRecord
  has_many :items, dependent: :delete_all
end

这里是item.rb模型文件,其中有验证功能,检查项目的价格总和是否超过了预算分配额:

class Item < ApplicationRecord
  belongs_to :budget

  validate :total_price_of_items, if: proc { |t| t.budget.present? }

  def total_price_of_items
    used_budget = 0
    # get the sum of the existing items price in the budget, if they exist
    used_budget = budget.items.map(&:price_cents).reduce(:+) if budget.items.count.positive?

    # add the sum with the current item's price (which will be added into the budget)
    # if the sum is more than the budget allocation, raise error
    if used_budget + price_cents > budget.allocation_cents
      errors.add(:base, 'Items total price has exceeded budget allocation')
    end
  end
end

例如,如果我们有10000美分的预算,而我们不断地在预算中加入一个价格为3000美分的项目,那么在第四个项目上就会产生错误(4 x 3000 = 12000,大于10000)。

这种验证工作很好......直到你有多个用户几乎同时向预算添加项目。

当多个用户在(几乎)同一时间向预算添加项目时,模型验证通过,因为验证是在几乎同一时间完成的,而每个用户的项目还没有添加进去。

时间轴:

  1. 用户1检查项目的总价格是否大于预算分配,验证通过。
  2. 用户2检查项目总价是否大于预算分配,验证通过。
  3. 用户3检查项目总价是否大于预算分配,验证通过。
  4. 用户1将该项目加入预算
  5. 用户2将项目添加到预算中
  6. 用户3将项目添加到预算中

我们可以通过创建多个线程并加入它们来模拟这个由多个用户并发添加项目的场景,在同一时间运行多个添加项目语句:

threads = concurrency_level.times.map do
  Thread.new do
   # dont start execution until we allow it to
   true while wait_for_it
   # the thread will keep looping the above line until we change 'wait_for_it' to false  
     
   budget.items.create(name: 'beers', price_cents: 4000)
  end
end

wait_for_it = false
# this will create 4 beer item, simultaneously
threads.each(&:join)

执行后,预算中会有4个啤酒瓶,总价12000,超过了10000的分配额度!为了防止这种情况,我们可以使用一个新的线程,将其加入预算中。

为了防止这种情况,我们可以利用数据库函数,在数据库中创建记录之前就进行验证。

用于验证的数据库函数和触发器

由于没有ActiveRecord函数来创建SQL函数/触发器,我们必须自己编写原始SQL:

rails g migration CreateTriggerItemTotalCheck

然后在迁移文件中:

class CreateTriggerItemTotalCheck < ActiveRecord::Migration[6.0]
  def up
    execute <<-SQL
      CREATE OR REPLACE FUNCTION check_item_total()
        RETURNS TRIGGER 
      AS $func$
      DECLARE
        allowed_total BIGINT;
        new_total     BIGINT;
      BEGIN
        SELECT INTO allowed_total allocation_cents
        FROM budgets
        WHERE id = NEW.budget_id;
       
        SELECT INTO new_total SUM(price_cents)
        FROM items
        WHERE budget_id = NEW.budget_id;
       
        IF new_total > allowed_total
        THEN
          RAISE EXCEPTION 'Items total price [%] is larger than budget allocation [%]',
          new_total,
          allowed_total;
        END IF;
        RETURN NEW;
      END;
      $func$ 
      LANGUAGE plpgsql;

      CREATE TRIGGER item_total_trigger
      AFTER INSERT OR UPDATE ON items
          FOR EACH ROW EXECUTE PROCEDURE check_item_total();
    SQL
  end

  def down
    execute <<-SQL
      DROP TRIGGER item_total_trigger ON items;
    SQL
  end
end

CREATE_AND_AND_REPLACE_FUNCTION将创建一个名为check_item_total()的函数,如果该函数已经存在,则替换它。这个函数不接受任何参数,并返回一个触发器类型(RETURNS TRIGGER),你可以认为触发器就像ActiveRecord的回调,我们可以设置它在一个模型对象(行)被创建(插入)后执行。

AS funcfuncfuncfunc之间是实际的SQL函数,我们可以认为AS funcfuncfuncfunc是Ruby中多行字符串的分隔符,类似于"-SQLSQL部分。

DECLARE部分,我们可以声明两个变量allowed_totalnew_total,它们的类型都是BIGINT。

实际的函数位于BEGINEND语句之间。

函数中的 "NEW"指的是我们要插入到项目表中的新行,或者我们要更新项目表中的现有行。

new row

SELECT INTO allowed_total allocation_cents
FROM budgets
WHERE id = NEW.budget_id;

上面的语句将从新项目行所属的预算中选择 "allocation_cents "值,并将其保存到变量allowed_total中:

SELECT INTO new_total SUM(price_cents)
FROM items
WHERE budget_id = NEW.budget_id;

上面的语句将选择属于预算的所有项目的价格之和,并保存到变量new_total中。由于函数(触发器)是在新项目插入后运行的,这个新总数包括新行的价格:

IF new_total > allowed_total
THEN
  RAISE EXCEPTION 'Items total price [%] is larger than budget allocation [%]',
  new_total,
  allowed_total;
END IF;

如果新的总数(所有项目价格的总和)大于允许的总数(预算的分配),则引发一个异常,这将使新项目行的插入回滚。

如果没有异常,即创建成功,我们通过返回新行来结束这个函数(RETURN NEW)。

在这个函数之后,我们创建了一个触发器,每当有新的行插入到items表中,或者items表中的现有行被更新时,这个触发器就会被执行。

CREATE TRIGGER item_total_trigger
AFTER INSERT OR UPDATE ON items
    FOR EACH ROW EXECUTE PROCEDURE check_item_total();

在items表中每插入或更新一行后,check_item_total函数将被运行。

对于向下迁移,我们可以选择删除触发器作为反转。

现在,当我们试图添加更多的项目时,当预算分配超过时,它将回滚并抛出一个 "ActiveRecord::StatementInvalid"异常。

在你的控制器动作中,你可以拯救这个异常并显示错误信息:

# budgets_controller.rb

def update
  # ...
rescue ActiveRecord::StatementInvalid => e
  flash[:error] = "Items total exceeded budget allocation, please try again"
  render 'edit'
end

使用fx gem向schema.rb添加触发器

如果你在迁移SQL触发器后打开schema.rb,你会发现里面没有关于触发器函数的信息。当你使用测试数据库运行测试时,这将导致问题,因为rake db:migrate RAILS_ENV=test ,只是将schema.rb的结构复制到数据库中(使用rake db:schema:load),而没有逐一查看所有的迁移文件。

我们可以通过使用**fx gem**来解决这个问题,这个gem会在迁移后将SQL函数纳入schema.rb。

在继续之前,请确保回滚到添加SQL函数之前的迁移,并删除SQL函数的迁移文件,以确保与fx gem创建的迁移没有冲突。

在你的Gemfile中包含'fx'gem。

gem 'fx'

然后运行bundle install

安装完fx gem后,我们可以使用它的生成器来创建一个函数,在终端中,运行.NET。

rails generate fx:function check_item_total

这将生成一个新的迁移文件和一个空白的sql文件,在db/functions/check_item_total_v01.sql

然后我们就可以把之前迁移中的SQL函数(不包括触发器)移到这个文件中。

sql_file

接下来我们将使用这个命令为这个函数创建一个触发器。

rails generate fx:trigger item_total_trigger table_name:items

这将生成一个新的迁移文件和一个在db/triggers/item_total_trigger_v01.sql的空白sql文件。

然后我们就可以把之前迁移中的SQL触发器移到这个文件中:

trigger sql

现在我们可以运行rake db:migrate ,将该函数和触发器添加到模式和数据库中。

这一次,SQL函数和触发器出现在schema.rb中!

测试竞赛条件

Arkency在这里写了一篇关于如何测试竞赛条件的好文章。我们可以参考它来模拟竞赛条件,通过使用多个线程(使用rspec)同时添加多个项目:

require 'rails_helper'

describe 'Create multiple items for the budget at the same time' do
  context 'items price more than allocation' do

    let!(:budget) { create(:budget, allocation_cents: 10_000) }

    it 'should fail' do
      # https://blog.arkency.com/2015/09/testing-race-conditions/
      expect(budget.allocation_cents).to eq(10_000)

      fail_occurred = false
      wait_for_it = true

      # create multiple threads to create the same payment at the same time
      threads = 4.times.map do
       Thread.new do
         # halt execution until we allow it to
         true while wait_for_it
         begin
           budget.items.create(name: 'beers', price_cents: 3000)
         rescue ActiveRecord::StatementInvalid => e
           # SQL exception will be thrown, then ActiveRecord will throw statement invalid
           # and we will catch it here
           fail_occurred = true
         end
       end
      end
      wait_for_it = false
      threads.each(&:join)

      # Add delay in case the previous database transaction rollback is not finished yet
      sleep 2
      
      # only 3 items should pass through, 3 x 3000 = 9000
      expect(budget.items.count).to eq(3)
      
      # sum of budget items price should be equal or smaller than budget allocation
      expect(budget.items.map(&:price_cents).reduce(:+)).to be <= budget.allocation_cents
      
      # ActiveRecord should throw an error
      expect(fail_occurred).to eq(true)
    end
  end
end

在添加了数据库触发器和函数后,这个规范应该通过。但是在一些旧的Postgresql版本(12.2和更早的版本)中,我遇到了一个错误,在运行rake db:migrate后,模式导出了 "FOR EACH ROW EXECUTEFUNCTION"而不是 "FOR EACH ROW EXECUTEPROCEDURE",尽管我们在SQL文件中输入了 "FOR EACH ROW EXECUTE PROCEDURE"。这可能会导致数据库函数没有被调用。

一个简单的解决方法是在schema.rb文件中手动将 "FUNCTION "改为 "PROCEDURE",然后运行rake db:migrate RAILS_ENV=test,要求测试数据库再次复制模式。

有了数据库触发器,我们可以更好地防止竞赛条件,因为验证是在数据库级别而不是应用程序级别完成的。