gongxiangp2/lib/__tests__/image-processing-limiter.test.ts

/** @jest-environment node */

import { EventEmitter } from 'events'
import { mkdtemp, rm } from 'fs/promises'
import os from 'os'
import path from 'path'

describe('image-processing-limiter', () => {
  const originalPlatform = Object.getOwnPropertyDescriptor(process, 'platform')
  const originalNodeEnv = process.env.NODE_ENV
  const originalUploadDir = process.env.UPLOAD_DIR

  afterEach(() => {
    delete process.env.IMAGE_PROCESSING_MAX_CONCURRENT
    const mutableEnv = process.env as Record<string, string | undefined>
    if (originalNodeEnv === undefined) delete mutableEnv.NODE_ENV
    else mutableEnv.NODE_ENV = originalNodeEnv
    if (originalUploadDir === undefined) delete process.env.UPLOAD_DIR
    else process.env.UPLOAD_DIR = originalUploadDir
    if (originalPlatform) Object.defineProperty(process, 'platform', originalPlatform)
    jest.dontMock('child_process')
    jest.resetModules()
  })

  test('并发为 2 时同时释放两个任务，名额仍先直接转交给既有等待者', async () => {
    jest.resetModules()
    process.env.IMAGE_PROCESSING_MAX_CONCURRENT = '2'
    const { withImageProcessingSlot } = await import('@/lib/image-processing-limiter')

    const releases: Array<() => void> = []
    const started: number[] = []
    let active = 0
    let peak = 0

    const runTask = (index: number) => withImageProcessingSlot(() => {
      active += 1
      peak = Math.max(peak, active)
      started.push(index)
      return new Promise<void>((resolve) => {
        releases[index] = () => {
          active -= 1
          resolve()
        }
      })
    })

    const firstTask = runTask(0)
    const secondTask = runTask(1)
    const firstQueuedTask = runTask(2)
    const secondQueuedTask = runTask(3)

    await new Promise<void>((resolve) => setImmediate(resolve))
    expect(started).toEqual([0, 1])

    releases[0]()
    releases[1]()
    // 两个 release continuation 已排入微任务后再发起新请求。旧实现会先把
    // active 降到 0，新请求 4 因而抢入，随后等待者 2、3 继续自增到 3。
    const lateTask = Promise.resolve().then(() => runTask(4))
    await new Promise<void>((resolve) => setImmediate(resolve))
    expect(started).toEqual([0, 1, 2, 3])
    expect(peak).toBe(2)

    releases[2]()
    releases[3]()
    await new Promise<void>((resolve) => setImmediate(resolve))
    expect(started).toEqual([0, 1, 2, 3, 4])

    releases[4]()
    await Promise.all([
      firstTask,
      secondTask,
      firstQueuedTask,
      secondQueuedTask,
      lateTask,
    ])
    expect(peak).toBe(2)

  })

  function createFlockChild() {
    const child = new EventEmitter() as EventEmitter & {
      stdin: EventEmitter & { end: jest.Mock }
      stdout: EventEmitter
      stderr: EventEmitter
      exitCode: number | null
      signalCode: NodeJS.Signals | null
      kill: jest.Mock
    }
    child.stdin = Object.assign(new EventEmitter(), { end: jest.fn() })
    child.stdout = new EventEmitter()
    child.stderr = new EventEmitter()
    child.exitCode = null
    child.signalCode = null
    child.kill = jest.fn(() => true)
    child.stdin.end.mockImplementation(() => {
      queueMicrotask(() => {
        child.exitCode = 0
        child.emit('close', 0, null)
      })
    })
    queueMicrotask(() => child.stdout.emit('data', 'IMAGE_PROCESSING_LOCK_READY\n'))
    return child
  }

  test('生产 Linux 使用共享 uploads 上的同一把 exclusive flock，且忽略进程内并发配置', async () => {
    const uploadRoot = await mkdtemp(path.join(os.tmpdir(), 'image-processing-lock-'))
    const children: ReturnType<typeof createFlockChild>[] = []
    const spawnMock = jest.fn(() => {
      const child = createFlockChild()
      children.push(child)
      return child
    })
    jest.doMock('child_process', () => ({ spawn: spawnMock }))
    Object.defineProperty(process, 'platform', { value: 'linux', configurable: true })
    ;(process.env as Record<string, string | undefined>).NODE_ENV = 'production'
    process.env.UPLOAD_DIR = uploadRoot
    process.env.IMAGE_PROCESSING_MAX_CONCURRENT = '2'

    try {
      const { withImageProcessingSlot, __imageProcessingLimiterTesting } = await import(
        '@/lib/image-processing-limiter'
      )
      await withImageProcessingSlot(async () => withImageProcessingSlot(async () => 'done'))

      const expectedLockPath = path.join(uploadRoot, '.locks', 'image-processing-global.lock')
      expect(__imageProcessingLimiterTesting.imageProcessingLockPath()).toBe(expectedLockPath)
      expect(spawnMock).toHaveBeenCalledWith(
        'flock',
        expect.arrayContaining(['--exclusive', expectedLockPath]),
        expect.objectContaining({ shell: false }),
      )
      expect(spawnMock).toHaveBeenCalledTimes(1)
      expect(children[0].stdin.end).toHaveBeenCalledTimes(1)
    } finally {
      await rm(uploadRoot, { recursive: true, force: true })
    }
  })

  test('生产 flock 在任务期间丢失时失败关闭，不返回本次图片结果', async () => {
    const uploadRoot = await mkdtemp(path.join(os.tmpdir(), 'image-processing-lock-lost-'))
    let child: ReturnType<typeof createFlockChild>
    jest.doMock('child_process', () => ({
      spawn: jest.fn(() => {
        child = createFlockChild()
        return child
      }),
    }))
    Object.defineProperty(process, 'platform', { value: 'linux', configurable: true })
    ;(process.env as Record<string, string | undefined>).NODE_ENV = 'production'
    process.env.UPLOAD_DIR = uploadRoot

    try {
      const { withImageProcessingSlot, ImageProcessingLockLostError } = await import(
        '@/lib/image-processing-limiter'
      )
      await expect(withImageProcessingSlot(async () => {
        expect(child!).toBeDefined()
        child.exitCode = 1
        child.emit('close', 1, null)
        await new Promise<void>((resolve) => setImmediate(resolve))
        return 'must-not-escape'
      })).rejects.toBeInstanceOf(ImageProcessingLockLostError)
    } finally {
      await rm(uploadRoot, { recursive: true, force: true })
    }
  })
})