#!/usr/bin/env python3 import argparse import concurrent.futures import os import random import re import signal import subprocess import sys import tempfile import threading import time from dataclasses import dataclass, field from datetime import datetime, timezone from typing import Any, IO, List, Optional, Protocol, Sequence, Tuple from urllib.parse import urlsplit import boto3 import botocore from boto3.session import Session from botocore.config import Config DUMPLING_PATH = "/dumpling" CSE_CTL_PATH = "/cse-ctl" MAX_ATTEMPTS = 4 RETRY_BACKOFF_SECONDS = (5, 10, 20) DUMPLING_TERMINATE_GRACE_SECONDS = 10 EXIT_LOG_DRAIN_SECONDS = 20 COMPLETE_MARKER_NAME = "complete" COMPLETE_MARKER_JOB_INSTANCE_METADATA = "job-instance-name" COMPLETE_MARKER_BACKUP_TIME_METADATA = "backup-time" COMPLETE_MARKER_BACKUP_TIME_FORMAT = "%Y-%m-%dT%H:%M:%SZ" RUN_ID_REGEXP = re.compile(r"[A-Za-z0-9][A-Za-z0-9._-]*") JOB_INSTANCE_NAME_ENV = "DUMPLING_PACKED_BACKUP_JOB_INSTANCE_NAME" AWS_MAX_ATTEMPTS = 10 LOG_LOCK = threading.Lock() SCRIPT_START_UNIX_NANO = time.time_ns() class ExportError(Exception): pass class DumplingFailed(Exception): def __init__(self, returncode: int) -> None: super().__init__(f"dumpling exited with status {returncode}") self.returncode = returncode class DumplingTimedOut(Exception): def __init__(self, process_id: int, timeout_seconds: int) -> None: super().__init__( f"dumpling pid {process_id} exceeded {timeout_seconds}s timeout" ) self.process_id = process_id self.timeout_seconds = timeout_seconds class ReceivedSignal(Exception): def __init__(self, signum: int) -> None: super().__init__(signal.Signals(signum).name) self.signum = signum class DumplingProcess(Protocol): pid: int def poll(self) -> Optional[int]: ... def wait(self, timeout: Optional[float] = None) -> int: ... class StoreOnce(argparse.Action): def __call__( self, parser: argparse.ArgumentParser, namespace: argparse.Namespace, values: object, option_string: Optional[str] = None, ) -> None: marker = f"_seen_{self.dest}" if getattr(namespace, marker, False): raise argparse.ArgumentError( self, f"{option_string or self.dest} may only be specified once" ) setattr(namespace, marker, True) setattr(namespace, self.dest, values) @dataclass(frozen=True) class SourcePrefix: bucket: str object_prefix: str backup_time: datetime @dataclass(frozen=True) class OutputPrefix: bucket: str object_prefix: str @dataclass class Progress: completed: int = 0 skipped: int = 0 total: Optional[int] = None lock: threading.Lock = field(default_factory=threading.Lock, repr=False) def set_total(self, total: int) -> None: with self.lock: self.total = total def mark_completed(self) -> int: with self.lock: self.completed += 1 return self.completed def mark_skipped(self) -> int: with self.lock: self.skipped += 1 return self.skipped def snapshot(self) -> Tuple[int, int, Optional[int]]: with self.lock: return self.completed, self.skipped, self.total @dataclass(frozen=True) class ExportItem: source: str target: str output: str command: Sequence[str] @dataclass(frozen=True) class CompleteMarker: uri: str job_instance_name: Optional[str] backup_time: Optional[datetime] etag: str def emit(message: str, *, error: bool = False) -> None: with LOG_LOCK: print(message, file=sys.stderr if error else sys.stdout, flush=True) def flush_log_streams() -> None: with LOG_LOCK: sys.stdout.flush() sys.stderr.flush() def wait_for_log_collection() -> None: emit( f"log drain: flushing streams and waiting " f"{EXIT_LOG_DRAIN_SECONDS}s before exit" ) flush_log_streams() try: time.sleep(EXIT_LOG_DRAIN_SECONDS) except (ReceivedSignal, KeyboardInterrupt): pass finally: flush_log_streams() def emit_progress_summary(progress: Progress, *, error: bool = False) -> None: completed, skipped, total = progress.snapshot() total_text = "unknown" if total is None else str(total) emit( f"export progress: completed={completed} skipped={skipped} " f"total={total_text}", error=error, ) def parse_non_negative_integer(value: str) -> int: if not value.isascii() or not value.isdigit(): raise argparse.ArgumentTypeError("must be a non-negative integer") return int(value) def parse_positive_integer(value: str) -> int: parsed = parse_non_negative_integer(value) if parsed == 0: raise argparse.ArgumentTypeError("must be a positive integer") return parsed def backup_time_from_name(backup_name: str) -> datetime: try: parsed = datetime.strptime(backup_name, "%Y%m%d/%H%M%S.meta") except ValueError as exc: raise ExportError( f"backup_name does not contain an expected UTC timestamp: " f"{backup_name!r}" ) from exc return parsed.replace(tzinfo=timezone.utc) def format_complete_marker_backup_time(backup_time: datetime) -> str: return backup_time.astimezone(timezone.utc).strftime( COMPLETE_MARKER_BACKUP_TIME_FORMAT ) def validate_run_id(run_id: str, description: str) -> str: if RUN_ID_REGEXP.fullmatch(run_id) is None: raise ExportError( f"{description} must contain only letters, digits, '.', '_', or '-'" ) return run_id def output_version(backup_time: datetime, run_id: str) -> str: validate_run_id(run_id, "run ID") return f"{format_complete_marker_backup_time(backup_time)}:{run_id}" def parse_output_version(value: str) -> Optional[Tuple[datetime, str]]: timestamp_length = len("2000-01-01T00:00:00Z") if len(value) <= timestamp_length or value[timestamp_length] != ":": return None timestamp = value[:timestamp_length] run_id = value[timestamp_length + 1 :] if RUN_ID_REGEXP.fullmatch(run_id) is None: return None try: backup_time = datetime.strptime( timestamp, COMPLETE_MARKER_BACKUP_TIME_FORMAT, ).replace(tzinfo=timezone.utc) except ValueError: return None return backup_time, run_id def complete_marker_body(backup_time: datetime, run_id: str) -> bytes: return output_version(backup_time, run_id).encode("ascii") def parse_complete_marker_backup_time( value: Any, marker_uri: str, ) -> Optional[datetime]: if value is None: return None if not isinstance(value, str): raise ExportError( f"complete marker has invalid " f"{COMPLETE_MARKER_BACKUP_TIME_METADATA!r} metadata: {marker_uri}" ) try: parsed = datetime.strptime( value, COMPLETE_MARKER_BACKUP_TIME_FORMAT, ) except ValueError as exc: raise ExportError( f"complete marker has invalid " f"{COMPLETE_MARKER_BACKUP_TIME_METADATA!r} metadata: {marker_uri}" ) from exc return parsed.replace(tzinfo=timezone.utc) def parse_regexp(value: str) -> re.Pattern: try: return re.compile(value) except re.error as exc: raise argparse.ArgumentTypeError(f"invalid regular expression: {exc}") from exc def build_parser() -> argparse.ArgumentParser: parser = argparse.ArgumentParser( prog="dumpling_packed_backup_init.py", usage=( "%(prog)s --s3-prefix S3_PREFIX --out-prefix OUT_PREFIX " "[options] [-- DUMPLING_ARGS...]" ), description="Export CSE packed-backup metadata entries with dumpling.", epilog="Arguments after -- are appended unchanged to every dumpling command.", ) parser.add_argument( "--s3-prefix", required=True, action=StoreOnce, help=( "s3:////packed_backup/native/" "//meta" ), ) parser.add_argument( "--out-prefix", required=True, action=StoreOnce, help="s3:/// used as the dumpling output root", ) parser.add_argument( "--exclude-meta-name-regexp", default=None, type=parse_regexp, action=StoreOnce, metavar="REGEXP", help=( "exclude backups whose meta filename (including .meta) matches " "REGEXP; applied before task shuffling" ), ) parser.add_argument( "--parallel", default=1, type=parse_positive_integer, action=StoreOnce, help="number of concurrent dumpling processes (default: 1)", ) parser.add_argument( "--dumpling-timeout", default=None, type=parse_positive_integer, action=StoreOnce, metavar="SECONDS", help=( "skip one dump when its dumpling process exceeds SECONDS; " "disabled by default" ), ) parser.add_argument( "--gpg", required=True, action=StoreOnce, metavar="S3_URI_OR_PATH", help=( "S3 URI or local path of an OpenPGP public key used to encrypt " "dumpling output" ), ) parser.add_argument( "--flat", action="store_true", help="omit the run timestamp and write to out-prefix//", ) return parser def parse_args(argv: Optional[Sequence[str]] = None) -> argparse.Namespace: raw_args = list(sys.argv[1:] if argv is None else argv) if "--" in raw_args: separator = raw_args.index("--") starter_args = raw_args[:separator] dumpling_args = raw_args[separator + 1 :] else: starter_args = raw_args dumpling_args = [] args = build_parser().parse_args(starter_args) args.dumpling_args = dumpling_args return args def parse_s3_prefix(value: str) -> SourcePrefix: parsed = urlsplit(value) if parsed.scheme != "s3": raise ExportError("s3-prefix must start with s3://") if not parsed.netloc: raise ExportError("s3-prefix must include a bucket") if parsed.query or parsed.fragment: raise ExportError("s3-prefix must not contain a query or fragment") object_prefix = parsed.path.lstrip("/").rstrip("/") if not object_prefix: raise ExportError("s3-prefix must include an object prefix") if any(not component for component in object_prefix.split("/")): raise ExportError("s3-prefix must not contain empty path components") root_marker = "packed_backup/native/" nested_marker = f"/{root_marker}" if object_prefix.startswith(root_marker): packed_prefix = object_prefix[len(root_marker) :] elif nested_marker in object_prefix: _, packed_prefix = object_prefix.split(nested_marker, 1) else: raise ExportError( "s3-prefix must match s3:////packed_backup/native/" "//meta" ) if nested_marker in packed_prefix: raise ExportError("s3-prefix contains an ambiguous packed_backup/native path") if not packed_prefix.endswith("/meta"): raise ExportError("s3-prefix must end with /meta") meta_parent = packed_prefix[: -len("/meta")] if "/" not in meta_parent: raise ExportError("s3-prefix does not contain both backup_name and batch_id") backup_name, batch_id = meta_parent.rsplit("/", 1) if not backup_name or not batch_id: raise ExportError("s3-prefix contains an empty backup_name or batch_id") if not backup_name.endswith(".meta"): raise ExportError("backup_name parsed from s3-prefix must end with .meta") backup_time = backup_time_from_name(backup_name) return SourcePrefix( bucket=parsed.netloc, object_prefix=object_prefix, backup_time=backup_time, ) def parse_output_prefix(value: str) -> OutputPrefix: parsed = urlsplit(value) if parsed.scheme != "s3": raise ExportError("out-prefix must start with s3://") if not parsed.netloc: raise ExportError("out-prefix must include a bucket") if parsed.query or parsed.fragment: raise ExportError("out-prefix must not contain a query or fragment") object_prefix = parsed.path.lstrip("/").rstrip("/") if not object_prefix: raise ExportError("out-prefix must include an object prefix") if any(not component for component in object_prefix.split("/")): raise ExportError("out-prefix must not contain empty path components") return OutputPrefix(bucket=parsed.netloc, object_prefix=object_prefix) def parse_s3_object_uri(value: str, option_name: str) -> Tuple[str, str]: parsed = urlsplit(value) if parsed.scheme != "s3" or not parsed.netloc: raise ExportError(f"{option_name} must be an s3:// URI") if parsed.query or parsed.fragment: raise ExportError(f"{option_name} must not contain a query or fragment") object_key = parsed.path.lstrip("/") if not object_key or object_key.endswith("/"): raise ExportError(f"{option_name} must identify an S3 object") return parsed.netloc, object_key def validate_gpg_source(value: str) -> bool: if value.startswith("s3://"): parse_s3_object_uri(value, "--gpg") return True if not os.path.isfile(value): raise ExportError(f"--gpg local path is not a regular file: {value}") if not os.access(value, os.R_OK): raise ExportError(f"--gpg local file is not readable: {value}") return False def parse_meta_key(prefix: SourcePrefix, object_key: str) -> str: key_prefix = f"{prefix.object_prefix}/" if not object_key.startswith(key_prefix): raise ExportError( f"cannot parse meta key s3://{prefix.bucket}/{object_key}" ) relative_key = object_key[len(key_prefix) :] meta_filename = relative_key.rsplit("/", 1)[-1] if not meta_filename.endswith(".meta"): raise ExportError( "meta filename must end with .meta: " f"s3://{prefix.bucket}/{object_key}" ) meta_identity = meta_filename[: -len(".meta")] if not meta_identity: raise ExportError( "meta filename must contain a name before .meta: " f"s3://{prefix.bucket}/{object_key}" ) return meta_identity def target_uri( prefix: OutputPrefix, script_start_unix_nano: int, meta_identity: str, flat: bool = False, ) -> str: root = f"s3://{prefix.bucket}/{prefix.object_prefix}/" if flat: return f"{root}{meta_identity}/" return f"{root}{script_start_unix_nano}/{meta_identity}/" def dumpling_output_uri( target: str, backup_time: datetime, run_id: str, ) -> str: return f"{target.rstrip('/')}/{output_version(backup_time, run_id)}/" def require_executables() -> None: if not os.path.isfile(DUMPLING_PATH) or not os.access(DUMPLING_PATH, os.X_OK): raise ExportError(f"{DUMPLING_PATH} is not executable") if not os.path.isfile(CSE_CTL_PATH) or not os.access(CSE_CTL_PATH, os.X_OK): raise ExportError(f"{CSE_CTL_PATH} is not executable") def require_job_instance_name() -> str: job_instance_name = os.environ.get(JOB_INSTANCE_NAME_ENV, "").strip() if not job_instance_name: raise ExportError(f"{JOB_INSTANCE_NAME_ENV} is required") return job_instance_name def create_s3_client( region: str, parallel: int, ) -> Tuple[Any, str, str]: # Share one client across all export threads so credentials and # connections are reused. try: session = Session() client = session.client( "s3", region_name=region, config=Config( max_pool_connections=parallel, retries={"total_max_attempts": AWS_MAX_ATTEMPTS}, ), ) except Exception as exc: raise ExportError( f"cannot initialize S3 client: {aws_error_message(exc)}" ) from exc return ( client, getattr(boto3, "__version__", "unknown"), getattr(botocore, "__version__", "unknown"), ) def aws_error_message(exc: Exception) -> str: detail = " ".join(str(exc).splitlines()) return detail or type(exc).__name__ def aws_error_code(exc: Exception) -> Tuple[Optional[str], Optional[int]]: response = getattr(exc, "response", None) if not isinstance(response, dict): return None, None error = response.get("Error") metadata = response.get("ResponseMetadata") code = error.get("Code") if isinstance(error, dict) else None status = ( metadata.get("HTTPStatusCode") if isinstance(metadata, dict) else None ) return ( str(code) if code is not None else None, status if isinstance(status, int) else None, ) def list_meta_keys(s3_client: Any, prefix: SourcePrefix) -> List[str]: object_keys = [] try: paginator = s3_client.get_paginator("list_objects_v2") pages = paginator.paginate( Bucket=prefix.bucket, Prefix=f"{prefix.object_prefix}/", ) for page in pages: if not isinstance(page, dict): raise ExportError("S3 returned an unexpected object page") contents = page.get("Contents", []) if contents is None: contents = [] if not isinstance(contents, list): raise ExportError("S3 returned unexpected object contents") for item in contents: if not isinstance(item, dict): raise ExportError("S3 returned an unexpected object") object_key = item.get("Key") if not isinstance(object_key, str): raise ExportError("S3 returned an object without a key") if object_key.endswith(".meta"): object_keys.append(object_key) except ExportError: raise except Exception as exc: raise ExportError( f"cannot list meta objects: {aws_error_message(exc)}" ) from exc if len(object_keys) != len(set(object_keys)): raise ExportError("S3 returned duplicate meta keys") return sorted(object_keys) def download_gpg_key(s3_client: Any, s3_uri: str, destination: str) -> None: bucket, object_key = parse_s3_object_uri(s3_uri, "--gpg") body = None try: response = s3_client.get_object(Bucket=bucket, Key=object_key) body = response.get("Body") if body is None or not callable(getattr(body, "read", None)): raise ExportError("S3 returned a GPG object without a body") with open(destination, "wb") as output: while True: chunk = body.read(1024 * 1024) if not chunk: break output.write(chunk) except ExportError: raise except Exception as exc: raise ExportError( f"cannot download --gpg {s3_uri}: {aws_error_message(exc)}" ) from exc finally: if body is not None: close = getattr(body, "close", None) if callable(close): close() os.chmod(destination, 0o600) def complete_marker_location(target: str) -> Tuple[str, str, str]: marker_uri = f"{target.rstrip('/')}/{COMPLETE_MARKER_NAME}" bucket, object_key = parse_s3_object_uri(marker_uri, "complete marker") return marker_uri, bucket, object_key def existing_complete_marker( s3_client: Any, target: str, ) -> Optional[CompleteMarker]: marker_uri, bucket, object_key = complete_marker_location(target) try: response = s3_client.head_object(Bucket=bucket, Key=object_key) except Exception as exc: error_code, http_status = aws_error_code(exc) if ( error_code in {"404", "NoSuchKey", "NotFound"} or http_status == 404 ): return None raise ExportError( f"cannot check complete marker {marker_uri}: " f"{aws_error_message(exc)}" ) from exc if not isinstance(response, dict): raise ExportError( f"S3 returned an unexpected complete marker response for {marker_uri}" ) metadata = response.get("Metadata", {}) if metadata is None: metadata = {} if not isinstance(metadata, dict): raise ExportError( f"S3 returned unexpected complete marker metadata for {marker_uri}" ) job_instance_name = metadata.get(COMPLETE_MARKER_JOB_INSTANCE_METADATA) backup_time = parse_complete_marker_backup_time( metadata.get(COMPLETE_MARKER_BACKUP_TIME_METADATA), marker_uri, ) etag = response.get("ETag") if not isinstance(etag, str) or not etag: raise ExportError( f"S3 returned a complete marker without an ETag for {marker_uri}" ) return CompleteMarker( uri=marker_uri, job_instance_name=( job_instance_name if isinstance(job_instance_name, str) else None ), backup_time=backup_time, etag=etag, ) def write_complete_marker( s3_client: Any, target: str, job_instance_name: str, backup_time: datetime, original_marker: Optional[CompleteMarker], ) -> str: marker_uri, bucket, object_key = complete_marker_location(target) if original_marker is not None and original_marker.uri != marker_uri: raise ExportError( "cannot conditionally update a complete marker using state read " f"from another object: {original_marker.uri}" ) request = { "Bucket": bucket, "Key": object_key, "Body": complete_marker_body(backup_time, job_instance_name), "Metadata": { COMPLETE_MARKER_JOB_INSTANCE_METADATA: job_instance_name, COMPLETE_MARKER_BACKUP_TIME_METADATA: ( format_complete_marker_backup_time(backup_time) ), }, } if original_marker is None: request["IfNoneMatch"] = "*" else: request["IfMatch"] = original_marker.etag try: s3_client.put_object(**request) except Exception as exc: error_code, http_status = aws_error_code(exc) if ( error_code in {"409", "412", "ConditionalRequestConflict", "PreconditionFailed"} or http_status in {409, 412} ): raise ExportError( f"complete marker changed since it was first read; " f"candidate output was retained: {marker_uri}" ) from exc raise ExportError( f"cannot write complete marker {marker_uri}: " f"{aws_error_message(exc)}" ) from exc return marker_uri def list_target_keys( s3_client: Any, target: str, ) -> Tuple[str, str, List[str]]: _, bucket, marker_key = complete_marker_location(target) key_prefix = marker_key[: -len(COMPLETE_MARKER_NAME)] object_keys = [] try: paginator = s3_client.get_paginator("list_objects_v2") pages = paginator.paginate(Bucket=bucket, Prefix=key_prefix) for page in pages: if not isinstance(page, dict): raise ExportError("S3 returned an unexpected target object page") contents = page.get("Contents", []) if contents is None: contents = [] if not isinstance(contents, list): raise ExportError("S3 returned unexpected target object contents") for item in contents: if not isinstance(item, dict): raise ExportError("S3 returned an unexpected target object") object_key = item.get("Key") if not isinstance(object_key, str): raise ExportError( "S3 returned a target object without a key" ) object_keys.append(object_key) except ExportError: raise except Exception as exc: raise ExportError( f"cannot list export target {target}: {aws_error_message(exc)}" ) from exc if len(object_keys) != len(set(object_keys)): raise ExportError(f"S3 returned duplicate keys under export target {target}") return bucket, marker_key, sorted(object_keys) def stale_export_object_keys( marker_key: str, object_keys: Sequence[str], backup_time: datetime, ) -> List[str]: key_prefix = marker_key[: -len(COMPLETE_MARKER_NAME)] stale_keys = [] for object_key in object_keys: if not object_key.startswith(key_prefix): continue relative_key = object_key[len(key_prefix) :] if relative_key == COMPLETE_MARKER_NAME: continue version_component, separator, _ = relative_key.partition("/") parsed_version = ( parse_output_version(version_component) if separator else None ) if parsed_version is not None: object_backup_time, _ = parsed_version if object_backup_time < backup_time: stale_keys.append(object_key) continue # Before versioned output was introduced, Dumpling wrote encrypted # files directly beside the marker. Delete only that known legacy # shape; preserve unknown objects for manual inspection. if ( relative_key and "/" not in relative_key and relative_key.endswith(".gpg") ): stale_keys.append(object_key) return sorted(stale_keys) def remove_stale_export_objects( s3_client: Any, target: str, backup_time: datetime, delete_executor: concurrent.futures.Executor, ) -> None: bucket, marker_key, object_keys = list_target_keys(s3_client, target) stale_keys = stale_export_object_keys( marker_key, object_keys, backup_time, ) emit( f"export target cleanup: stale_objects={len(stale_keys)} " f"retained_objects={len(object_keys) - len(stale_keys)} " f"before_backup_time={format_complete_marker_backup_time(backup_time)} " f"target={target}" ) futures = { delete_executor.submit( s3_client.delete_object, Bucket=bucket, Key=object_key, ): object_key for object_key in stale_keys } failures = [] for future in concurrent.futures.as_completed(futures): object_key = futures[future] try: future.result() except Exception as exc: failures.append((object_key, exc)) if failures: object_key, failure = min(failures, key=lambda entry: entry[0]) raise ExportError( f"cannot delete export target object s3://{bucket}/{object_key}: " f"{aws_error_message(failure)}" ) from failure emit(f"export target cleaned: objects={len(stale_keys)} target={target}") def dumpling_command( source_uri: str, target: str, region: str, gpg_key_file: str, extra_args: Sequence[str], ) -> List[str]: command = [ DUMPLING_PATH, "--packed-backup", source_uri, "--cse-ctl-path", CSE_CTL_PATH, "--filter", "*.*", "--filter", "!mysql.*", "--filter", "!sys.*", "--filetype", "csv", "--output", target, "--s3.region", region, "--status-addr", "127.0.0.1:0", ] command.extend(["--gpg-key-file", gpg_key_file]) command.extend(extra_args) return command def forward_dumpling_stream( stream: IO[str], process_id: int, *, error: bool ) -> None: with stream: for raw_line in stream: line = raw_line.rstrip("\r\n") emit( f"[dumpling pid={process_id}] {line}", error=error, ) def signal_dumpling_process_group( process: DumplingProcess, signum: signal.Signals, ) -> None: if process.poll() is not None: return try: os.killpg(process.pid, signum) except ProcessLookupError: pass def stop_timed_out_dumpling( process: DumplingProcess, timeout_seconds: int, ) -> int: emit( f"dumpling timeout: pid={process.pid} timeout={timeout_seconds}s " f"signal={signal.SIGTERM.name}", error=True, ) signal_dumpling_process_group(process, signal.SIGTERM) try: return process.wait(timeout=DUMPLING_TERMINATE_GRACE_SECONDS) except subprocess.TimeoutExpired: emit( f"dumpling timeout: pid={process.pid} did not exit within " f"{DUMPLING_TERMINATE_GRACE_SECONDS}s; " f"signal={signal.SIGKILL.name}", error=True, ) signal_dumpling_process_group(process, signal.SIGKILL) return process.wait() def run_dumpling_process( command: Sequence[str], total: int, attempt: int, source: str, target: str, timeout_seconds: Optional[int], ) -> Tuple[int, int]: process = subprocess.Popen( list(command), stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, encoding="utf-8", errors="replace", bufsize=1, start_new_session=True, ) process_id = process.pid emit( f"export attempt: total={total} attempt={attempt}/{MAX_ATTEMPTS} " f"pid={process_id} " f"source={source} target={target}" ) assert process.stdout is not None assert process.stderr is not None with concurrent.futures.ThreadPoolExecutor(max_workers=2) as log_executor: stdout_future = log_executor.submit( forward_dumpling_stream, process.stdout, process_id, error=False, ) stderr_future = log_executor.submit( forward_dumpling_stream, process.stderr, process_id, error=True, ) timed_out = False try: returncode = process.wait(timeout=timeout_seconds) except subprocess.TimeoutExpired: assert timeout_seconds is not None timed_out = True returncode = stop_timed_out_dumpling(process, timeout_seconds) stdout_future.result() stderr_future.result() if timed_out: assert timeout_seconds is not None raise DumplingTimedOut(process_id, timeout_seconds) return returncode, process_id def run_dumpling( command: Sequence[str], total: int, source: str, target: str, output: str, s3_client: Any, progress: Progress, job_instance_name: str, backup_time: datetime, dumpling_timeout: Optional[int], delete_executor: concurrent.futures.Executor, ) -> Optional[int]: complete_marker = existing_complete_marker(s3_client, target) if complete_marker is not None: if ( complete_marker.backup_time is not None and complete_marker.backup_time > backup_time ): completed = progress.mark_completed() emit( f"export skipped for newer backup: completed={completed} " f"total={total} marker={complete_marker.uri} " f"marker_backup_time=" f"{format_complete_marker_backup_time(complete_marker.backup_time)} " f"current_backup_time=" f"{format_complete_marker_backup_time(backup_time)} " f"source={source} target={target}" ) return None if complete_marker.job_instance_name == job_instance_name: completed = progress.mark_completed() emit( f"export skipped: completed={completed} total={total} " f"marker={complete_marker.uri} " f"source={source} target={target}" ) return None emit( f"complete marker mismatch: marker={complete_marker.uri} " f"expected_job_instance={job_instance_name!r} " f"actual_job_instance={complete_marker.job_instance_name!r}" ) for attempt in range(1, MAX_ATTEMPTS + 1): try: returncode, process_id = run_dumpling_process( command, total, attempt, source, output, dumpling_timeout, ) except DumplingTimedOut as exc: skipped = progress.mark_skipped() completed, _, _ = progress.snapshot() emit( f"export skipped after dumpling timeout: " f"completed={completed} skipped={skipped} total={total} " f"pid={exc.process_id} timeout={exc.timeout_seconds}s " f"source={source} target={target} output={output}", error=True, ) return None if returncode == 0: marker_uri = write_complete_marker( s3_client, target, job_instance_name, backup_time, complete_marker, ) remove_stale_export_objects( s3_client, target, backup_time, delete_executor, ) completed = progress.mark_completed() emit( f"export marked complete: completed={completed} total={total} " f"marker={marker_uri} output={output}" ) return process_id if attempt == MAX_ATTEMPTS: emit( f"dumpling failed: total={total} attempt={attempt}/{MAX_ATTEMPTS} " f"pid={process_id} " f"exit_status={returncode}; retries exhausted", error=True, ) raise DumplingFailed(returncode) backoff_seconds = RETRY_BACKOFF_SECONDS[attempt - 1] emit( f"dumpling failed: total={total} attempt={attempt}/{MAX_ATTEMPTS} " f"pid={process_id} " f"exit_status={returncode}; retrying in {backoff_seconds}s", error=True, ) time.sleep(backoff_seconds) raise AssertionError("unreachable") def run_export_items( items: Sequence[ExportItem], parallel: int, total: int, s3_client: Any, progress: Progress, job_instance_name: str, backup_time: datetime, dumpling_timeout: Optional[int], ) -> None: next_item = 0 failures = [] with ( concurrent.futures.ThreadPoolExecutor(max_workers=parallel) as executor, concurrent.futures.ThreadPoolExecutor( max_workers=parallel ) as delete_executor, ): futures = {} def fill_workers() -> None: nonlocal next_item while ( not failures and next_item < len(items) and len(futures) < parallel ): item = items[next_item] next_item += 1 future = executor.submit( run_dumpling, item.command, total, item.source, item.target, item.output, s3_client, progress, job_instance_name, backup_time, dumpling_timeout, delete_executor, ) futures[future] = item fill_workers() while futures: done, _ = concurrent.futures.wait( tuple(futures), return_when=concurrent.futures.FIRST_COMPLETED ) for future in sorted(done, key=lambda entry: futures[entry].source): item = futures.pop(future) try: process_id = future.result() except Exception as exc: failures.append((item.source, exc)) else: if process_id is None: continue emit( f"export completed: total={total} pid={process_id} " f"source={item.source} target={item.target}" ) fill_workers() if failures: _, failure = min(failures, key=lambda entry: entry[0]) raise failure def run_export(args: argparse.Namespace, progress: Progress) -> None: source_prefix = parse_s3_prefix(args.s3_prefix) output_prefix = parse_output_prefix(args.out_prefix) gpg_is_s3 = validate_gpg_source(args.gpg) job_instance_name = require_job_instance_name() region = ( os.environ.get("AWS_REGION") or os.environ.get("AWS_DEFAULT_REGION") or os.environ.get("DFS_S3_REGION") ) if not region: raise ExportError( "AWS_REGION, AWS_DEFAULT_REGION, or DFS_S3_REGION is required" ) require_executables() dumpling_timeout_text = ( "disabled" if args.dumpling_timeout is None else f"{args.dumpling_timeout}s" ) emit( f"export start: parallel={args.parallel} region={region} " f"dumpling_timeout={dumpling_timeout_text} " f"job_instance={job_instance_name} " f"backup_time=" f"{format_complete_marker_backup_time(source_prefix.backup_time)}" ) s3_client, boto3_version, botocore_version = create_s3_client( region, args.parallel, ) emit( f"export S3 client ready: boto3={boto3_version} " f"botocore={botocore_version} " f"max_pool_connections={args.parallel} " f"max_attempts={AWS_MAX_ATTEMPTS}" ) object_keys = list_meta_keys(s3_client, source_prefix) meta_identities = [ parse_meta_key(source_prefix, object_key) for object_key in object_keys ] emit( f"export plan: total={len(object_keys)} " f"parallel={args.parallel} flat={args.flat} " f"out_prefix={args.out_prefix.rstrip('/')} " f"script_start_unix_nano={SCRIPT_START_UNIX_NANO}" ) if not object_keys: progress.set_total(0) return export_indices = range(len(object_keys)) if args.exclude_meta_name_regexp is not None: export_indices = [] for item_index, object_key in enumerate(object_keys): meta_filename = object_key.rsplit("/", 1)[-1] if args.exclude_meta_name_regexp.search(meta_filename) is not None: emit( f"export filtered: total={len(object_keys)} " f"meta_name={meta_filename}" ) continue export_indices.append(item_index) emit( f"export filter: regexp={args.exclude_meta_name_regexp.pattern!r} " f"excluded={len(object_keys) - len(export_indices)} " f"selected={len(export_indices)}" ) progress.set_total(len(export_indices)) if not export_indices: return temporary_directory = None try: gpg_key_file = args.gpg if gpg_is_s3: temporary_directory = tempfile.TemporaryDirectory( prefix="dumpling-packed-backup-gpg-" ) gpg_key_file = os.path.join(temporary_directory.name, "public-key") download_gpg_key(s3_client, args.gpg, gpg_key_file) items = [] for item_index in export_indices: object_key = object_keys[item_index] meta_identity = meta_identities[item_index] source = f"s3://{source_prefix.bucket}/{object_key}" target = target_uri( output_prefix, SCRIPT_START_UNIX_NANO, meta_identity, args.flat, ) output = dumpling_output_uri( target, source_prefix.backup_time, job_instance_name, ) command = dumpling_command( source, output, region, gpg_key_file, args.dumpling_args, ) items.append( ExportItem( source=source, target=target, output=output, command=command, ) ) random.shuffle(items) emit(f"export queue shuffled: queued={len(items)}") run_export_items( items, args.parallel, len(items), s3_client, progress, job_instance_name, source_prefix.backup_time, args.dumpling_timeout, ) finally: if temporary_directory is not None: temporary_directory.cleanup() def install_signal_handlers() -> None: def handle_signal(signum: int, _frame: object) -> None: raise ReceivedSignal(signum) for signum in (signal.SIGHUP, signal.SIGINT, signal.SIGTERM): signal.signal(signum, handle_signal) def normalize_returncode(returncode: int) -> int: if returncode < 0: return min(255, 128 - returncode) return min(255, max(1, returncode)) def main(argv: Optional[Sequence[str]] = None) -> int: args = parse_args(argv) progress = Progress() exit_code = 0 try: install_signal_handlers() run_export(args, progress) except ReceivedSignal as exc: emit(f"error: received signal {signal.Signals(exc.signum).name}", error=True) exit_code = 128 + exc.signum except KeyboardInterrupt: emit("error: interrupted", error=True) exit_code = 130 except DumplingFailed as exc: exit_code = normalize_returncode(exc.returncode) except Exception as exc: emit(f"error: {exc}", error=True) exit_code = 1 emit_progress_summary(progress, error=exit_code != 0) wait_for_log_collection() return exit_code if __name__ == "__main__": sys.exit(main())