fix: index the chat queries that make large SQLite instances unusable (#27663)
The timer scheduler polls once a second and cancels on every message send and chat open, the sidebar lists chats ordered by `updated_at`, and the folder badges count unread chats per folder. None of those could be served by an index, so each call read most of the `chat` table, and because `meta` sits after the chat payload column SQLite had to walk every row's overflow pages to get there. On a large history that stalls the sidebar, every chat switch and every send, and the idle poll alone burns about a quarter of a CPU core. Timers now keep their due time in a dedicated `chat.timer_at` column behind a partial index, and the chat list, unread and unfinished-reply queries each get an index matching their filter and ordering. Existing pending timers are backfilled from their meta by the migration. Dropping the `internal` and `type` checks also makes a forked timer chat inert, where a fork used to copy `meta` verbatim and become a second claim target that could fire a duplicate timer. Measured on SQLite, same rows returned: | query | before | after | |---|---|---| | idle timer poll (2000 chats, 0.43 GB) | 170 ms | 0.04 ms | | cancel on send and chat open (4000 chats, 377 MB) | 200 ms | 0.04 ms | | sidebar chat list (15000 chats, 1.26 GB) | 157 ms | 1.8 ms | | folder unread badges (15000 chats, 1.4 GB) | 54 ms | 0.2 ms | PostgreSQL 17 serves all of them as index-only scans with no sort node. Exercised through fresh install, upgrade with seeded data, downgrade and re-upgrade on SQLite and PostgreSQL 17. Fixes #27622
This commit is contained in:
+70
@@ -0,0 +1,70 @@
|
||||
"""add chat timer_at and chat list, unread and timer indexes
|
||||
|
||||
Revision ID: d4c1a8e37b62
|
||||
Revises: 6d09d1bf1f23
|
||||
Create Date: 2026-08-23 18:05:12.441907
|
||||
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
import sqlalchemy as sa
|
||||
from alembic import op
|
||||
|
||||
# revision identifiers, used by Alembic.
|
||||
revision: str = 'd4c1a8e37b62'
|
||||
down_revision: str | None = '6d09d1bf1f23'
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
op.add_column('chat', sa.Column('timer_at', sa.BigInteger(), nullable=True))
|
||||
op.create_index(
|
||||
'timer_at_idx',
|
||||
'chat',
|
||||
['timer_at'],
|
||||
sqlite_where=sa.text('timer_at IS NOT NULL'),
|
||||
postgresql_where=sa.text('timer_at IS NOT NULL'),
|
||||
)
|
||||
|
||||
# Timers created before this migration carry their due time in meta only, and would never fire.
|
||||
chat = sa.table(
|
||||
'chat', sa.column('id', sa.String), sa.column('meta', sa.JSON), sa.column('timer_at', sa.BigInteger)
|
||||
)
|
||||
conn = op.get_bind()
|
||||
pending = conn.execute(
|
||||
sa.select(chat.c.id, chat.c.meta)
|
||||
.where(chat.c.meta['type'].as_string() == 'timer')
|
||||
.where(chat.c.meta['status'].as_string() == 'pending')
|
||||
).all()
|
||||
for chat_id, meta in pending:
|
||||
try: # imported chats can carry any meta, and a non-numeric due time must not abort the migration
|
||||
due_at = int(meta.get('timer_at'))
|
||||
except (TypeError, ValueError):
|
||||
continue
|
||||
conn.execute(chat.update().where(chat.c.id == chat_id).values(timer_at=due_at))
|
||||
|
||||
op.create_index('user_id_updated_at_id_idx', 'chat', ['user_id', sa.text('updated_at DESC'), 'id'])
|
||||
op.create_index(
|
||||
'user_id_timer_at_idx',
|
||||
'chat',
|
||||
['user_id', 'timer_at'],
|
||||
sqlite_where=sa.text('timer_at IS NOT NULL'),
|
||||
postgresql_where=sa.text('timer_at IS NOT NULL'),
|
||||
)
|
||||
op.create_index(
|
||||
'user_id_folder_unread_idx',
|
||||
'chat',
|
||||
['user_id', 'folder_id', 'archived', 'updated_at', 'last_read_at', 'id'],
|
||||
)
|
||||
op.create_index('chat_message_chat_role_done_idx', 'chat_message', ['chat_id', 'role', 'done'])
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
op.drop_index('chat_message_chat_role_done_idx', table_name='chat_message')
|
||||
op.drop_index('user_id_folder_unread_idx', table_name='chat')
|
||||
op.drop_index('user_id_timer_at_idx', table_name='chat')
|
||||
op.drop_index('user_id_updated_at_id_idx', table_name='chat')
|
||||
op.drop_index('timer_at_idx', table_name='chat')
|
||||
op.drop_column('chat', 'timer_at')
|
||||
@@ -168,6 +168,7 @@ class ChatMessage(Base):
|
||||
Index('chat_message_chat_parent_idx', 'chat_id', 'parent_id'),
|
||||
Index('chat_message_model_created_idx', 'model_id', 'created_at'),
|
||||
Index('chat_message_user_created_idx', 'user_id', 'created_at'),
|
||||
Index('chat_message_chat_role_done_idx', 'chat_id', 'role', 'done'), # unfinished-assistant probe
|
||||
)
|
||||
|
||||
|
||||
|
||||
@@ -149,6 +149,7 @@ class Chat(Base): # database table mapping for chat entity
|
||||
current_message_id = Column(Text, nullable=True)
|
||||
|
||||
last_read_at = Column(BigInteger, nullable=True)
|
||||
timer_at = Column(BigInteger, nullable=True) # ns due time, set only while a timer chat waits to be claimed
|
||||
|
||||
__table_args__ = (
|
||||
# Performance indexes for common queries
|
||||
@@ -157,6 +158,23 @@ class Chat(Base): # database table mapping for chat entity
|
||||
Index('user_id_archived_idx', 'user_id', 'archived'),
|
||||
Index('updated_at_user_id_idx', 'updated_at', 'user_id'),
|
||||
Index('folder_id_user_id_idx', 'folder_id', 'user_id'),
|
||||
Index('user_id_updated_at_id_idx', 'user_id', updated_at.desc(), 'id'),
|
||||
Index(
|
||||
'timer_at_idx',
|
||||
'timer_at',
|
||||
sqlite_where=text('timer_at IS NOT NULL'),
|
||||
postgresql_where=text('timer_at IS NOT NULL'),
|
||||
),
|
||||
# timer_at key column turns the IS NOT NULL into a seek, so this beats the plain user_id indexes
|
||||
Index(
|
||||
'user_id_timer_at_idx',
|
||||
'user_id',
|
||||
'timer_at',
|
||||
sqlite_where=text('timer_at IS NOT NULL'),
|
||||
postgresql_where=text('timer_at IS NOT NULL'),
|
||||
),
|
||||
# covering index: lets SQLite serve count_unread_by_folder_ids without reading chat rows
|
||||
Index('user_id_folder_unread_idx', 'user_id', 'folder_id', 'archived', 'updated_at', 'last_read_at', 'id'),
|
||||
)
|
||||
|
||||
|
||||
@@ -187,6 +205,7 @@ class ChatModel(BaseModel):
|
||||
current_message_id: str | None = None
|
||||
|
||||
last_read_at: int | None = None
|
||||
timer_at: int | None = None
|
||||
|
||||
@field_validator('variables', mode='before')
|
||||
@classmethod
|
||||
@@ -502,6 +521,7 @@ class ChatTable:
|
||||
db: AsyncSession | None = None,
|
||||
*,
|
||||
internal_meta: dict | None = None,
|
||||
timer_at: int | None = None,
|
||||
) -> ChatModel | None:
|
||||
async with get_async_db_context(db) as session:
|
||||
chat = ChatModel(
|
||||
@@ -514,6 +534,7 @@ class ChatTable:
|
||||
'chat': self._clean_null_bytes(form_data.chat),
|
||||
'folder_id': form_data.folder_id,
|
||||
'meta': internal_meta or {},
|
||||
'timer_at': timer_at,
|
||||
'variables': form_data.variables or {},
|
||||
'current_message_id': self.get_current_message_id(form_data.chat),
|
||||
'created_at': int(time.time()),
|
||||
|
||||
@@ -152,6 +152,7 @@ async def create_timer(
|
||||
'cancel_on': selected_events,
|
||||
'run': run,
|
||||
},
|
||||
timer_at=due_at,
|
||||
)
|
||||
if not chat:
|
||||
return 'Error: failed to create timer.'
|
||||
@@ -171,20 +172,18 @@ async def claim_due_timers(now_ns: int, limit: int = 10) -> list[tuple[str, str]
|
||||
async with get_async_db() as db:
|
||||
stmt = (
|
||||
select(Chat)
|
||||
.where(Chat.meta['internal'].as_boolean().is_(True))
|
||||
.where(Chat.meta['type'].as_string() == 'timer')
|
||||
.where(Chat.timer_at <= now_ns)
|
||||
.where(Chat.meta['status'].as_string() == 'pending')
|
||||
.order_by(Chat.timer_at)
|
||||
.limit(limit)
|
||||
)
|
||||
if db.bind.dialect.name == 'postgresql':
|
||||
stmt = stmt.with_for_update(skip_locked=True)
|
||||
|
||||
result = await db.execute(stmt)
|
||||
rows = [row for row in result.scalars().all() if int((row.meta or {}).get('timer_at') or 0) <= now_ns]
|
||||
rows.sort(key=lambda row: int((row.meta or {}).get('timer_at') or 0))
|
||||
rows = rows[:limit]
|
||||
|
||||
claimed = []
|
||||
for row in rows:
|
||||
for row in result.scalars().all():
|
||||
claim_id = str(uuid4())
|
||||
row.meta = {
|
||||
**(row.meta or {}),
|
||||
@@ -192,6 +191,7 @@ async def claim_due_timers(now_ns: int, limit: int = 10) -> list[tuple[str, str]
|
||||
'timer_started_at': now_ns,
|
||||
'timer_claim_id': claim_id,
|
||||
}
|
||||
row.timer_at = None
|
||||
row.updated_at = int(time.time())
|
||||
claimed.append((row.id, claim_id))
|
||||
await db.commit()
|
||||
@@ -205,9 +205,8 @@ async def cancel_timers_for_chat(
|
||||
async with get_async_db() as db:
|
||||
result = await db.execute(
|
||||
select(Chat)
|
||||
.where(Chat.timer_at.isnot(None))
|
||||
.where(Chat.user_id == user_id)
|
||||
.where(Chat.meta['internal'].as_boolean().is_(True))
|
||||
.where(Chat.meta['type'].as_string() == 'timer')
|
||||
.where(Chat.meta['parent_chat_id'].as_string() == parent_chat_id)
|
||||
.where(Chat.meta['status'].as_string() == 'pending')
|
||||
)
|
||||
@@ -222,6 +221,7 @@ async def cancel_timers_for_chat(
|
||||
'timer_cancelled_at': now_ns,
|
||||
'timer_cancelled_by': event,
|
||||
}
|
||||
row.timer_at = None
|
||||
row.updated_at = int(time.time())
|
||||
await db.commit()
|
||||
|
||||
@@ -289,12 +289,14 @@ async def execute_due_timer(app, timer_id: str, claim_id: str | None = None) ->
|
||||
if await has_active_tasks(app.state.redis, parent_chat_id):
|
||||
timer_row = await db.get(Chat, timer_id)
|
||||
if timer_row:
|
||||
timer_meta = timer_row.meta or {}
|
||||
timer_row.meta = {
|
||||
**(timer_row.meta or {}),
|
||||
**timer_meta,
|
||||
'status': 'pending',
|
||||
'timer_claim_id': None,
|
||||
'timer_started_at': None,
|
||||
}
|
||||
timer_row.timer_at = timer_meta.get('timer_at')
|
||||
timer_row.updated_at = int(time.time())
|
||||
await db.commit()
|
||||
return
|
||||
|
||||
Reference in New Issue
Block a user