feat: implement distributed budget locking and database migration hea…

src/common/locks/budget-lock.decorator.spec.ts

@@ -0,0 +1,43 @@
+import { describe, expect, it } from 'vitest';
+import { Reflector } from '@nestjs/core';
+import { BUDGET_LOCK_KEY, UseBudgetLock } from './budget-lock.decorator';
+
+describe('UseBudgetLock', () => {
+  it('sets the budget-lock metadata with default options', () => {
+    class Controller {
+      @UseBudgetLock()
+      allocate(): void {}
+    }
+
+    const reflector = new Reflector();
+    expect(reflector.get(BUDGET_LOCK_KEY, Controller.prototype.allocate)).toEqual({});
+  });
+
+  it('stores custom key resolver and ttl options', () => {
+    const keyResolver = () => 'budget:custom';
+
+    class Controller {
+      @UseBudgetLock({ key: keyResolver, ttl: 1000 })
+      allocate(): void {}
+    }
+
+    const reflector = new Reflector();
+    expect(reflector.get(BUDGET_LOCK_KEY, Controller.prototype.allocate)).toEqual({
+      key: keyResolver,
+      ttl: 1000,
+    });
+  });
+
+  it('is only applied to the decorated method', () => {
+    class Controller {
+      @UseBudgetLock()
+      allocate(): void {}
+
+      list(): void {}
+    }
+
+    const reflector = new Reflector();
+    expect(reflector.get(BUDGET_LOCK_KEY, Controller.prototype.allocate)).toEqual({});
+    expect(reflector.get(BUDGET_LOCK_KEY, Controller.prototype.list)).toBeUndefined();
+  });
+});

src/common/locks/budget-lock.decorator.ts

@@ -0,0 +1,37 @@
+import { applyDecorators, SetMetadata } from '@nestjs/common';
+import { Request } from 'express';
+
+/** Reflector metadata key used by {@link BudgetLockInterceptor}. */
+export const BUDGET_LOCK_KEY = 'budgetLock';
+
+export interface BudgetLockOptions {
+  /**
+   * Static lock key or a resolver producing the lock key from the incoming
+   * request. Defaults to `budget:{params.id}`.
+   */
+  key?: string | ((request: Request) => string);
+  /** Lock time-to-live in milliseconds (defaults to 5000). */
+  ttl?: number;
+}
+
+/**
+ * Decorator that serializes concurrent mutations on the same budget resource
+ * by acquiring a Redis distributed lock around the handler.
+ *
+ * This prevents race conditions when multiple agents or requests simultaneously
+ * attempt to allocate, deduct, or modify budgets — ensuring atomic
+ * checks-and-balances for the same budget resource.
+ *
+ * Usage:
+ * ```ts
+ * @UseBudgetLock()
+ * @Post(':id/allocate')
+ * allocate(...) { ... }
+ * ```
+ *
+ * When the lock cannot be acquired (another request is already mutating the
+ * same budget), the request fails with a `409 LOCK_ACQUISITION_FAILED` error.
+ */
+export function UseBudgetLock(options: BudgetLockOptions = {}): MethodDecorator {
+  return applyDecorators(SetMetadata(BUDGET_LOCK_KEY, options));
+}

src/common/locks/budget-lock.interceptor.spec.ts

@@ -0,0 +1,136 @@
+import { beforeEach, describe, expect, it, vi } from 'vitest';
+import { ExecutionContext, CallHandler } from '@nestjs/common';
+import { Reflector } from '@nestjs/core';
+import { of, throwError } from 'rxjs';
+import { BudgetLockInterceptor } from './budget-lock.interceptor';
+import { RedisLock } from './redis-lock.util';
+import { BUDGET_LOCK_KEY, BudgetLockOptions } from './budget-lock.decorator';
+import { LockNotAcquiredException } from '../exceptions/domain.exception';
+
+function buildContext(req: Record<string, unknown>): ExecutionContext {
+  return {
+    switchToHttp: () => ({ getRequest: () => req }),
+    getHandler: () => handler,
+    getClass: () => Controller,
+  } as unknown as ExecutionContext;
+}
+
+const handler = (): void => {};
+class Controller {}
+
+function decorate(options: BudgetLockOptions | undefined): void {
+  if (options === undefined) {
+    Reflect.deleteMetadata(BUDGET_LOCK_KEY, handler);
+    return;
+  }
+  Reflect.defineMetadata(BUDGET_LOCK_KEY, options, handler);
+}
+
+describe('BudgetLockInterceptor', () => {
+  let reflector: Reflector;
+  let redisLock: { withLock: ReturnType<typeof vi.fn> };
+  let interceptor: BudgetLockInterceptor;
+
+  beforeEach(() => {
+    vi.clearAllMocks();
+    Reflect.deleteMetadata(BUDGET_LOCK_KEY, handler);
+    reflector = new Reflector();
+    redisLock = {
+      withLock: vi.fn().mockImplementation(async (_key: string, fn: () => Promise<unknown>) => fn()),
+    };
+    interceptor = new BudgetLockInterceptor(reflector, redisLock as unknown as RedisLock);
+  });
+
+  it('passes through when the handler is not decorated with @UseBudgetLock()', async () => {
+    decorate(undefined);
+    const ctx = buildContext({ params: { id: 'budget-1' } });
+    const next = { handle: () => of('ok') } as unknown as CallHandler;
+
+    const result = await interceptor.intercept(ctx, next).toPromise();
+
+    expect(result).toBe('ok');
+    expect(redisLock.withLock).not.toHaveBeenCalled();
+  });
+
+  it('acquires a lock on the default budget key for decorated handlers', async () => {
+    decorate({});
+    const ctx = buildContext({ params: { id: 'budget-1' } });
+    const next = { handle: () => of('ok') } as unknown as CallHandler;
+
+    await interceptor.intercept(ctx, next).toPromise();
+
+    expect(redisLock.withLock).toHaveBeenCalledWith('budget:budget-1', expect.any(Function), 5000);
+  });
+
+  it('supports a custom static lock key', async () => {
+    decorate({ key: 'budget:custom' });
+    const ctx = buildContext({ params: { id: 'budget-1' } });
+    const next = { handle: () => of('ok') } as unknown as CallHandler;
+
+    await interceptor.intercept(ctx, next).toPromise();
+
+    expect(redisLock.withLock).toHaveBeenCalledWith('budget:custom', expect.any(Function), 5000);
+  });
+
+  it('supports a custom key resolver function receiving the request', async () => {
+    const resolver = vi.fn().mockReturnValue('budget:resolved');
+    decorate({ key: resolver });
+    const req = { params: { id: 'budget-1' } };
+    const ctx = buildContext(req);
+    const next = { handle: () => of('ok') } as unknown as CallHandler;
+
+    await interceptor.intercept(ctx, next).toPromise();
+
+    expect(resolver).toHaveBeenCalledWith(req);
+    expect(redisLock.withLock).toHaveBeenCalledWith('budget:resolved', expect.any(Function), 5000);
+  });
+
+  it('falls back to body.budgetId when no route param is present', async () => {
+    decorate({});
+    const ctx = buildContext({ params: {}, body: { budgetId: 'budget-9' } });
+    const next = { handle: () => of('ok') } as unknown as CallHandler;
+
+    await interceptor.intercept(ctx, next).toPromise();
+
+    expect(redisLock.withLock).toHaveBeenCalledWith('budget:budget-9', expect.any(Function), 5000);
+  });
+
+  it('uses the configured ttl when provided', async () => {
+    decorate({ ttl: 250 });
+    const ctx = buildContext({ params: { id: 'budget-1' } });
+    const next = { handle: () => of('ok') } as unknown as CallHandler;
+
+    await interceptor.intercept(ctx, next).toPromise();
+
+    expect(redisLock.withLock).toHaveBeenCalledWith('budget:budget-1', expect.any(Function), 250);
+  });
+
+  it('rejects the request when the budget id cannot be resolved', async () => {
+    decorate({});
+    const ctx = buildContext({ params: {}, body: {} });
+    const next = { handle: () => of('ok') } as unknown as CallHandler;
+
+    await expect(interceptor.intercept(ctx, next).toPromise()).rejects.toThrow('budget id');
+    expect(redisLock.withLock).not.toHaveBeenCalled();
+  });
+
+  it('propagates lock acquisition failures to the caller', async () => {
+    decorate({});
+    redisLock.withLock.mockRejectedValue(new LockNotAcquiredException('budget-1'));
+    const ctx = buildContext({ params: { id: 'budget-1' } });
+    const next = { handle: () => of('ok') } as unknown as CallHandler;
+
+    await expect(interceptor.intercept(ctx, next).toPromise()).rejects.toBeInstanceOf(
+      LockNotAcquiredException,
+    );
+  });
+
+  it('propagates handler errors through the lock boundary', async () => {
+    decorate({});
+    const ctx = buildContext({ params: { id: 'budget-1' } });
+    const next = { handle: () => throwError(() => new Error('boom')) } as unknown as CallHandler;
+
+    await expect(interceptor.intercept(ctx, next).toPromise()).rejects.toThrow('boom');
+    expect(redisLock.withLock).toHaveBeenCalledWith('budget:budget-1', expect.any(Function), 5000);
+  });
+});

src/common/locks/budget-lock.interceptor.ts

@@ -0,0 +1,66 @@
+import {
+  BadRequestException,
+  CallHandler,
+  ExecutionContext,
+  Injectable,
+  NestInterceptor,
+} from '@nestjs/common';
+import { Reflector } from '@nestjs/core';
+import { from, lastValueFrom, Observable, throwError } from 'rxjs';
+import { Request } from 'express';
+import { RedisLock, DEFAULT_LOCK_TTL_MS } from './redis-lock.util';
+import { BUDGET_LOCK_KEY, BudgetLockOptions } from './budget-lock.decorator';
+
+/**
+ * Resolves the budget resource id from the request: `req.params.id` with a
+ * fallback to `req.body.budgetId`.
+ */
+function defaultBudgetKeyResolver(request: Request): string {
+  const params = request.params as { id?: string } | undefined;
+  const body = request.body as { budgetId?: string } | undefined;
+  const budgetId = params?.id ?? body?.budgetId;
+  if (!budgetId) {
+    throw new BadRequestException('A budget id is required to acquire the budget lock');
+  }
+  return `budget:${budgetId}`;
+}
+
+/**
+ * Global interceptor enforcing `@UseBudgetLock()`. For handlers decorated with
+ * the decorator it acquires a Redis distributed lock for the budget resource
+ * before invoking the handler and releases it afterwards (including on error),
+ * so concurrent mutations of the same budget are serialized across instances.
+ *
+ * Handlers without the decorator pass straight through untouched.
+ */
+@Injectable()
+export class BudgetLockInterceptor implements NestInterceptor {
+  constructor(
+    private readonly reflector: Reflector,
+    private readonly redisLock: RedisLock,
+  ) {}
+
+  intercept(context: ExecutionContext, next: CallHandler): Observable<unknown> {
+    const options = this.reflector.getAllAndOverride<BudgetLockOptions | undefined>(BUDGET_LOCK_KEY, [
+      context.getHandler(),
+      context.getClass(),
+    ]);
+
+    if (!options) {
+      return next.handle();
+    }
+
+    try {
+      const request = context.switchToHttp().getRequest<Request>();
+      const resourceKey =
+        typeof options.key === 'function'
+          ? options.key(request)
+          : options.key ?? defaultBudgetKeyResolver(request);
+      const ttl = options.ttl ?? DEFAULT_LOCK_TTL_MS;
+
+      return from(this.redisLock.withLock(resourceKey, () => lastValueFrom(next.handle()), ttl));
+    } catch (error) {
+      return throwError(() => error);
+    }
+  }
+}

src/common/locks/locks.module.ts

@@ -6,13 +6,15 @@ import { RedisConfig } from '../../config/redis.config';
 import { REDIS_CLIENT } from './locks.constants';
 import { RedisLock } from './redis-lock.util';
 import { AgentLockInterceptor } from './agent-lock.interceptor';
+import { BudgetLockInterceptor } from './budget-lock.interceptor';
 
 /**
  * Global distributed-locking infrastructure.
  *
  * Provides a single shared ioredis client and the {@link RedisLock} service to
- * every module, and registers the {@link AgentLockInterceptor} that enforces
- * `@UseAgentLock()` on any decorated controller method.
+ * every module, and registers the {@link AgentLockInterceptor} and
+ * {@link BudgetLockInterceptor} that enforce `@UseAgentLock()` and
+ * `@UseBudgetLock()` on any decorated controller method.
  */
 @Global()
 @Module({
@@ -27,6 +29,7 @@ import { AgentLockInterceptor } from './agent-lock.interceptor';
     },
     RedisLock,
     { provide: APP_INTERCEPTOR, useClass: AgentLockInterceptor },
+    { provide: APP_INTERCEPTOR, useClass: BudgetLockInterceptor },
   ],
   exports: [REDIS_CLIENT, RedisLock],
 })

src/modules/budgets/budget.controller.ts

@@ -19,6 +19,7 @@ import {
 } from '@nestjs/swagger';
 import { UserRole } from '@prisma/client';
 import { BudgetService } from './budget.service';
+import { UseBudgetLock } from '../../common/locks/budget-lock.decorator';
 import {
   allocateBudgetSchema,
   AllocateBudgetInput,
@@ -95,6 +96,7 @@ export class BudgetController {
   }
 
   @Patch(':id')
+  @UseBudgetLock()
   @Roles(UserRole.OWNER, UserRole.ADMIN, UserRole.FINANCE)
   @ApiOperation({
     summary: 'Update a budget',
@@ -116,6 +118,7 @@ export class BudgetController {
   }
 
   @Post(':id/allocate')
+  @UseBudgetLock()
   @Roles(UserRole.OWNER, UserRole.ADMIN, UserRole.FINANCE)
   @ApiOperation({
     summary: 'Allocate funds from the parent budget to this child',
@@ -139,6 +142,7 @@ export class BudgetController {
   }
 
   @Delete(':id')
+  @UseBudgetLock()
   @Roles(UserRole.OWNER, UserRole.ADMIN, UserRole.FINANCE)
   @ApiOperation({
     summary: 'Delete (soft) a budget',