312 lines
12 KiB
Python
312 lines
12 KiB
Python
import asyncio
|
|
import aiohttp
|
|
import aiofiles
|
|
import tarfile
|
|
import shutil
|
|
from pathlib import Path
|
|
from typing import List, Set
|
|
import logging
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
import os
|
|
|
|
# --- configuration ---
|
|
BASE_DIR = Path(__file__).parent.resolve() / Path("../../../..")
|
|
DATA_DIR = BASE_DIR / "data"
|
|
CSV_DIR = BASE_DIR / "data/csv"
|
|
EXTRACT_ADSB = BASE_DIR / "Code/cpp/extract-adsb"
|
|
|
|
# Performance tuning
|
|
MAX_CONCURRENT_DOWNLOADS = 2 # Limit concurrent HTTP requests
|
|
MAX_CONCURRENT_EXTRACTIONS = 4 # Limit concurrent tar extractions (CPU intensive)
|
|
MAX_CONCURRENT_DAYS = 3 # Limit concurrent day processing
|
|
# -----------------------
|
|
|
|
# Set up logging
|
|
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
|
|
logger = logging.getLogger(__name__)
|
|
|
|
def pretty_path(path: os.PathLike | str) -> Path:
|
|
"""Return a path relative to the current working directory."""
|
|
path = Path(path).resolve()
|
|
return Path(os.path.relpath(path, Path.cwd()))
|
|
|
|
def pretty_url(url: str) -> str:
|
|
"""Return a shortened version of the URL for logging."""
|
|
if len(url) <= 60:
|
|
return url
|
|
return f"{url[:30]}...{url[-30:]}"
|
|
|
|
|
|
async def download(session: aiohttp.ClientSession, url: str, dest: Path, semaphore: asyncio.Semaphore) -> None:
|
|
"""Download a file asynchronously with semaphore limiting and retry logic."""
|
|
async with semaphore:
|
|
dest.parent.mkdir(parents=True, exist_ok=True)
|
|
logger.info(f"Downloading {pretty_url(url)} → {pretty_path(dest)}")
|
|
|
|
max_retries = 3
|
|
retry_delay = 10 # seconds
|
|
timeout = 600 # seconds
|
|
|
|
for attempt in range(max_retries + 1):
|
|
try:
|
|
async with session.get(url, timeout=aiohttp.ClientTimeout(total=timeout)) as response:
|
|
if response.status == 429: # Too Many Requests
|
|
wait_time = int(response.headers.get('Retry-After', retry_delay * (attempt + 1)))
|
|
logger.warning(f"Rate limited for {pretty_url(url)}. Waiting {wait_time}s before retry {attempt + 1}/{max_retries}")
|
|
await asyncio.sleep(wait_time)
|
|
continue
|
|
|
|
response.raise_for_status()
|
|
|
|
async with aiofiles.open(dest, 'wb') as f:
|
|
async for chunk in response.content.iter_chunked(8192 * 4):
|
|
await f.write(chunk)
|
|
|
|
|
|
logger.info(f"Download completed: {pretty_path(dest)}")
|
|
return # Success, exit function
|
|
|
|
except aiohttp.ClientResponseError as e:
|
|
if e.status >= 500 and attempt < max_retries: # Server errors are retryable
|
|
logger.warning(f"Server error {e.status} for {pretty_url(url)} on attempt {attempt + 1}/{max_retries}: {e}")
|
|
await asyncio.sleep(retry_delay * (2 ** attempt)) # Exponential backoff
|
|
continue
|
|
else:
|
|
logger.error(f"HTTP error {e.status} for {pretty_url(url)}: {e}")
|
|
raise
|
|
|
|
except aiohttp.ClientConnectorError as e:
|
|
if attempt < max_retries:
|
|
logger.warning(f"Connection error for {pretty_url(url)} on attempt {attempt + 1}/{max_retries}: {e}")
|
|
await asyncio.sleep(retry_delay * (2 ** attempt))
|
|
continue
|
|
else:
|
|
logger.error(f"Connection failed for {pretty_url(url)}: {e}")
|
|
raise
|
|
|
|
except asyncio.TimeoutError as e:
|
|
if attempt < max_retries:
|
|
logger.warning(f"Timeout for {pretty_url(url)} on attempt {attempt + 1}/{max_retries}: {e}")
|
|
await asyncio.sleep(retry_delay * (2 ** attempt))
|
|
continue
|
|
else:
|
|
logger.error(f"Download timeout for {pretty_url(url)}: {e}")
|
|
raise
|
|
|
|
except Exception as e:
|
|
if attempt < max_retries:
|
|
logger.warning(f"Unexpected error for {pretty_url(url)} on attempt {attempt + 1}/{max_retries}: {e}")
|
|
await asyncio.sleep(retry_delay * (2 ** attempt))
|
|
continue
|
|
else:
|
|
logger.error(f"Download failed for {pretty_url(url)} after {max_retries} attempts: {e}")
|
|
raise
|
|
|
|
# This should never be reached, but just in case
|
|
raise Exception(f"Download failed for {pretty_url(url)} after {max_retries} attempts")
|
|
|
|
def get_name_from_url(url: str) -> str:
|
|
"""Extract filename from URL."""
|
|
return url.split("/")[-1]
|
|
|
|
|
|
async def extract_adsb_data(cmd: List[str]) -> None:
|
|
"""Run extract-adsb binary asynchronously."""
|
|
logger.info(f"Running: {' '.join(cmd)}")
|
|
process = await asyncio.create_subprocess_exec(*cmd)
|
|
await process.wait()
|
|
if process.returncode != 0:
|
|
raise Exception(f"Command failed with return code {process.returncode}: {' '.join(cmd)}")
|
|
|
|
|
|
def extract_tar_sync(tar_paths: List[Path], extract_dir: Path) -> Path:
|
|
"""Synchronous tar extraction (CPU-bound, run in thread pool)."""
|
|
extract_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
if len(tar_paths) == 1:
|
|
tar_path = tar_paths[0]
|
|
else:
|
|
# Keep only multipart fragments like .aa, .ab, .ac ...
|
|
parts = [p for p in tar_paths if len(p.suffix) == 3]
|
|
|
|
# Derive correct base archive name: strip only the multipart suffix
|
|
base = parts[0].with_suffix("") # "foo.tar.aa" → "foo.tar"
|
|
tar_path = extract_dir / base.name
|
|
|
|
logger.info(f"Concatenating {len(parts)} parts → {pretty_path(tar_path)}")
|
|
with open(tar_path, "wb") as outfile:
|
|
for part in sorted(parts):
|
|
with open(part, "rb") as infile:
|
|
shutil.copyfileobj(infile, outfile)
|
|
|
|
logger.info(f"Extracting {pretty_path(tar_path)} → {pretty_path(extract_dir)}")
|
|
with tarfile.open(tar_path) as tf:
|
|
tf.extractall(path=extract_dir)
|
|
|
|
return extract_dir
|
|
|
|
|
|
async def extract_tar(tar_paths: List[Path], extract_dir: Path, executor: ThreadPoolExecutor) -> Path:
|
|
"""Run tar extraction in thread pool to avoid blocking event loop."""
|
|
loop = asyncio.get_event_loop()
|
|
return await loop.run_in_executor(executor, extract_tar_sync, tar_paths, extract_dir)
|
|
|
|
|
|
async def move_csv_files(src_dir: Path, dest_dir: Path) -> None:
|
|
"""Move all CSV files from source to destination directory."""
|
|
dest_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
for p in src_dir.rglob("*.csv"):
|
|
dest = dest_dir / p.name
|
|
shutil.move(str(p), str(dest))
|
|
|
|
|
|
async def remove_dir(dir_path: Path) -> None:
|
|
"""Remove directory and all its contents."""
|
|
if dir_path.exists():
|
|
shutil.rmtree(dir_path)
|
|
|
|
|
|
async def remove_file(file_path: Path) -> None:
|
|
"""Remove a single file."""
|
|
if file_path.exists():
|
|
file_path.unlink()
|
|
|
|
|
|
async def download_day_parts(session: aiohttp.ClientSession, urls: List[str],
|
|
download_semaphore: asyncio.Semaphore) -> List[Path]:
|
|
"""Download all parts for a single day concurrently."""
|
|
tar_files = []
|
|
download_tasks = []
|
|
|
|
for url in urls:
|
|
filename = DATA_DIR / get_name_from_url(url)
|
|
tar_files.append(filename)
|
|
download_tasks.append(
|
|
download(session, url, filename, download_semaphore)
|
|
)
|
|
|
|
await asyncio.gather(*download_tasks)
|
|
return tar_files
|
|
|
|
|
|
async def process_single_day(session: aiohttp.ClientSession, line: str,
|
|
download_semaphore: asyncio.Semaphore,
|
|
extraction_executor: ThreadPoolExecutor) -> None:
|
|
"""Process a single day's data with concurrent downloads and parallel extraction."""
|
|
parts = line.strip().split(",")
|
|
if not parts or not parts[0]:
|
|
return
|
|
|
|
if len(parts) > 1:
|
|
# Remove any part that is not of ending .tar.aa, .tar.ab, etc.
|
|
parts = [p for p in parts if len(p.split(".")[-1]) == 2]
|
|
|
|
logger.info(f"Processing day with {len(parts)} parts")
|
|
|
|
try:
|
|
# Download all parts for this day concurrently
|
|
tar_files = await download_day_parts(session, parts, download_semaphore)
|
|
|
|
# Extract the tar files
|
|
extract_dir = tar_files[0].parent / tar_files[0].stem
|
|
await extract_tar(tar_files, extract_dir, extraction_executor)
|
|
|
|
# Process ADSB data
|
|
traces_dir = extract_dir / "traces"
|
|
if traces_dir.exists():
|
|
await extract_adsb_data([str(EXTRACT_ADSB), str(traces_dir)])
|
|
|
|
# Move CSV files
|
|
csv_files_dir = CSV_DIR / extract_dir.name
|
|
await move_csv_files(extract_dir, csv_files_dir)
|
|
|
|
# Cleanup
|
|
cleanup_tasks = [
|
|
asyncio.create_task(remove_file(tar_file))
|
|
for tar_file in tar_files
|
|
if tar_file.exists()
|
|
]
|
|
cleanup_tasks.append(asyncio.create_task(remove_dir(extract_dir)))
|
|
|
|
await asyncio.gather(*cleanup_tasks)
|
|
logger.info(f"Completed processing day: {extract_dir.name}")
|
|
|
|
except Exception as e:
|
|
logger.error(f"Failed to process day: {e}")
|
|
raise
|
|
else:
|
|
# If successful, write to ignore file
|
|
await write_line_to_ignore(BASE_DIR / "Code/python/adsbpy/.ignore", line.strip())
|
|
|
|
async def write_line_to_ignore(file: Path, line: str) -> None:
|
|
"""Append a line to the ignore file asynchronously."""
|
|
async with aiofiles.open(file, 'a') as f:
|
|
await f.write(line + '\n')
|
|
|
|
|
|
async def process_days_concurrently(lines: List[str]) -> None:
|
|
"""Process multiple days concurrently with controlled parallelism."""
|
|
# Create semaphores for controlling concurrency
|
|
download_semaphore = asyncio.Semaphore(MAX_CONCURRENT_DOWNLOADS)
|
|
day_semaphore = asyncio.Semaphore(MAX_CONCURRENT_DAYS)
|
|
|
|
# Thread pool for CPU-bound tar extraction
|
|
with ThreadPoolExecutor(max_workers=MAX_CONCURRENT_EXTRACTIONS) as extraction_executor:
|
|
connector = aiohttp.TCPConnector(limit=MAX_CONCURRENT_DOWNLOADS)
|
|
async with aiohttp.ClientSession(connector=connector) as session:
|
|
|
|
async def process_day_with_semaphore(line: str) -> None:
|
|
"""Wrapper to limit concurrent day processing."""
|
|
async with day_semaphore:
|
|
await process_single_day(session, line, download_semaphore, extraction_executor)
|
|
|
|
# Process all days concurrently with limits
|
|
tasks = [
|
|
process_day_with_semaphore(line)
|
|
for line in lines
|
|
if line.strip()
|
|
]
|
|
|
|
# Process in batches to avoid overwhelming the system
|
|
batch_size = MAX_CONCURRENT_DAYS * 5
|
|
for i in range(0, len(tasks), batch_size):
|
|
batch = tasks[i:i + batch_size]
|
|
logger.info(f"Processing batch {i//batch_size + 1}/{(len(tasks)-1)//batch_size + 1}")
|
|
await asyncio.gather(*batch, return_exceptions=True)
|
|
|
|
|
|
async def fetch_lines(lines: List[str]) -> None:
|
|
"""Process multiple lines (days) concurrently."""
|
|
await process_days_concurrently(lines)
|
|
|
|
|
|
async def fetch_file(file: Path) -> None:
|
|
"""Fetch and process URLs from a file."""
|
|
async with aiofiles.open(file, 'r') as f:
|
|
lines = await f.readlines()
|
|
|
|
logger.info(f"Processing {len(lines)} days")
|
|
await fetch_lines(lines)
|
|
|
|
|
|
async def main() -> None:
|
|
"""Main entry point if running as a script."""
|
|
import sys
|
|
|
|
if len(sys.argv) > 1:
|
|
input_file = Path(sys.argv[1])
|
|
else:
|
|
input_file = Path("urls.txt")
|
|
|
|
if input_file.exists():
|
|
logger.info(f"Starting processing of {input_file}")
|
|
await fetch_file(input_file)
|
|
logger.info("Processing completed successfully")
|
|
else:
|
|
logger.error(f"Input file {input_file} not found")
|
|
sys.exit(1)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
asyncio.run(main()) |