Skip to content

chore: optimize index task callbacks #964

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
iziang merged 2 commits into from
Jun 20, 2025
+83 −93
Merged

chore: optimize index task callbacks #964

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
6934494
chore: optimize index task callbacks
iziang Jun 19, 2025
2d868fe
Merge remote-tracking branch 'origin/main' into bugfix/document-index…
iziang Jun 20, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
51 changes: 0 additions & 51 deletions aperag/db/models.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -727,19 +727,6 @@ class DocumentIndex(Base):
def __repr__(self):
return f"<DocumentIndex(id={self.id}, document_id={self.document_id}, type={self.index_type}, desired={self.desired_state}, actual={self.actual_state})>"

def needs_reconciliation(self) -> bool:
"""Check if this index needs reconciliation"""
# Check version first - if observed_version < version, need reconciliation
if self.observed_version < self.version:
return True

# Then check state consistency
if self.desired_state == IndexDesiredState.PRESENT:
return self.actual_state in [IndexActualState.ABSENT, IndexActualState.FAILED]
elif self.desired_state == IndexDesiredState.ABSENT:
return self.actual_state in [IndexActualState.CREATING, IndexActualState.PRESENT]
return False

def is_in_sync(self) -> bool:
"""Check if desired and actual states are in sync"""
if self.observed_version < self.version:
Expand All @@ -759,41 +746,3 @@ def update_spec(self, desired_state: IndexDesiredState = None, created_by: str =
self.created_by = created_by
self.version += 1
self.gmt_updated = utc_now()

def mark_creating(self):
"""Mark as creating"""
self.actual_state = IndexActualState.CREATING
self.gmt_updated = utc_now()
self.gmt_last_reconciled = utc_now()

def mark_present(self, index_data: str = None):
"""Mark as present (successfully created)"""
self.actual_state = IndexActualState.PRESENT
self.observed_version = self.version # Mark as processed
self.error_message = None
if index_data:
self.index_data = index_data
self.gmt_updated = utc_now()
self.gmt_last_reconciled = utc_now()

def mark_deleting(self):
"""Mark as deleting"""
self.actual_state = IndexActualState.DELETING
self.gmt_updated = utc_now()
self.gmt_last_reconciled = utc_now()

def mark_absent(self):
"""Mark as absent (successfully deleted)"""
self.actual_state = IndexActualState.ABSENT
self.observed_version = self.version # Mark as processed
self.index_data = None
self.error_message = None
self.gmt_updated = utc_now()
self.gmt_last_reconciled = utc_now()

def mark_failed(self, error_message: str):
"""Mark as failed"""
self.actual_state = IndexActualState.FAILED
self.error_message = error_message
self.gmt_updated = utc_now()
self.gmt_last_reconciled = utc_now()
123 changes: 83 additions & 40 deletions aperag/index/reconciler.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,7 @@
IndexDesiredState,
)
from aperag.tasks.scheduler import TaskScheduler, create_task_scheduler
from aperag.utils.utils import utc_now

logger = logging.getLogger(__name__)

Expand Down Expand Up @@ -208,7 +209,11 @@ def _claim_document_indexes(self, session: Session, document_id: str, indexes_to
update_stmt = (
update(DocumentIndex)
.where(and_(*where_conditions))
.values(actual_state=target_state, gmt_updated=func.now(), gmt_last_reconciled=func.now())
.values(
actual_state=target_state,
gmt_updated=utc_now(),
gmt_last_reconciled=utc_now()
)
)

result = session.execute(update_stmt)
Expand Down Expand Up @@ -290,61 +295,99 @@ def _update_document_status(document_id: str, session: Session):
def on_index_created(document_id: str, index_type: str, index_data: str = None):
"""Called when index creation succeeds"""
for session in get_sync_session():
stmt = select(DocumentIndex).where(
and_(
DocumentIndex.document_id == document_id,
DocumentIndex.index_type == DocumentIndexType(index_type),
# Use atomic update with state validation
update_stmt = (
update(DocumentIndex)
.where(
and_(
DocumentIndex.document_id == document_id,
DocumentIndex.index_type == DocumentIndexType(index_type),
DocumentIndex.actual_state == IndexActualState.CREATING, # Only allow transition from CREATING
)
)
.values(
actual_state=IndexActualState.PRESENT,
observed_version=DocumentIndex.version, # Mark as processed
index_data=index_data,
error_message=None,
gmt_updated=utc_now(),
gmt_last_reconciled=utc_now(),
)
)
result = session.execute(stmt)
doc_index = result.scalar_one_or_none()

if doc_index:
doc_index.mark_present(index_data)
IndexTaskCallbacks._update_document_status(document_id, session)

logger.info(f"{index_type} index creation completed for document {document_id}")
session.commit()

result = session.execute(update_stmt)
if result.rowcount > 0:
IndexTaskCallbacks._update_document_status(document_id, session)
logger.info(f"{index_type} index creation completed for document {document_id}")
session.commit()
else:
logger.warning(f"Index creation callback ignored for document {document_id} type {index_type} - not in CREATING state")
session.rollback()

@staticmethod
def on_index_failed(document_id: str, index_type: str, error_message: str):
"""Called when index operation fails"""
for session in get_sync_session():
stmt = select(DocumentIndex).where(
and_(
DocumentIndex.document_id == document_id,
DocumentIndex.index_type == DocumentIndexType(index_type),
# Use atomic update with state validation
update_stmt = (
update(DocumentIndex)
.where(
and_(
DocumentIndex.document_id == document_id,
DocumentIndex.index_type == DocumentIndexType(index_type),
# Only allow transition from CREATING or DELETING states
DocumentIndex.actual_state.in_([IndexActualState.CREATING, IndexActualState.DELETING]),
)
)
.values(
actual_state=IndexActualState.FAILED,
error_message=error_message,
gmt_updated=utc_now(),
gmt_last_reconciled=utc_now(),
)
)
result = session.execute(stmt)
doc_index = result.scalar_one_or_none()

if doc_index:
doc_index.mark_failed(error_message)
IndexTaskCallbacks._update_document_status(document_id, session)

logger.error(f"{index_type} index operation failed for document {document_id}: {error_message}")
session.commit()

result = session.execute(update_stmt)
if result.rowcount > 0:
IndexTaskCallbacks._update_document_status(document_id, session)
logger.error(f"{index_type} index operation failed for document {document_id}: {error_message}")
session.commit()
else:
logger.warning(f"Index failure callback ignored for document {document_id} type {index_type} - not in CREATING or DELETING state")
session.rollback()

@staticmethod
def on_index_deleted(document_id: str, index_type: str):
"""Called when index deletion succeeds"""
for session in get_sync_session():
stmt = select(DocumentIndex).where(
and_(
DocumentIndex.document_id == document_id,
DocumentIndex.index_type == DocumentIndexType(index_type),
# Use atomic update with state validation
update_stmt = (
update(DocumentIndex)
.where(
and_(
DocumentIndex.document_id == document_id,
DocumentIndex.index_type == DocumentIndexType(index_type),
DocumentIndex.actual_state == IndexActualState.DELETING, # Only allow transition from DELETING
)
)
.values(
actual_state=IndexActualState.ABSENT,
observed_version=DocumentIndex.version, # Mark as processed
index_data=None,
error_message=None,
gmt_updated=utc_now(),
gmt_last_reconciled=utc_now(),
)
)
result = session.execute(stmt)
doc_index = result.scalar_one_or_none()

if doc_index:
doc_index.mark_absent()
IndexTaskCallbacks._update_document_status(document_id, session)

logger.info(f"{index_type} index deletion completed for document {document_id}")
session.commit()

result = session.execute(update_stmt)
if result.rowcount > 0:
IndexTaskCallbacks._update_document_status(document_id, session)
logger.info(f"{index_type} index deletion completed for document {document_id}")
session.commit()
else:
logger.warning(f"Index deletion callback ignored for document {document_id} type {index_type} - not in DELETING state")
session.rollback()


# Global instance
Expand Down
1 change: 0 additions & 1 deletion docs/indexing_architecture.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -354,7 +354,6 @@ class IndexTaskCallbacks:
def on_index_failed(document_id: str, index_type: str, error_message: str):
"""Task failure callback"""
# Update database state
doc_index.mark_failed(error_message)
doc_index.actual_state = IndexActualState.FAILED
doc_index.error_message = error_message

Expand Down
1 change: 0 additions & 1 deletion docs/indexing_architecture_zh.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -354,7 +354,6 @@ class IndexTaskCallbacks:
def on_index_failed(document_id: str, index_type: str, error_message: str):
"""任务失败回调"""
# 更新数据库状态
doc_index.mark_failed(error_message)
doc_index.actual_state = IndexActualState.FAILED
doc_index.error_message = error_message

Expand Down
Loading