63 lines
2.6 KiB
Python
63 lines
2.6 KiB
Python
from pathlib import Path
|
|
import os
|
|
import shutil
|
|
|
|
wrapper = '''
|
|
|
|
_key_acquisition_lock = threading.Lock()
|
|
_key_acquisition_inflight = {}
|
|
|
|
|
|
def acquire_local_keys(timeout: float = 60.0) -> None:
|
|
"""Share concurrent initialization of the same local data directory only.
|
|
|
|
There is no completed-result cache: a later explicit request scans again.
|
|
Every caller still validates its own account database after acquisition.
|
|
"""
|
|
root = application_data_dir()
|
|
try:
|
|
config = json.loads((root / "wxwork_gui_config.json").read_text(encoding="utf-8-sig"))
|
|
source = str(config.get("db_dir") or "") if isinstance(config, dict) else ""
|
|
except (OSError, ValueError):
|
|
source = ""
|
|
identity = (str(root.resolve()).casefold(), source.casefold(),
|
|
os.environ.get("WECOM_ARCHIVE_SOURCE_DIR", "").casefold())
|
|
with _key_acquisition_lock:
|
|
future = _key_acquisition_inflight.get(identity)
|
|
owner = future is None
|
|
if owner:
|
|
future = Future()
|
|
_key_acquisition_inflight[identity] = future
|
|
if not owner:
|
|
try:
|
|
return future.result(timeout=max(0.0, float(timeout)))
|
|
except FutureTimeoutError as exc:
|
|
raise ValueError("正在获取本机密钥，等待超时。请稍后重试，或导入已有密钥文件") from exc
|
|
try:
|
|
result = _acquire_local_keys_once(timeout)
|
|
except BaseException as exc:
|
|
future.set_exception(exc)
|
|
raise
|
|
else:
|
|
future.set_result(result)
|
|
return result
|
|
finally:
|
|
with _key_acquisition_lock:
|
|
if _key_acquisition_inflight.get(identity) is future:
|
|
_key_acquisition_inflight.pop(identity, None)
|
|
'''
|
|
|
|
for root in (Path(r'C:\kefu\wechat_rpa'), Path(r'C:\wechat_rpa')):
|
|
path = root / 'wxwork_local_setup.py'
|
|
backup = root / 'backups/loading-performance-20260917'
|
|
backup.mkdir(parents=True, exist_ok=True)
|
|
assert not (backup / path.name).exists()
|
|
shutil.copy2(path, backup / path.name)
|
|
text = path.read_text(encoding='utf-8')
|
|
text = text.replace('import tempfile\n', 'import tempfile\nimport threading\nfrom concurrent.futures import Future, TimeoutError as FutureTimeoutError\n', 1)
|
|
text = text.replace('def acquire_local_keys(timeout: float = 60.0) -> None:', 'def _acquire_local_keys_once(timeout: float = 60.0) -> None:', 1)
|
|
text += wrapper
|
|
compile(text, str(path), 'exec')
|
|
tmp = path.with_suffix('.performance-tmp');tmp.write_text(text, encoding='utf-8');os.replace(tmp, path)
|
|
print(str(root) + ': simultaneous key requests share one worker')
|