From 1318efc2488ed540083983eb80429173d34b0c14 Mon Sep 17 00:00:00 2001 From: BiFangKNT <1320414964@qq.com> Date: Sat, 17 Oct 2026 19:41:55 +0800 Subject: [PATCH] perf(persistence): index workspace execution leases Add composite indexes on Workspace, status, and lease expiry for AgentRun and MonitoringMessage so lease maintenance can locate active executions without scanning completed history. Keep ORM metadata and Alembic migration 0039 consistent, including fresh installations where the indexes already exist. Cover migration round trips, data preservation, and SQLite query plans with regression tests. --- .../pkg/entity/persistence/agent_run.py | 1 + .../pkg/entity/persistence/monitoring.py | 6 + .../versions/0039_execution_lease_indexes.py | 29 ++++ .../test_execution_lease_indexes.py | 134 ++++++++++++++++++ 4 files changed, 170 insertions(+) create mode 100644 src/langbot/pkg/persistence/alembic/versions/0039_execution_lease_indexes.py create mode 100644 tests/unit_tests/persistence/test_execution_lease_indexes.py diff --git a/src/langbot/pkg/entity/persistence/agent_run.py b/src/langbot/pkg/entity/persistence/agent_run.py index 7d5dc0a8c..af1f9f6b6 100644 --- a/src/langbot/pkg/entity/persistence/agent_run.py +++ b/src/langbot/pkg/entity/persistence/agent_run.py @@ -116,6 +116,7 @@ class AgentRun(Base): ), sqlalchemy.Index('ix_agent_run_runner_status', 'runner_id', 'status'), sqlalchemy.Index('ix_agent_run_queue_claim', 'queue_name', 'status', 'priority', 'id'), + sqlalchemy.Index('ix_agent_run_workspace_status_lease', 'workspace_id', 'status', 'execution_lease_expires_at'), ) diff --git a/src/langbot/pkg/entity/persistence/monitoring.py b/src/langbot/pkg/entity/persistence/monitoring.py index ecf032986..e1a92159e 100644 --- a/src/langbot/pkg/entity/persistence/monitoring.py +++ b/src/langbot/pkg/entity/persistence/monitoring.py @@ -40,6 +40,12 @@ class MonitoringMessage(Base): __table_args__ = ( sqlalchemy.Index('ix_monitoring_messages_workspace_timestamp', 'workspace_uuid', 'timestamp'), sqlalchemy.Index('ix_monitoring_messages_workspace_session', 'workspace_uuid', 'session_id'), + sqlalchemy.Index( + 'ix_monitoring_messages_workspace_status_lease', + 'workspace_uuid', + 'status', + 'execution_lease_expires_at', + ), ) diff --git a/src/langbot/pkg/persistence/alembic/versions/0039_execution_lease_indexes.py b/src/langbot/pkg/persistence/alembic/versions/0039_execution_lease_indexes.py new file mode 100644 index 000000000..4d66427e4 --- /dev/null +++ b/src/langbot/pkg/persistence/alembic/versions/0039_execution_lease_indexes.py @@ -0,0 +1,29 @@ +"""Bound host lease maintenance to active executions within each Workspace.""" + +from alembic import op +import sqlalchemy as sa + +revision = '0039_execution_lease_indexes' +down_revision = '0038_llm_call_details' +branch_labels = None +depends_on = None + + +def upgrade(): + inspector = sa.inspect(op.get_bind()) + for table, workspace_column in ( + ('agent_run', 'workspace_id'), + ('monitoring_messages', 'workspace_uuid'), + ): + index = f'ix_{table}_workspace_status_lease' + # Fresh installations create these indexes from ORM metadata before migrations. + if index not in {item['name'] for item in inspector.get_indexes(table)}: + op.create_index(index, table, [workspace_column, 'status', 'execution_lease_expires_at']) + + +def downgrade(): + inspector = sa.inspect(op.get_bind()) + for table in ('agent_run', 'monitoring_messages'): + index = f'ix_{table}_workspace_status_lease' + if index in {item['name'] for item in inspector.get_indexes(table)}: + op.drop_index(index, table_name=table) diff --git a/tests/unit_tests/persistence/test_execution_lease_indexes.py b/tests/unit_tests/persistence/test_execution_lease_indexes.py new file mode 100644 index 000000000..468871087 --- /dev/null +++ b/tests/unit_tests/persistence/test_execution_lease_indexes.py @@ -0,0 +1,134 @@ +"""Lease maintenance must not scan a Workspace's completed execution history.""" + +import datetime +import importlib + +import pytest +import sqlalchemy as sa +from alembic.migration import MigrationContext +from alembic.operations import Operations +from sqlalchemy.ext.asyncio import create_async_engine + +from langbot.pkg.entity.persistence import user, workspace # noqa: F401 +from langbot.pkg.entity.persistence.agent_run import AgentRun +from langbot.pkg.entity.persistence.base import Base +from langbot.pkg.entity.persistence.monitoring import MonitoringMessage +from langbot.pkg.utils.execution_lease import reconcile_execution_leases + + +MIGRATION = importlib.import_module('langbot.pkg.persistence.alembic.versions.0039_execution_lease_indexes') +LEASE_INDEXES = { + 'agent_run': ('ix_agent_run_workspace_status_lease', ['workspace_id', 'status', 'execution_lease_expires_at']), + 'monitoring_messages': ( + 'ix_monitoring_messages_workspace_status_lease', + ['workspace_uuid', 'status', 'execution_lease_expires_at'], + ), +} + + +def _create_schema(connection, source): + Base.metadata.create_all(connection) + if source == 'migration': + for table, (index, _) in LEASE_INDEXES.items(): + next(item for item in Base.metadata.tables[table].indexes if item.name == index).drop(connection) + with Operations.context(MigrationContext.configure(connection)): + MIGRATION.upgrade() + + +@pytest.mark.parametrize('source', ['metadata', 'migration']) +def test_lease_index_migration_preserves_data_and_existing_indexes(source): + engine = sa.create_engine('sqlite:///:memory:') + try: + with engine.begin() as connection: + _create_schema(connection, source) + connection.execute(sa.insert(AgentRun).values(run_id='history', runner_id='runner', status='completed')) + original_indexes = { + table: {index['name'] for index in sa.inspect(connection).get_indexes(table)} for table in LEASE_INDEXES + } + with Operations.context(MigrationContext.configure(connection)): + MIGRATION.upgrade() + for table, (index, columns) in LEASE_INDEXES.items(): + indexes = {item['name']: item['column_names'] for item in sa.inspect(connection).get_indexes(table)} + assert indexes[index] == columns + assert original_indexes[table] <= indexes.keys() + + MIGRATION.downgrade() + for table, (index, _) in LEASE_INDEXES.items(): + remaining = {item['name'] for item in sa.inspect(connection).get_indexes(table)} + assert remaining == original_indexes[table] - {index} + assert connection.scalar(sa.select(AgentRun.status)) == 'completed' + + MIGRATION.upgrade() + assert connection.scalar(sa.select(AgentRun.status)) == 'completed' + finally: + engine.dispose() + + +@pytest.mark.asyncio +@pytest.mark.parametrize('source', ['metadata', 'migration']) +@pytest.mark.parametrize('model', [AgentRun, MonitoringMessage]) +async def test_lease_queries_use_workspace_status_expiry_index(model, source): + engine = create_async_engine('sqlite+aiosqlite:///:memory:') + now = datetime.datetime(2026, 10, 17) + statements = [] + + def capture_statement(connection, cursor, statement, parameters, context, executemany): + if statement.startswith(('SELECT ', 'UPDATE ')) and model.__tablename__ in statement: + statements.append((statement, parameters)) + + try: + async with engine.begin() as connection: + await connection.run_sync(lambda sync: _create_schema(sync, source)) + common = { + 'execution_owner_id': 'live-host', + 'execution_lease_expires_at': now - datetime.timedelta(seconds=1), + } + if model is AgentRun: + rows = [ + dict( + common, + run_id=str(i), + runner_id='runner', + workspace_id='workspace', + status='completed', + created_at=now, + updated_at=now, + ) + for i in range(2000) + ] + else: + rows = [ + dict( + common, + id=str(i), + workspace_uuid='workspace', + timestamp=now, + bot_id='bot', + bot_name='bot', + pipeline_id='pipeline', + pipeline_name='pipeline', + message_content='Historical message ' * 64, + session_id='session', + status='success', + level='info', + ) + for i in range(2000) + ] + await connection.execute(sa.insert(model), rows) + + sa.event.listen(engine.sync_engine, 'before_cursor_execute', capture_statement) + try: + assert await reconcile_execution_leases(engine, 'workspace', now=now, owner_id='live-host') == 0 + finally: + sa.event.remove(engine.sync_engine, 'before_cursor_execute', capture_statement) + + assert len(statements) == 2 + async with engine.connect() as connection: + for statement, parameters in statements: + plan = (await connection.exec_driver_sql('EXPLAIN QUERY PLAN ' + statement, parameters)).all() + details = ' '.join(row[3] for row in plan) + assert LEASE_INDEXES[model.__tablename__][0] in details, details + assert 'execution_lease_expires_at' in details, details + assert 'USE TEMP B-TREE' not in details, details + finally: + await engine.dispose()