You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@tooling.apache.org by sb...@apache.org on 2025/03/12 17:59:09 UTC
(tooling-trusted-release) branch main updated: Start to use async more consistently
This is an automated email from the ASF dual-hosted git repository.
sbp pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tooling-trusted-release.git
The following commit(s) were added to refs/heads/main by this push:
new 5f8665c Start to use async more consistently
5f8665c is described below
commit 5f8665cdb06ce64859b0f05c0fe2885f0380bd74
Author: Sean B. Palmer <se...@miscoranda.com>
AuthorDate: Wed Mar 12 19:58:54 2025 +0200
Start to use async more consistently
---
atr/db/__init__.py | 20 +++++++++--
atr/server.py | 66 +++++++++++++++++++++----------------
atr/tasks/archive.py | 6 ++--
atr/worker.py | 93 +++++++++++++++++++++++++++-------------------------
4 files changed, 106 insertions(+), 79 deletions(-)
diff --git a/atr/db/__init__.py b/atr/db/__init__.py
index 7f59e86..8209976 100644
--- a/atr/db/__init__.py
+++ b/atr/db/__init__.py
@@ -84,10 +84,24 @@ def create_database(app: QuartApp) -> None:
def create_async_db_session() -> sqlalchemy.ext.asyncio.AsyncSession:
"""Create a new asynchronous database session."""
- async_session = util.validate_as_type(
- quart.current_app.extensions["async_session"](), sqlalchemy.ext.asyncio.AsyncSession
+ if quart.has_app_context():
+ extensions = quart.current_app.extensions
+ return util.validate_as_type(extensions["async_session"](), sqlalchemy.ext.asyncio.AsyncSession)
+
+ import atr.config as config
+
+ conf = config.get()
+ sqlite_url = f"sqlite+aiosqlite://{conf.SQLITE_DB_PATH}"
+ engine = sqlalchemy.ext.asyncio.create_async_engine(
+ sqlite_url,
+ connect_args={
+ "check_same_thread": False,
+ "timeout": 30,
+ },
)
- return async_session
+ return sqlalchemy.ext.asyncio.async_sessionmaker(
+ bind=engine, class_=sqlalchemy.ext.asyncio.AsyncSession, expire_on_commit=False
+ )()
def create_sync_db_engine() -> None:
diff --git a/atr/server.py b/atr/server.py
index 3df1727..1ff9dc9 100644
--- a/atr/server.py
+++ b/atr/server.py
@@ -22,10 +22,10 @@ import os
from collections.abc import Iterable
from typing import Any
-from blockbuster import BlockBuster
-from quart import render_template
-from quart_schema import OpenAPIProvider, QuartSchema
-from werkzeug.routing import Rule
+import blockbuster
+import quart
+import quart_schema
+import werkzeug.routing as routing
import asfquart
import asfquart.base as base
@@ -33,11 +33,10 @@ import asfquart.generics
import asfquart.session
import atr.blueprints as blueprints
import atr.config as config
+import atr.db as db
+import atr.manager as manager
+import atr.preload as preload
import atr.ssh as ssh
-from asfquart.base import ASFQuartException
-from atr.db import create_database
-from atr.manager import get_worker_manager
-from atr.preload import setup_template_preloading
# TODO: Technically this is a global variable
# We should probably find a cleaner way to do this
@@ -48,15 +47,24 @@ asfquart.generics.OAUTH_URL_INIT = "https://oauth.apache.org/auth?state=%s&redir
asfquart.generics.OAUTH_URL_CALLBACK = "https://oauth.apache.org/token?code=%s"
-class ApiOnlyOpenAPIProvider(OpenAPIProvider):
- def generate_rules(self) -> Iterable[Rule]:
+class ApiOnlyOpenAPIProvider(quart_schema.OpenAPIProvider):
+ def generate_rules(self) -> Iterable[routing.Rule]:
for rule in super().generate_rules():
if rule.rule.startswith("/api"):
yield rule
def register_routes(app: base.QuartApp) -> tuple[str, ...]:
- from atr.routes import candidate, dev, docs, download, keys, package, project, release, root
+ # NOTE: These imports are for their side effects only
+ import atr.routes.candidate as candidate
+ import atr.routes.dev as dev
+ import atr.routes.docs as docs
+ import atr.routes.download as download
+ import atr.routes.keys as keys
+ import atr.routes.package as package
+ import atr.routes.project as project
+ import atr.routes.release as release
+ import atr.routes.root as root
# Add a global error handler to show helpful error messages with tracebacks.
@app.errorhandler(Exception)
@@ -65,21 +73,21 @@ def register_routes(app: base.QuartApp) -> tuple[str, ...]:
tb = traceback.format_exc()
app.logger.error(f"Unhandled exception: {error}\n{tb}")
- return await render_template("error.html", error=str(error), traceback=tb, status_code=500), 500
+ return await quart.render_template("error.html", error=str(error), traceback=tb, status_code=500), 500
- @app.errorhandler(ASFQuartException)
- async def handle_asfquart_exception(error: ASFQuartException) -> Any:
+ @app.errorhandler(base.ASFQuartException)
+ async def handle_asfquart_exception(error: base.ASFQuartException) -> Any:
# TODO: Figure out why pyright doesn't know about this attribute
if not hasattr(error, "errorcode"):
errorcode = 500
else:
errorcode = getattr(error, "errorcode")
- return await render_template("error.html", error=str(error), status_code=errorcode), errorcode
+ return await quart.render_template("error.html", error=str(error), status_code=errorcode), errorcode
# Add a global error handler in case a page does not exist.
@app.errorhandler(404)
async def handle_not_found(error: Exception) -> Any:
- return await render_template("notfound.html", error="404 Not Found", traceback="", status_code=404), 404
+ return await quart.render_template("notfound.html", error="404 Not Found", traceback="", status_code=404), 404
# Must do this otherwise ruff "fixes" this function by removing the imports
return (
@@ -119,7 +127,7 @@ def app_setup_api_docs(app: base.QuartApp) -> None:
import atr.metadata as metadata
- QuartSchema(
+ quart_schema.QuartSchema(
app,
info=quart_schema.Info(
title="ATR API",
@@ -154,7 +162,7 @@ def app_setup_lifecycle(app: base.QuartApp) -> None:
@app.before_serving
async def startup() -> None:
"""Start services before the app starts serving requests."""
- worker_manager = get_worker_manager()
+ worker_manager = manager.get_worker_manager()
await worker_manager.start()
ssh_server = await ssh.server_start()
@@ -163,7 +171,7 @@ def app_setup_lifecycle(app: base.QuartApp) -> None:
@app.after_serving
async def shutdown() -> None:
"""Clean up services after the app stops serving requests."""
- worker_manager = get_worker_manager()
+ worker_manager = manager.get_worker_manager()
await worker_manager.stop()
ssh_server = app.extensions.get("ssh_server")
@@ -197,7 +205,7 @@ def create_app(app_config: type[config.AppConfig]) -> base.QuartApp:
app = app_create_base(app_config)
app_setup_api_docs(app)
- create_database(app)
+ db.create_database(app)
register_routes(app)
blueprints.register(app)
@@ -209,24 +217,24 @@ def create_app(app_config: type[config.AppConfig]) -> base.QuartApp:
# do not enable template pre-loading if we explicitly want to reload templates
if not app_config.TEMPLATES_AUTO_RELOAD:
- setup_template_preloading(app)
+ preload.setup_template_preloading(app)
@app.before_serving
async def start_blockbuster() -> None:
# "I'll have a P, please, Bob."
- blockbuster: BlockBuster | None = None
+ bb: blockbuster.BlockBuster | None = None
if config_mode == config.Mode.Profiling:
- blockbuster = BlockBuster()
- app.extensions["blockbuster"] = blockbuster
- if blockbuster is not None:
- blockbuster.activate()
+ bb = blockbuster.BlockBuster()
+ app.extensions["blockbuster"] = bb
+ if bb is not None:
+ bb.activate()
app.logger.info("Blockbuster activated to detect blocking calls")
@app.after_serving
async def stop_blockbuster() -> None:
- blockbuster = app.extensions.get("blockbuster")
- if blockbuster is not None:
- blockbuster.deactivate()
+ bb = app.extensions.get("blockbuster")
+ if bb is not None:
+ bb.deactivate()
app.logger.info("Blockbuster deactivated")
return app
diff --git a/atr/tasks/archive.py b/atr/tasks/archive.py
index ae057be..677cae0 100644
--- a/atr/tasks/archive.py
+++ b/atr/tasks/archive.py
@@ -24,7 +24,7 @@ import pydantic
import atr.tasks.task as task
-_LOGGER = logging.getLogger(__name__)
+_LOGGER: Final = logging.getLogger(__name__)
class CheckIntegrity(pydantic.BaseModel):
@@ -34,7 +34,7 @@ class CheckIntegrity(pydantic.BaseModel):
chunk_size: int = pydantic.Field(default=4096, description="Size of chunks to read when checking the file")
-def check_integrity(args: dict[str, Any]) -> tuple[task.Status, str | None, tuple[Any, ...]]:
+async def check_integrity(args: dict[str, Any]) -> tuple[task.Status, str | None, tuple[Any, ...]]:
"""Check the integrity of a .tar.gz file."""
# TODO: We should standardise the "ERROR" mechanism here in the data
# Then we can have a single task wrapper for all tasks
@@ -73,7 +73,7 @@ def root_directory(tgz_path: str) -> str:
return root
-def _check_integrity_core(tgz_path: str, chunk_size: int = 4096) -> int:
+async def _check_integrity_core(tgz_path: str, chunk_size: int = 4096) -> int:
"""Verify a .tar.gz file and compute its uncompressed size."""
total_size = 0
diff --git a/atr/worker.py b/atr/worker.py
index 1886bca..0214d68 100644
--- a/atr/worker.py
+++ b/atr/worker.py
@@ -22,6 +22,7 @@
# the task to status='QUEUED'. For this to work, ideally we
# need to check wall clock time as well as CPU time.
+import asyncio
import datetime
import json
import logging
@@ -29,8 +30,7 @@ import os
import resource
import signal
import sys
-import time
-from typing import Any
+from typing import Any, Final
import sqlalchemy
@@ -45,11 +45,11 @@ import atr.tasks.signature as signature
import atr.tasks.task as task
import atr.tasks.vote as vote
-_LOGGER = logging.getLogger(__name__)
+_LOGGER: Final = logging.getLogger(__name__)
# Resource limits, 5 minutes and 1GB
-CPU_LIMIT_SECONDS = 300
-MEMORY_LIMIT_BYTES = 1024 * 1024 * 1024
+# _CPU_LIMIT_SECONDS: Final = 300
+_MEMORY_LIMIT_BYTES: Final = 1024 * 1024 * 1024
# # Create tables if they don't exist
# SQLModel.metadata.create_all(engine)
@@ -59,23 +59,23 @@ def main() -> None:
"""Main entry point."""
import atr.config as config
- signal.signal(signal.SIGTERM, worker_signal_handle)
- signal.signal(signal.SIGINT, worker_signal_handle)
+ signal.signal(signal.SIGTERM, _worker_signal_handle)
+ signal.signal(signal.SIGINT, _worker_signal_handle)
conf = config.get()
if os.path.isdir(conf.STATE_DIR):
os.chdir(conf.STATE_DIR)
- setup_logging()
+ _setup_logging()
_LOGGER.info(f"Starting worker process with pid {os.getpid()}")
- db.create_sync_db_engine()
+ # db.create_sync_db_engine()
- worker_resources_limit_set()
- worker_loop_run()
+ _worker_resources_limit_set()
+ asyncio.run(_worker_loop_run())
-def setup_logging() -> None:
+def _setup_logging() -> None:
# Configure logging
log_format = "[%(asctime)s.%(msecs)03d] [%(process)d] [%(levelname)s] %(message)s"
date_format = "%Y-%m-%d %H:%M:%S"
@@ -86,14 +86,14 @@ def setup_logging() -> None:
# Task functions
-def task_error_handle(task_id: int, e: Exception) -> None:
+async def _task_error_handle(task_id: int, e: Exception) -> None:
"""Handle task error by updating the database with error information."""
if isinstance(e, task.Error):
_LOGGER.error(f"Task {task_id} failed: {e.message}")
result = json.dumps(e.result)
- with db.create_sync_db_session() as session:
- with session.begin():
- session.execute(
+ async with db.create_async_db_session() as session:
+ async with session.begin():
+ await session.execute(
sqlalchemy.text("""
UPDATE task
SET status = 'FAILED', completed = :now, error = :error, result = :result
@@ -108,9 +108,9 @@ def task_error_handle(task_id: int, e: Exception) -> None:
)
else:
_LOGGER.error(f"Task {task_id} failed: {e}")
- with db.create_sync_db_session() as session:
- with session.begin():
- session.execute(
+ async with db.create_async_db_session() as session:
+ async with session.begin():
+ await session.execute(
sqlalchemy.text("""
UPDATE task
SET status = 'FAILED', completed = :now, error = :error
@@ -120,17 +120,17 @@ def task_error_handle(task_id: int, e: Exception) -> None:
)
-def task_next_claim() -> tuple[int, str, str] | None:
+async def _task_next_claim() -> tuple[int, str, str] | None:
"""
Attempt to claim the oldest unclaimed task.
Returns (task_id, task_type, task_args) if successful.
Returns None if no tasks are available.
"""
- with db.create_sync_db_session() as session:
- with session.begin():
+ async with db.create_async_db_session() as session:
+ async with session.begin():
# Find and claim the oldest unclaimed task
# We have an index on (status, added)
- result = session.execute(
+ result = await session.execute(
sqlalchemy.text("""
UPDATE task
SET started = :now, pid = :pid, status = 'ACTIVE'
@@ -153,15 +153,15 @@ def task_next_claim() -> tuple[int, str, str] | None:
return None
-def task_result_process(
+async def _task_result_process(
task_id: int, task_results: tuple[Any, ...], status: str = "COMPLETED", error: str | None = None
) -> None:
"""Process and store task results in the database."""
- with db.create_sync_db_session() as session:
+ async with db.create_async_db_session() as session:
result = json.dumps(task_results)
- with session.begin():
+ async with session.begin():
if status == "FAILED" and error:
- session.execute(
+ await session.execute(
sqlalchemy.text("""
UPDATE task
SET status = :status, completed = :now, result = :result, error = :error
@@ -176,7 +176,7 @@ def task_result_process(
},
)
else:
- session.execute(
+ await session.execute(
sqlalchemy.text("""
UPDATE task
SET status = :status, completed = :now, result = :result
@@ -191,7 +191,7 @@ def task_result_process(
)
-def task_process(task_id: int, task_type: str, task_args: str) -> None:
+async def _task_process(task_id: int, task_type: str, task_args: str) -> None:
"""Process a claimed task."""
_LOGGER.info(f"Processing task {task_id} ({task_type}) with args {task_args}")
try:
@@ -202,6 +202,8 @@ def task_process(task_id: int, task_type: str, task_args: str) -> None:
dict_task_handlers = {
"verify_archive_integrity": archive.check_integrity,
}
+ # TODO: These are synchronous
+ # We plan to convert these to async dict handlers
list_task_handlers = {
"verify_archive_structure": archive.check_structure,
"verify_license_files": license.check_files,
@@ -220,7 +222,7 @@ def task_process(task_id: int, task_type: str, task_args: str) -> None:
msg = f"Unknown task type: {task_type}"
_LOGGER.error(msg)
raise Exception(msg)
- status, error, task_results = dict_handler(args)
+ status, error, task_results = await dict_handler(args)
else:
list_handler = list_task_handlers.get(task_type)
if not list_handler:
@@ -229,13 +231,13 @@ def task_process(task_id: int, task_type: str, task_args: str) -> None:
raise Exception(msg)
status, error, task_results = list_handler(args)
- task_result_process(task_id, task_results, status=status.value.upper(), error=error)
+ await _task_result_process(task_id, task_results, status=status.value.upper(), error=error)
except Exception as e:
- task_error_handle(task_id, e)
+ await _task_error_handle(task_id, e)
-def task_process_wrap(item: Any) -> tuple[Any, ...]:
+async def _task_process_wrap(item: Any) -> tuple[Any, ...]:
"""Ensure that returned results are structured as a tuple."""
if not isinstance(item, tuple):
return (item,)
@@ -245,27 +247,30 @@ def task_process_wrap(item: Any) -> tuple[Any, ...]:
# Worker functions
-def worker_loop_run() -> None:
+async def _worker_loop_run() -> None:
"""Main worker loop."""
while True:
try:
- task = task_next_claim()
+ task = await _task_next_claim()
if task:
task_id, task_type, task_args = task
- task_process(task_id, task_type, task_args)
+ await _task_process(task_id, task_type, task_args)
# Only process one task and then exit
# This prevents memory leaks from accumulating
break
else:
# No tasks available, wait 20ms before checking again
- time.sleep(0.02)
- except Exception as e:
+ await asyncio.sleep(0.02)
+ except Exception:
# TODO: Should probably be more robust about this
- _LOGGER.error(f"Worker loop error: {e}")
- time.sleep(1)
+ # Extract the traceback and log it
+ import traceback
+ _LOGGER.error(f"Worker loop error: {traceback.format_exc()}")
+ await asyncio.sleep(1)
-def worker_resources_limit_set() -> None:
+
+def _worker_resources_limit_set() -> None:
"""Set CPU and memory limits for this process."""
# # Set CPU time limit
# try:
@@ -276,13 +281,13 @@ def worker_resources_limit_set() -> None:
# Set memory limit
try:
- resource.setrlimit(resource.RLIMIT_AS, (MEMORY_LIMIT_BYTES, MEMORY_LIMIT_BYTES))
- _LOGGER.info(f"Set memory limit to {MEMORY_LIMIT_BYTES} bytes")
+ resource.setrlimit(resource.RLIMIT_AS, (_MEMORY_LIMIT_BYTES, _MEMORY_LIMIT_BYTES))
+ _LOGGER.info(f"Set memory limit to {_MEMORY_LIMIT_BYTES} bytes")
except ValueError as e:
_LOGGER.warning(f"Could not set memory limit: {e}")
-def worker_signal_handle(signum: int, frame: object) -> None:
+def _worker_signal_handle(signum: int, frame: object) -> None:
"""Handle termination signals gracefully."""
# For RLIMIT_AS we'll generally get a SIGKILL
# For RLIMIT_CPU we'll get a SIGXCPU, which we can catch
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@tooling.apache.org
For additional commands, e-mail: commits-help@tooling.apache.org