284 lines
13 KiB
Python
284 lines
13 KiB
Python
import logging
|
|
import os
|
|
import shlex
|
|
import signal
|
|
import subprocess
|
|
import sys
|
|
import threading
|
|
import time
|
|
|
|
try:
|
|
import docker
|
|
except ImportError:
|
|
docker = None
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# Worker ID tracking
|
|
worker_id_map = {}
|
|
worker_id_counter = 0
|
|
worker_id_lock = threading.Lock()
|
|
|
|
def get_worker_id():
|
|
"""Assigns a stable, sequential ID to each worker thread."""
|
|
global worker_id_counter
|
|
thread_id = threading.get_ident()
|
|
with worker_id_lock:
|
|
if thread_id not in worker_id_map:
|
|
worker_id_map[thread_id] = worker_id_counter
|
|
worker_id_counter += 1
|
|
return worker_id_map[thread_id]
|
|
|
|
def run_command(cmd, running_processes, process_lock, input_data=None, binary_stdout=False, stream_output=False, stream_prefix="", env=None):
|
|
"""
|
|
Runs a command, captures its output, and returns status.
|
|
If binary_stdout is True, stdout is returned as bytes. Otherwise, both are decoded strings.
|
|
If stream_output is True, the command's stdout/stderr are printed to the console in real-time.
|
|
"""
|
|
logger.debug(f"Running command: {' '.join(shlex.quote(s) for s in cmd)}")
|
|
if env:
|
|
logger.debug(f"With custom environment: {env}")
|
|
process = None
|
|
try:
|
|
# Combine with os.environ to ensure PATH etc. are inherited.
|
|
process_env = os.environ.copy()
|
|
if env:
|
|
# Ensure all values in the custom env are strings
|
|
process_env.update({k: str(v) for k, v in env.items()})
|
|
|
|
# Always open in binary mode to handle both cases. We will decode later.
|
|
process = subprocess.Popen(
|
|
cmd,
|
|
stdin=subprocess.PIPE if input_data else None,
|
|
stdout=subprocess.PIPE,
|
|
stderr=subprocess.PIPE,
|
|
preexec_fn=os.setsid, # Start in a new process group to isolate from terminal signals
|
|
env=process_env
|
|
)
|
|
with process_lock:
|
|
running_processes.add(process)
|
|
|
|
stdout_capture = []
|
|
stderr_capture = []
|
|
|
|
def read_pipe(pipe, capture_list, display_pipe=None, prefix=""):
|
|
"""Reads a pipe line by line (as bytes), appending to a list and optionally displaying."""
|
|
for line in iter(pipe.readline, b''):
|
|
capture_list.append(line)
|
|
if display_pipe:
|
|
# Decode for display
|
|
display_line = line.decode('utf-8', errors='replace')
|
|
# Use print to ensure atomicity and proper handling of newlines
|
|
print(f"{prefix}{display_line.strip()}", file=display_pipe)
|
|
|
|
stdout_display_pipe = sys.stdout if stream_output else None
|
|
stderr_display_pipe = sys.stderr if stream_output else None
|
|
|
|
# We must read stdout and stderr in parallel to prevent deadlocks.
|
|
stdout_thread = threading.Thread(target=read_pipe, args=(process.stdout, stdout_capture, stdout_display_pipe, stream_prefix))
|
|
stderr_thread = threading.Thread(target=read_pipe, args=(process.stderr, stderr_capture, stderr_display_pipe, stream_prefix))
|
|
|
|
stdout_thread.start()
|
|
stderr_thread.start()
|
|
|
|
# Handle stdin after starting to read outputs to avoid deadlocks.
|
|
if input_data:
|
|
try:
|
|
process.stdin.write(input_data.encode('utf-8'))
|
|
process.stdin.close()
|
|
except (IOError, BrokenPipeError):
|
|
# This can happen if the process exits quickly or doesn't read stdin.
|
|
logger.debug(f"Could not write to stdin for command: {' '.join(cmd)}. Process may have already exited.")
|
|
|
|
# Wait for the process to finish and for all output to be read.
|
|
# Add a timeout to prevent indefinite hangs. 15 minutes should be enough for any single download.
|
|
timeout_seconds = 15 * 60
|
|
try:
|
|
retcode = process.wait(timeout=timeout_seconds)
|
|
except subprocess.TimeoutExpired:
|
|
logger.error(f"Command timed out after {timeout_seconds} seconds: {' '.join(cmd)}")
|
|
# Kill the entire process group to ensure child processes (like yt-dlp or ffmpeg) are also terminated.
|
|
try:
|
|
os.killpg(os.getpgid(process.pid), signal.SIGKILL)
|
|
except (ProcessLookupError, PermissionError):
|
|
pass # Process already finished or we lack permissions
|
|
retcode = -1 # Indicate failure
|
|
# Wait a moment for pipes to close after killing.
|
|
try:
|
|
process.wait(timeout=5)
|
|
except subprocess.TimeoutExpired:
|
|
logger.warning("Process did not terminate gracefully after SIGKILL.")
|
|
|
|
stdout_thread.join(timeout=5)
|
|
stderr_thread.join(timeout=5)
|
|
|
|
stdout_bytes = b"".join(stdout_capture)
|
|
stderr_bytes = b"".join(stderr_capture)
|
|
|
|
# If we timed out, create a synthetic stderr message to ensure the failure is reported upstream.
|
|
if retcode == -1 and not stderr_bytes.strip():
|
|
stderr_bytes = f"Command timed out after {timeout_seconds} seconds".encode('utf-8')
|
|
|
|
stdout = stdout_bytes if binary_stdout else stdout_bytes.decode('utf-8', errors='replace')
|
|
stderr = stderr_bytes.decode('utf-8', errors='replace')
|
|
|
|
return retcode, stdout, stderr
|
|
|
|
except FileNotFoundError:
|
|
logger.error(f"Command not found: {cmd[0]}. Make sure it's in your PATH.")
|
|
return -1, "", f"Command not found: {cmd[0]}"
|
|
except Exception as e:
|
|
logger.error(f"An error occurred while running command: {' '.join(cmd)}. Error: {e}")
|
|
return -1, "", str(e)
|
|
finally:
|
|
if process:
|
|
with process_lock:
|
|
running_processes.discard(process)
|
|
|
|
|
|
def run_docker_container(image_name, command, volumes, stream_prefix="", network_name=None, log_callback=None, profile_manager=None, profile_name=None, environment=None, log_command_override=None):
|
|
"""
|
|
Runs a command in a new, ephemeral Docker container using docker-py.
|
|
Streams logs in real-time, allows for live log processing, and ensures cleanup.
|
|
Can monitor a profile and stop the container if the profile is BANNED or RESTING.
|
|
Returns a tuple of (exit_code, stdout_str, stderr_str, stop_reason).
|
|
"""
|
|
if not docker:
|
|
# This should be caught earlier, but as a safeguard:
|
|
return -1, "", "Docker SDK for Python is not installed. Please run: pip install docker", None
|
|
|
|
logger.debug(f"Running docker container. Image: {image_name}, Command: {command}, Volumes: {volumes}, Network: {network_name}")
|
|
|
|
# --- Construct and log the equivalent CLI command for debugging ---
|
|
try:
|
|
user_id = f"{os.getuid()}:{os.getgid()}" if os.name != 'nt' else None
|
|
cli_cmd = ['docker', 'run', '--rm']
|
|
if user_id:
|
|
cli_cmd.extend(['-u', user_id])
|
|
if network_name:
|
|
cli_cmd.extend(['--network', network_name])
|
|
if environment:
|
|
for k, v in sorted(environment.items()):
|
|
cli_cmd.extend(['-e', f"{k}={v}"])
|
|
if volumes:
|
|
for host_path, container_config in sorted(volumes.items()):
|
|
bind = container_config.get('bind')
|
|
mode = container_config.get('mode', 'rw')
|
|
cli_cmd.extend(['-v', f"{os.path.abspath(host_path)}:{bind}:{mode}"])
|
|
cli_cmd.append(image_name)
|
|
cli_cmd.extend(command)
|
|
logger.info(f"Full docker command: {' '.join(shlex.quote(s) for s in cli_cmd)}")
|
|
if log_command_override:
|
|
# Build a more comprehensive, runnable command for logging
|
|
env_prefix_parts = []
|
|
if environment:
|
|
for k, v in sorted(environment.items()):
|
|
env_prefix_parts.append(f"{k}={shlex.quote(str(v))}")
|
|
|
|
env_prefix = ' '.join(env_prefix_parts)
|
|
|
|
equivalent_ytdlp_cmd = ' '.join(shlex.quote(s) for s in log_command_override)
|
|
|
|
full_equivalent_cmd = f"{env_prefix} {equivalent_ytdlp_cmd}".strip()
|
|
logger.info(f"Equivalent host command: {full_equivalent_cmd}")
|
|
except Exception as e:
|
|
logger.warning(f"Could not construct equivalent docker command for logging: {e}")
|
|
# --- End of logging ---
|
|
|
|
container = None
|
|
monitor_thread = None
|
|
stop_monitor_event = threading.Event()
|
|
# Use a mutable object (dict) to share the stop reason between threads
|
|
stop_reason_obj = {'reason': None}
|
|
try:
|
|
client = docker.from_env()
|
|
|
|
# Run container as current host user to avoid permission issues with volume mounts
|
|
user_id = f"{os.getuid()}:{os.getgid()}" if os.name != 'nt' else None
|
|
|
|
container = client.containers.run(
|
|
image_name,
|
|
command=command,
|
|
volumes=volumes,
|
|
detach=True,
|
|
network=network_name,
|
|
user=user_id,
|
|
environment=environment,
|
|
# We use `remove` in `finally` instead of `auto_remove` to ensure we can get logs
|
|
# even if the container fails to start.
|
|
)
|
|
|
|
# Thread to monitor profile status and stop container if BANNED or RESTING
|
|
def monitor_profile():
|
|
while not stop_monitor_event.is_set():
|
|
try:
|
|
profile_info = profile_manager.get_profile(profile_name)
|
|
if profile_info:
|
|
state = profile_info.get('state')
|
|
if state in ['BANNED', 'RESTING']:
|
|
logger.warning(f"Profile '{profile_name}' is {state}. Stopping container {container.short_id}.")
|
|
stop_reason_obj['reason'] = f"Profile became {state}"
|
|
try:
|
|
container.stop(timeout=5)
|
|
except docker.errors.APIError as e:
|
|
logger.warning(f"Could not stop container {container.short_id}: {e}")
|
|
break # Stop monitoring
|
|
except Exception as e:
|
|
logger.error(f"Error in profile monitor thread: {e}")
|
|
|
|
# Wait for 2 seconds or until stop event is set
|
|
stop_monitor_event.wait(2)
|
|
|
|
if profile_manager and profile_name:
|
|
monitor_thread = threading.Thread(target=monitor_profile, daemon=True)
|
|
monitor_thread.start()
|
|
|
|
# Stream logs in a separate thread to avoid blocking.
|
|
log_stream = container.logs(stream=True, follow=True, stdout=True, stderr=True)
|
|
for line_bytes in log_stream:
|
|
line_str = line_bytes.decode('utf-8', errors='replace').strip()
|
|
# Use logger.info to ensure output is captured by all handlers
|
|
logger.info(f"{stream_prefix}{line_str}")
|
|
if log_callback:
|
|
# The callback can return True to signal an immediate stop.
|
|
if log_callback(line_str):
|
|
logger.warning(f"Log callback requested to stop container {container.short_id}.")
|
|
stop_reason_obj['reason'] = "Stopped by log callback (fatal error)"
|
|
try:
|
|
container.stop(timeout=5)
|
|
except docker.errors.APIError as e:
|
|
logger.warning(f"Could not stop container {container.short_id}: {e}")
|
|
break # Stop reading logs
|
|
|
|
result = container.wait(timeout=15 * 60)
|
|
exit_code = result.get('StatusCode', -1)
|
|
|
|
# Get final logs to separate stdout and stderr.
|
|
final_stdout = container.logs(stdout=True, stderr=False)
|
|
final_stderr = container.logs(stdout=False, stderr=True)
|
|
|
|
stdout_str = final_stdout.decode('utf-8', errors='replace')
|
|
stderr_str = final_stderr.decode('utf-8', errors='replace')
|
|
|
|
return exit_code, stdout_str, stderr_str, stop_reason_obj['reason']
|
|
|
|
except docker.errors.ImageNotFound:
|
|
logger.error(f"Docker image not found: '{image_name}'. Please pull it first.")
|
|
return -1, "", f"Docker image not found: {image_name}", None
|
|
except docker.errors.APIError as e:
|
|
logger.error(f"Docker API error: {e}")
|
|
return -1, "", str(e), None
|
|
except Exception as e:
|
|
logger.error(f"An unexpected error occurred while running docker container: {e}", exc_info=True)
|
|
return -1, "", str(e), None
|
|
finally:
|
|
if monitor_thread:
|
|
stop_monitor_event.set()
|
|
monitor_thread.join(timeout=1)
|
|
if container:
|
|
try:
|
|
container.remove(force=True)
|
|
logger.debug(f"Removed container {container.short_id}")
|
|
except docker.errors.APIError as e:
|
|
logger.warning(f"Could not remove container {container.short_id}: {e}")
|