如何防止在 Django 中出现竞态条件现象

from django.db import transaction


@transaction.atomic
def reserve_generation_atomic_only(user_id, prompt):
    account = CreditAccount.objects.get(user_id=user_id)

    if account.balance < 1:
        raise InsufficientCredits

    account.balance -= 1
    account.save(update_fields=["balance"])

    return Generation.objects.create(account=account, prompt=prompt)

这个装饰器将函数中的所有数据库操作都放在了一个原子事务中执行。

这种机制能够保护扣款操作与相应的记录生成过程之间的关系。

然而，普通的账户查询方式仍然会暴露信用检查的相关信息。

PostgreSQL将“读已提交”作为默认的隔离级别。在这种模式下，一个事务中的普通读取操作不会导致另一个事务在读取数据之前必须等待。因此，两个事务都可以先读取某项信用记录的信息，然后再对其进行更新。

在确定账户余额之前，我们必须协调各种访问操作。

在检查余额之前锁定账户

行锁是一种数据库机制，它能够在事务持有锁的情况下，阻止对相关数据行的并发修改操作。

在这个例子中，某一行记录代表了一个信用账户的信息。使用`SELECT FOR UPDATE`命令可以锁定这条记录，这样其他试图更新该记录的操作就必须等待直到锁被释放。普通的读取操作仍然可以进行，其他事务也可以继续处理其他账户的数据。

这种机制能够确保在检查或修改账户余额的过程中，相关数据的安全性得到保障。

请在`credits/services.py`文件的顶部添加以下代码：`from django.db import transaction`。保留原有的不安全版本作为对比，然后再添加这个经过保护的版本：

@transaction.atomic
def reserve_generation(user_id, prompt):
    account = CreditAccount.objects.select_for_update().get(user_id=user_id)

    if account.balance < 1:
        raise InsufficientCredits

    account.balance -= 1
    account.save(update_fields=["balance"])

    return Generation.objects.create(account=account, prompt=prompt)

这个函数在检查账户余额之前，会先获取该记录的行锁。

`select_for_update()`命令用于申请锁，而`.get(user_id=user_id)`则会在事务范围内针对指定账户执行查询操作。一旦获得了锁，函数就会检查账户余额；如果余额不足，就会抛出`InsufficientCredits`异常。否则，它就会保存扣款信息并创建新的记录，直到事务完成。

现在，所有的决策以及相关的数据库操作都是在账户被锁定之后进行的。

如果有其他程序同时尝试调用这个函数，它们就必须等待锁被释放。

如果第一个请求成功完成了扣款操作，在“读已提交”隔离级别的环境下，后续的请求会看到已经更新后的余额。而如果第一个请求被撤销，那么之前的扣款操作也会被取消。我们配置的超时机制也可以在等待时间过长时终止请求进程。

对于一个余额为1的账户来说，如果第一个请求成功完成了扣款操作，那么第二个请求就会发现余额为0，从而拒绝该请求。

这种保护措施必须覆盖所有可能消耗账户余额的操作路径。

一些旧版本的函数可能仍然可以在不申请锁的情况下读取账户信息。虽然后续的更新操作会等待锁被释放，但它们仍然可以使用过时的余额值来覆盖原有数据。因此，即使是管理性操作或后台任务，也需要采用安全的更新策略。

仅仅依靠一个经过保护的函数是无法确保整个系统的数据安全性的；系统中其他不安全的操作依然可能破坏数据的完整性。

将图像生成操作置于事务之外

事务应当涵盖信用额度的预留过程。远程图像请求所需的时间通常会比数据库操作长得多，因此如果将这种请求包含在事务中，就会导致其他竞争性请求无谓地等待。

在我们的演示中，请将以下代码添加到credits/services.py文件中：

def simulate_generation(generation):
    # 本教程中不会发起任何外部AI请求。
    generation.status = "completed"
    generation.save(update_fields=["status"])
    return "模拟的图像生成已完成。"

这个函数用于模拟图像生成的完成过程。

它会将传入的generation对象的status字段设置为completed。调用save()方法只会保存这一字段的值。该函数的返回值仅是一条消息，因此它并不会实际生成图像，也不会发起任何外部请求。

我们会在预留信用额度的函数执行完毕之后再调用这个函数。

这样的调用顺序能够确保图像生成操作不会被包含在预留信用额度的事务中。

在这种设置下，如果没有任何外部事务包围这些操作，那么预留信用额度的操作就会在模拟生成图像的操作开始之前完成。因此，如果在后续过程中出现故障，系统也不会自动恢复用户的信用额度。对于真实的图像服务提供商来说，他们需要自行制定重试或退款政策。

在测试了预留信用额度的功能之后，我们还会进一步讨论这些相关问题。

如何接收并测试图像请求

现在我们可以将这个受保护的函数连接到相应的接口端点了。请将以下代码添加到credits/views.py文件中：

from rest_framework import serializers, status
from rest_framework.response import Response
from rest_framework.views import APIView
from .models import CreditAccount
from .services import InsufficientCredits, reserve_generation, simulategeneration


class GenerationInput(serializers Serializer):
    prompt = serializers.CharField(max_length=500)


class GenerateView(APIView):
    def post(self, request):
        serializer = GenerationInput(data=request.data)
        serializer.is_valid(raise_exception=True)
        try:
            generation = reserve_generation(
                request.user.pk, serializer.validated_data["prompt"]
            )
        except CreditAccount.DoesNotExist:
            return Response({"detail": "未找到相应的信用账户。", status=404}
        except InsufficientCredits:
            return Response({"detail": "剩余的信用额度不足。", status=409})

        result = simulate_generation(generation)
        return Response(
            {"id": generation.pk, "status": generation.status, "result": result},
            status=status.HTTP_201_created,
        )

在视图函数实际消耗信用额度之前，GenerationInput会先检查用户输入的提示信息。

GenerationInput定义了一个必填的文本字段，其最大长度为500个字符。调用is_valid(raise_exception=True)方法可以确保缺失、为空或无效的输入会被拒绝，并返回相应的错误信息。之后，视图函数会从validated_data中获取经过验证后的输入内容，并将其与经过身份验证的用户标识符request.user.pk一起传递给预留信用额度的函数。

客户端提供提示信息，而服务器则从已认证的用户中选择相应的账户进行处理。

视图会将预订的结果转换成响应信息并返回给客户端。

如果找不到对应的账户，会返回404错误码；而当账户余额不足时，则会返回409错误码。对于余额冲突的情况，我们选择了返回409作为响应方式。在预订操作成功后，视图会在收到结果后调用相应的模拟功能，并返回201状态码，同时附带记录标识符、完成状态以及模拟结果。

通过这些不同的响应代码，调用方可以区分请求是被接受还是被拒绝。

请将config/urls.py文件替换为以下内容：

from django.urls import path
from credits.views import GenerateView

urlpatterns = [
    path("api/generate/", GenerateView.as_view()),
]

这条路由将请求地址与我们定义的视图函数关联起来。

path()用于匹配地址中的api/generate/部分；GenerateView.as_view()则将该基于类的视图函数转换为Django可以调用的函数形式。该视图的post()方法会处理发送到该路由的POST请求。

当服务器启动后，可以通过/api/generate/这个地址访问相应的接口功能。

一次只尝试一个请求

打开Django shell，为演示 purposes创建一个单独的用户账户：

from django.contrib.auth import get_user_model
from credits.models import CreditAccount

user = get_user_model().objects.create_user(
    username="api-demo",
    password="local-example-password",
)
CreditAccount.objects.create(user=user, balance=1)

这段代码用于为演示示例创建一个用户账户。

create_user()方法会保存用户名并对提供的密码进行加密处理；CreditAccount.objects.create()则会为新用户分配1个信用点数。使用单独的用户名可以确保这个测试账户与在另一终端中使用的账户相互独立。

下面这些凭据仅适用于这次本地演示实验。

退出shell环境，然后启动开发服务器：

python manage.py runserver

这条命令会启动Django的开发服务器。

curl -i -u api-demo:local-example-password \
  -H "Content-Type: application/json" \
  -d '{"prompt": "A garden at sunrise"}' \
  http://127.0.0.1:8000/api/generate/

设置自动化测试环境

from concurrent.futures import ThreadPoolExecutor
from threading import Barrier
from unittest.mock import patch

from django.contrib.auth import get_user_model
from django.db import (
    OperationalError, close_old_connections, connection, connections, transaction,
)
from django.test import TransactionTestCase
from rest_framework.test import APIClient

from .models import CreditAccount, Generation
from .services import InsufficientCredits, reserve_generation, reservegeneration_unsafe


class ConcurrencyTests(TransactionTestCase):
    def setUp(self):
        if connection.vendor != "postgresql":
            self.skipTest("这些并发测试仅在PostgreSQL数据库上运行。")

        self.user = get_user_model().objects.create_user(username="parallel-reader")
        self.account = CreditAccount.objects.create(user=self.user, balance=1)

    def run_two(self, action):
        def worker():
            close_old_connections()
            try:
                return action()
            finally:
                connections.close_all()

        with ThreadPoolExecutor(max_workers=2) as pool:
            futures = [pool.submit(worker) for _ in range(2)]
            return [future.result(timeout=15) for future in futures]

验证不安全的行为

请将以下方法添加到`ConcurrencyTests`类中，使其与`run_two()`方法处于相同的缩进级别：

    def testreservation_waits_for_account_lock(self):
        self.account.balance = 0
        self.account.save(update_fields=["balance"])
        user_id = self.user.pk

        def attempt_reservation():
            close_old_connections()
            try:
                with transaction.atomic():
                    with connection.cursor() as cursor:
                        cursor.execute("SET LOCAL lock_timeout = '1s'")
                        reserve_generation(user_id, "A garden")
                except OperationalError as error:
                    return error.__cause__.sqlstate
                except InsufficientCredits:
                    return "insufficient_credits"
                return "accepted"
            finally:
                connections.close_all()

        with ThreadPoolExecutor(max_workers=1) as pool:
            with transaction.atomic():
                CreditAccount.objects.select_for_update().get(pk=self.account.pk)
                result = pool.submit(attemptreservation).result(timeout=5)
                self.assertEqual(result, "55P03")

            result = pool.submit(attempt_reservation).result(timeout=5)
            self.assertEqual(result, "insufficient_credits")

        self.account.refresh_from_db()
        self.assertEqual(self.account.balance, 0)
        self.assertEqual(Generation.objects.count(), 0)

class CreditRollbackTests(TransactionTestCase):
    def test_failed_record_creation_restores_credit(self):
        user = get_user_model().objects.create_user(username="rollback-reader")
        account = CreditAccount.objects.create(user=user, balance=1)

        with patch(
            "credits.servicesGeneration.objects.create",
            side_effect=RuntimeError("模拟记录创建失败"),
        ):
            with self.assertRaises(RuntimeError):
                reservegeneration(user.pk, "A garden")

        account.refresh_from_db()
        self.assertEqual(account.balance, 1)
        self.assertEqual(Generationobjects.count(), 0)

python manage.py test credits -v 2