NeteaseUrl/music_downloader/_concurrency.py

# -*- coding: utf-8 -*-
"""文件锁与并发闸（asyncio 原语）。

（自 music_downloader.py 拆出。）

下载链路已整体异步化：同一进程、同一事件循环里跑多个任务，互斥不再靠线程
阻塞，而是靠 ``await`` 让出控制权。故此处两个原语都换成 asyncio 实现——
用 threading 锁会「持锁方在等事件循环、等待方又在事件循环上」而直接死锁。
两个计数器字典仍用 threading.Lock 保护：临界区只有字典读写、内部无 await，
不会把事件循环卡住，也让仍在同步线程里跑的监控循环能安全读取。
"""

from __future__ import annotations

import asyncio
import threading
from pathlib import Path

# per-file 写锁：保证同一目标文件的 .part 临时文件同时只有一个写入者，
# 防止并发下载同名歌曲（不同用户/任务）互相覆盖导致文件损坏
# 锁表带引用计数与回收——无人引用时从表中移除，防止长期运行无界增长
# 计数在注册表守卫锁内完成（"取出引用"即计数 +1），消除"取引用后、acquire 前
# 被另一个任务回收、第三个任务拿到新锁对象"的互斥破坏窗口
_FILE_LOCKS: dict[str, RefCountedFileLock] = {}

_FILE_LOCKS_GUARD = threading.Lock()


class RefCountedFileLock:
    """带引用计数的文件锁：refs 表示"已持有或已登记即将持有"的数量，归零可回收"""

    def __init__(self):
        self._lock = asyncio.Lock()
        self.refs = 0

    async def acquire(self) -> None:
        await self._lock.acquire()

    def release(self) -> None:
        self._lock.release()

    def locked(self) -> bool:
        return self._lock.locked()


class _FileLockHandle:
    """`async with get_file_lock(path):` 句柄：语义与引用计数回收口径同旧实现"""

    def __init__(self, path: Path):
        self._path = path
        self._lock = None
        self._key = ''

    async def __aenter__(self) -> RefCountedFileLock:
        key = str(self._path)
        with _FILE_LOCKS_GUARD:
            lock = _FILE_LOCKS.get(key)
            if lock is None:
                lock = RefCountedFileLock()
                _FILE_LOCKS[key] = lock
            lock.refs += 1
        self._lock = lock
        self._key = key
        try:
            await lock.acquire()
        except BaseException:
            # acquire 被打断时回滚引用计数并摘除空表项，防锁表条目永久泄漏
            with _FILE_LOCKS_GUARD:
                lock.refs -= 1
                if lock.refs <= 0 and not lock.locked():
                    _FILE_LOCKS.pop(key, None)
            self._lock = None
            self._key = ''
            raise
        return lock

    async def __aexit__(self, exc_type, exc_val, exc_tb) -> None:
        if self._lock is None:
            return
        self._lock.release()
        with _FILE_LOCKS_GUARD:
            self._lock.refs -= 1
            if self._lock.refs <= 0 and not self._lock.locked():
                _FILE_LOCKS.pop(self._key, None)
        self._lock = None


def get_file_lock(path: Path) -> _FileLockHandle:
    """兼容入口：返回可用 `async with` 的锁句柄（引用计数在守卫锁内登记，无回收竞态）"""
    return _FileLockHandle(path)


class CountingSemaphore:
    """带活跃/等待计数的并发闸（asyncio），用于实时展示排队情况

    容量由 ``active`` 计数实现，等待靠 ``await asyncio.sleep`` 让出事件循环。
    与旧实现（threading.BoundedSemaphore + 轮询 ``acquire(blocking=False)``）
    同为非公平抢占、同样按 poll 粒度醒来，语义不变；差别只在等待不再阻塞线程。
    """

    def __init__(self, value: int = 1):
        self.capacity = max(1, int(value))
        self._lock = threading.Lock()
        self.active = 0
        self.waiting = 0

    async def acquire_cancellable(self, cancel_ev, poll: float = 1.0) -> bool:
        """可取消地等待名额：等待期间 cancel_ev 置位立即返回 False，
        不再无期限等待。返回 True 表示已获得名额。"""
        with self._lock:
            self.waiting += 1
        try:
            while True:
                if cancel_ev is not None and cancel_ev.is_set():
                    return False
                with self._lock:
                    if self.active < self.capacity:
                        self.active += 1
                        return True
                await asyncio.sleep(poll)
        finally:
            with self._lock:
                self.waiting = max(0, self.waiting - 1)

    def release(self) -> None:
        with self._lock:
            self.active = max(0, self.active - 1)