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.
This commit is contained in:
BiFangKNT
2026-10-17 19:41:55 +08:00
parent c05aac538f
commit 1318efc248
4 changed files with 170 additions and 0 deletions
@@ -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'),
)
@@ -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',
),
)
@@ -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)
@@ -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()