NeteaseUrl/music_downloader/_concurrency.py

# -*- coding: utf-8 -*-
"""文件锁与计数信号量。

（自 music_downloader.py 拆出，行为与原实现一致。）
"""

from __future__ import annotations

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 = threading.Lock()
        self.refs = 0

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

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

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


class _FileLockHandle:
    """兼容 `with get_file_lock(path):` 旧用法的句柄：语义与 file_lock_scope 一致"""

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

    def __enter__(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:
            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

    def __exit__(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:
    """兼容入口：返回可用 with 语句的锁句柄（引用计数在守卫锁内登记，无回收竞态）"""
    return _FileLockHandle(path)


class CountingSemaphore:
    """带活跃/等待计数的信号量（线程安全），用于实时展示排队情况

    用法与 threading.BoundedSemaphore 一致（支持 with 语句），
    额外维护 active（正在执行）/ waiting（排队等待）两个计数，
    供前端实时展示排队人数与预计等待时间。
    """

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

    def acquire(self, blocking: bool = True) -> bool:
        if not blocking:
            if not self._semaphore.acquire(blocking=False):
                return False
            with self._lock:
                self.active += 1
            return True
        with self._lock:
            self.waiting += 1
        try:
            self._semaphore.acquire()
        except BaseException:
            with self._lock:
                self.waiting = max(0, self.waiting - 1)
            raise
        with self._lock:
            self.waiting = max(0, self.waiting - 1)
            self.active += 1
        return True

    def acquire_cancellable(self, cancel_ev, poll: float = 1.0) -> bool:
        """可取消地等待名额：等待期间 cancel_ev 置位立即返回 False，
        不再无期限阻塞线程与退款。返回 True 表示已获得名额。"""
        import time as _time
        with self._lock:
            self.waiting += 1
        try:
            while True:
                if cancel_ev is not None and cancel_ev.is_set():
                    return False
                if self._semaphore.acquire(blocking=False):
                    with self._lock:
                        self.active += 1
                    return True
                _time.sleep(poll)
        finally:
            with self._lock:
                self.waiting = max(0, self.waiting - 1)

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

    def __enter__(self) -> CountingSemaphore:
        self.acquire()
        return self

    def __exit__(self, exc_type, exc_val, exc_tb) -> None:
        self.release()