diff --git a/providers/edge3/src/airflow/providers/edge3/plugins/edge_executor_plugin.py b/providers/edge3/src/airflow/providers/edge3/plugins/edge_executor_plugin.py index ad22956258c0e..eece750ece0c8 100644 --- a/providers/edge3/src/airflow/providers/edge3/plugins/edge_executor_plugin.py +++ b/providers/edge3/src/airflow/providers/edge3/plugins/edge_executor_plugin.py @@ -17,38 +17,176 @@ from __future__ import annotations +import logging +import random import sys +import time +from contextlib import asynccontextmanager from typing import TYPE_CHECKING, Any +from sqlalchemy import inspect +from sqlalchemy.exc import OperationalError + from airflow.configuration import conf from airflow.exceptions import AirflowConfigException from airflow.providers.common.compat.sdk import AirflowPlugin from airflow.providers.edge3.version_compat import AIRFLOW_V_3_1_PLUS from airflow.utils.session import NEW_SESSION, provide_session +from airflow.utils.sqlalchemy import is_lock_not_available_error if TYPE_CHECKING: + from fastapi import FastAPI + from sqlalchemy.engine import Engine from sqlalchemy.orm import Session from airflow.utils.db import DBLocks, create_global_lock +log = logging.getLogger(__name__) + +# Retry configuration for lock acquisition during table creation +MAX_LOCK_RETRIES = 5 +LOCK_RETRY_DELAY_BASE = 1.0 # Base delay in seconds for exponential backoff + +# Required edge tables +EDGE_TABLES = ("edge_job", "edge_logs", "edge_worker") + + +def _tables_exist(engine: Engine) -> bool: + """ + Check if all required edge tables already exist in the database. + + This is a fast path to avoid acquiring a global lock when tables + are already present (normal operation after initial setup). + """ + try: + inspector = inspect(engine) + existing_tables = set(inspector.get_table_names()) + return all(table in existing_tables for table in EDGE_TABLES) + except Exception: + # If we can't check, assume tables don't exist and proceed with creation + return False + @provide_session -def _get_api_endpoint(session: Session = NEW_SESSION) -> dict[str, Any]: - # Ensure all required DB modeals are created before starting the API - with create_global_lock(session=session, lock=DBLocks.MIGRATIONS): - engine = session.get_bind().engine - from airflow.providers.edge3.models.edge_job import EdgeJobModel - from airflow.providers.edge3.models.edge_logs import EdgeLogsModel - from airflow.providers.edge3.models.edge_worker import EdgeWorkerModel - - EdgeJobModel.metadata.create_all(engine) - EdgeLogsModel.metadata.create_all(engine) - EdgeWorkerModel.metadata.create_all(engine) +def _ensure_tables_created(session: Session = NEW_SESSION) -> None: + """ + Ensure all required DB models are created with retry logic. + + This is called lazily on FastAPI app startup, not at plugin import time, + to avoid blocking all API server processes on a global database lock + during concurrent startup. + + Uses exponential backoff with jitter to handle lock contention when + multiple API server processes start simultaneously. + + Fast path: If tables already exist, skips lock acquisition entirely. + """ + from airflow.providers.edge3.models.edge_job import EdgeJobModel + from airflow.providers.edge3.models.edge_logs import EdgeLogsModel + from airflow.providers.edge3.models.edge_worker import EdgeWorkerModel + engine = session.get_bind().engine + + # Fast path: skip lock acquisition if tables already exist + if _tables_exist(engine): + log.debug("Edge tables already exist, skipping creation.") + return + + last_error: OperationalError | RuntimeError | None = None + + for attempt in range(MAX_LOCK_RETRIES): + try: + log.debug("Ensuring edge tables exist (attempt %d/%d)...", attempt + 1, MAX_LOCK_RETRIES) + with create_global_lock(session=session, lock=DBLocks.MIGRATIONS): + # Double-check after acquiring lock (another process may have created them) + if _tables_exist(engine): + log.debug("Edge tables were created by another process.") + return + EdgeJobModel.metadata.create_all(engine) + EdgeLogsModel.metadata.create_all(engine) + EdgeWorkerModel.metadata.create_all(engine) + log.debug("Edge tables created successfully.") + return + except OperationalError as e: + # Use the existing function that handles PostgreSQL (55P03) and MySQL (1205, 3572) error codes + if is_lock_not_available_error(e): + last_error = e + if attempt < MAX_LOCK_RETRIES - 1: + # Exponential backoff with jitter to avoid thundering herd + delay = LOCK_RETRY_DELAY_BASE * (2**attempt) + random.uniform(0, 1) + log.warning( + "Could not acquire migration lock for edge tables (attempt %d/%d), " + "retrying in %.1f seconds...", + attempt + 1, + MAX_LOCK_RETRIES, + delay, + ) + time.sleep(delay) + # Get a fresh session for retry to avoid stale connection issues + session.rollback() + # Check if tables were created while we were waiting + if _tables_exist(engine): + log.debug("Edge tables were created by another process during retry wait.") + return + else: + # Re-raise non-lock-related errors immediately + raise + except RuntimeError as e: + # MySQL's create_global_lock raises RuntimeError when GET_LOCK times out + # Check if it's a lock-related error by examining the error message + error_str = str(e).lower() + if "lock" in error_str and ("could not acquire" in error_str or "timeout" in error_str): + last_error = e + if attempt < MAX_LOCK_RETRIES - 1: + # Exponential backoff with jitter to avoid thundering herd + delay = LOCK_RETRY_DELAY_BASE * (2**attempt) + random.uniform(0, 1) + log.warning( + "Could not acquire migration lock for edge tables (attempt %d/%d), " + "retrying in %.1f seconds...", + attempt + 1, + MAX_LOCK_RETRIES, + delay, + ) + time.sleep(delay) + # Get a fresh session for retry to avoid stale connection issues + session.rollback() + # Check if tables were created while we were waiting + if _tables_exist(engine): + log.debug("Edge tables were created by another process during retry wait.") + return + else: + # Re-raise non-lock-related RuntimeErrors immediately + raise + + # All retries exhausted + log.error( + "Failed to acquire migration lock for edge tables after %d attempts. " + "Edge worker API may not function correctly.", + MAX_LOCK_RETRIES, + ) + if last_error: + raise last_error + + +def _get_api_endpoint() -> dict[str, Any]: + """ + Get API endpoint configuration. + + Table creation is deferred to FastAPI startup event to avoid + blocking plugin import on database lock acquisition. + """ from airflow.providers.edge3.worker_api.app import create_edge_worker_api_app + @asynccontextmanager + async def lifespan(app: FastAPI): # type: ignore[type-arg] + """Create edge tables on app startup instead of at import time.""" + _ensure_tables_created() + yield + + app = create_edge_worker_api_app(lifespan=lifespan) + return { - "app": create_edge_worker_api_app(), + "app": app, "url_prefix": "/edge_worker", "name": "Airflow Edge Worker", } diff --git a/providers/edge3/src/airflow/providers/edge3/worker_api/app.py b/providers/edge3/src/airflow/providers/edge3/worker_api/app.py index 1e1c6d619ad22..044ee9e157738 100644 --- a/providers/edge3/src/airflow/providers/edge3/worker_api/app.py +++ b/providers/edge3/src/airflow/providers/edge3/worker_api/app.py @@ -17,6 +17,7 @@ from __future__ import annotations from pathlib import Path +from typing import Any from fastapi import FastAPI from fastapi.staticfiles import StaticFiles @@ -28,7 +29,9 @@ from airflow.providers.edge3.worker_api.routes.worker import worker_router -def create_edge_worker_api_app() -> FastAPI: +def create_edge_worker_api_app( + lifespan: Any = None, +) -> FastAPI: """Create FastAPI app for edge worker API.""" edge_worker_api_app = FastAPI( title="Airflow Edge Worker API", @@ -40,6 +43,7 @@ def create_edge_worker_api_app() -> FastAPI: "All endpoints under ``/edge_worker/ui`` are used by UI and can be accessed with normal authentication. " "Please assume UI endpoints to change and not be stable." ), + lifespan=lifespan, ) edge_worker_api_app.include_router(jobs_router, prefix="/v1")