from objstore.objstore import pick_object_store from datetime import datetime, timedelta import os from concurrent.futures import ThreadPoolExecutor, as_completed from utils.logger_utils import log_info, log_error # Define the interval in seconds for merged log files. # This constant specifies that log files are merged on an hourly basis. MERGED_FILE_INTERVAL_SECONDS = 3600 # Define a list of default components to download logs for if none are specified. DEFAULT_COMPONENTS = [ "pd", "tikv", "tidb", "tiflash", ] def _download_single_file(store, bucket, s3_key, local_path): """ Helper function to download a single file from an object store to a local path. Args: store: An instance of the object store client (e.g., S3 client). bucket (str): The name of the bucket where the file is stored. s3_key (str): The S3 key (path) of the file to download. local_path (str): The local file path where the downloaded file will be saved. Returns: bool: True if the download was successful or the file already exists, False otherwise. """ log_info(f"Downloading {s3_key} from S3 to {local_path}") # Ensure the local directory exists before attempting to download. os.makedirs(os.path.dirname(local_path), exist_ok=True) # Check if the file already exists locally to avoid re-downloading. if os.path.exists(local_path): log_info(f"File already exists locally: {local_path}. Skipping download.") return True else: # Attempt to download the file. if store.download_file(bucket, s3_key, local_path): log_info(f"Successfully downloaded {s3_key} to {local_path}") return True else: log_info(f"Failed to download {s3_key}") if os.path.exists(local_path): os.remove(local_path) return False def download_merged_logs(cluster_id, vendor, region, components, start_time, end_time, destination_path): """ Downloads merged log files for a given cluster within a specified time range. Logs are downloaded from an object store (e.g., S3) and saved to a local directory. The function supports concurrent downloads for efficiency. Args: cluster_id (str): The ID of the cluster. vendor (str): The cloud vendor (e.g., "aws", "gcp"). region (str): The region where the object store bucket is located. components (list[str]): A list of components for which to download logs (e.g., ["pd", "tikv"]). If the list is empty, logs for `DEFAULT_COMPONENTS` will be downloaded. start_time (str): The start time for log download in 'YYYYMMDDHH' format. end_time (str): The end time for log download in 'YYYYMMDDHH' format. destination_path (str): The base local path where the downloaded logs will be stored. """ # Initialize the object store client based on the vendor. store = pick_object_store(vendor) # Construct the bucket name using a standardized format. bucket = f"o11y-prod-shared-{region}" # Define the base prefix for log files within the object store. store_prefix = f"diagnosis/data/{cluster_id}" # Define the base local destination path for logs. local_base_dest = f"{destination_path}/logs/{cluster_id}" # Use default components if no specific components are provided. if not components: # Checks if the list is empty components = DEFAULT_COMPONENTS log_info( f"No components specified, using default components: {components}") # Parse start and end time strings into datetime objects for easier manipulation. try: start_dt = datetime.strptime(start_time, '%Y%m%d%H') end_dt = datetime.strptime(end_time, '%Y%m%d%H') except ValueError as e: log_error( f"Error parsing start_time or end_time. Please use 'YYYYMMDDHH' format: {e}") return # Generate a list of download tasks, each containing S3 key and local file path. # Support two path formats: # old format: {store_prefix}/merged-logs/{component}/{current_hour_str}.log # new format: {store_prefix}/merged-logs/{current_hour_str}/{component}/{instance}.log download_tasks = [] current_dt = start_dt while current_dt <= end_dt: current_hour_str = current_dt.strftime('%Y%m%d%H') for component in components: # Format 1: {store_prefix}/merged-logs/{component}/{current_hour_str}.log s3_key_format1 = f"{store_prefix}/merged-logs/{component}/{current_hour_str}.log" try: files_in_dir = store.list(bucket, s3_key_format1) for file in files_in_dir: local_file_format1 = f"{local_base_dest}/{component}/{current_hour_str}.log" download_tasks.append({"s3_key": file, "local_file": local_file_format1}) except Exception as e: log_error(f"Error listing files for prefix {s3_key_format1}: {e}") # Format 2: {store_prefix}/merged-logs/{current_hour_str}/{component}/{instance}.log # List all files under this prefix and download them prefix_format2 = f"{store_prefix}/merged-logs/{current_hour_str}/{component}/" try: files_in_dir = store.list(bucket, prefix_format2) for s3_key_format2 in files_in_dir: filename = os.path.basename(s3_key_format2) local_file_format2 = f"{local_base_dest}/{component}/{current_hour_str}/{filename}" download_tasks.append({"s3_key": s3_key_format2, "local_file": local_file_format2}) except Exception as e: log_error(f"Error listing files for prefix {prefix_format2}: {e}") current_dt += timedelta(hours=1) # Execute downloads concurrently using a thread pool. if download_tasks: # Determine the maximum number of worker threads. # It's set to the minimum of the number of tasks or a reasonable maximum (e.g., 30) # to prevent overwhelming the system or the object storage service. max_workers = min(len(download_tasks), 30) log_info( f"Initiating download of {len(download_tasks)} log files with {max_workers} concurrent threads.") with ThreadPoolExecutor(max_workers=max_workers) as executor: # Submit download tasks to the thread pool. futures = [ executor.submit( _download_single_file, store, bucket, task["s3_key"], task["local_file"] ) for task in download_tasks ] # Wait for all futures to complete and process results. for future in as_completed(futures): try: success = future.result() if not success: log_error("One or more downloads failed.") except Exception as exc: log_error(f"A download task generated an exception: {exc}") else: log_info("No log files to download for the specified criteria.")