# Copyright © Cloud Linux GmbH & Cloud Linux Software, Inc 2010-2021 All Rights Reserved
#
# Licensed under CLOUD LINUX LICENSE AGREEMENT
# http://cloudlinux.com/docs/LICENSE.TXT
import errno
import logging
import os
import threading
from collections.abc import Generator
from contextlib import contextmanager
from dataclasses import dataclass
from xray.internal.constants import request_data_storage, safe_id_pattern
logger = logging.getLogger(__name__)
@dataclass
class TaskCounterStorage:
"""
Typing that has lock and value storage with some syntax sugar.
"""
lock: threading.Lock
next_request_id: int
@property
def processed_requests(self):
return self.next_request_id - 1
_request_id_storage: dict[str, TaskCounterStorage] = {}
_global_storage_lock = threading.Lock()
@contextmanager
def open_local_storage(fake_task_id: str, flush=False) -> Generator[TaskCounterStorage, None, None]:
"""
Open local task information storage.
@param fake_task_id:
unique string, usually obtained as task.fake_id
@param flush:
whether to save data to file right after update
"""
logger.debug('Opening storage %s', fake_task_id)
while True:
storage = _get_or_create_record(fake_task_id)
with storage.lock:
with _global_storage_lock:
if _request_id_storage.get(fake_task_id) is not storage:
# flushed and removed while we were waiting for its lock,
# updating it would be lost: take the record anew
continue
yield storage
if flush:
logger.info('Updating task %s requests counter in file', fake_task_id)
_save_data_to_file(fake_task_id, storage)
return
def remove_local_storage(fake_task_id):
"""
Remove local storage record.
@param fake_task_id:
unique string, usually obtained as task.fake_id
"""
logger.info('Removing memory storage for task %s', fake_task_id)
with _global_storage_lock:
storage = _request_id_storage.get(fake_task_id)
if storage is None:
return
# same order as open_local_storage: record lock, then the global one
with storage.lock, _global_storage_lock:
if _request_id_storage.get(fake_task_id) is storage:
del _request_id_storage[fake_task_id]
def get_task_ids():
"""
List all fake task ids saved in local storage.
"""
return list(_request_id_storage)
def flush_memory_storage(remove=True) -> dict[str, int]:
"""
Save all in-memory task counters on disk.
Return the written values: {fake_task_id: next_request_id}
"""
written = {}
for fake_task_id in list(_request_id_storage.keys()):
logger.info('Flushing task id %s on disk', fake_task_id)
with open_local_storage(fake_task_id) as storage:
_save_data_to_file(fake_task_id, storage)
written[fake_task_id] = storage.next_request_id
if remove:
# under the record lock, so nobody holds the record being removed
# (see the check in open_local_storage)
with _global_storage_lock:
if _request_id_storage.get(fake_task_id) is storage:
del _request_id_storage[fake_task_id]
return written
def read_counter_file(fake_task_id: str) -> int | None:
"""
Read the next_request_id saved on disk for the task, without creating
an in-memory record. None if there is no (valid, readable) counter file:
it never raises, the SIGUSR2 handler calls it before its ack
"""
if not safe_id_pattern.match(fake_task_id):
return None
req_id_file = os.path.join(request_data_storage, fake_task_id)
try:
with open(req_id_file, 'r', opener=_nofollow_opener) as f:
return int(f.read())
except (FileNotFoundError, ValueError):
return None
except OSError as e:
if e.errno == errno.ELOOP:
logger.warning('Refusing to follow symlink: %s', req_id_file)
else:
logger.warning('Unable to read counter file %s: errno %s (%s)', req_id_file, e.errno, e.strerror)
return None
def _nofollow_opener(path, flags):
return os.open(path, flags | os.O_NOFOLLOW, 0o600)
def _save_data_to_file(fake_task_id: str, storage: TaskCounterStorage):
"""
Saves storage data from memory to file.
"""
if not fake_task_id or not safe_id_pattern.match(fake_task_id):
raise ValueError(f'Invalid task_id format: {fake_task_id}')
req_id_file = os.path.join(request_data_storage, fake_task_id)
try:
with open(req_id_file, 'w', opener=_nofollow_opener) as f:
f.write(str(storage.next_request_id))
except OSError as e:
if e.errno == errno.ELOOP:
logger.warning('Refusing to follow symlink: %s', req_id_file)
return
raise
def _get_or_create_record(fake_task_id) -> TaskCounterStorage:
"""
Takes record from local storage or creates new one and returns object
"""
with _global_storage_lock:
if fake_task_id not in _request_id_storage:
_request_id_storage[fake_task_id] = TaskCounterStorage(lock=threading.Lock(), next_request_id=1)
storage = _request_id_storage[fake_task_id]
req_id_file = os.path.join(request_data_storage, fake_task_id)
try:
with open(req_id_file, 'r', opener=_nofollow_opener) as f:
storage.next_request_id = int(f.read())
except FileNotFoundError:
pass
except ValueError:
logger.warning('Corrupted counter file for task %s, resetting to default', fake_task_id)
except OSError as e:
if e.errno == errno.ELOOP:
logger.warning('Refusing to follow symlink: %s', req_id_file)
else:
raise
else:
storage = _request_id_storage[fake_task_id]
return storage