mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-12 18:15:33 +00:00
Merge pull request #2551 from arc53/collisons-in-messages
fix: classify FK violations in message journal to skip pointless seq-retry
This commit is contained in:
2 files changed
+251
-50
No files matched your search
@@ -59,6 +59,27 @@ def _strip_null_bytes(value: Any) -> Any:
|
||||
return value
|
||||
|
||||
|
||||
# Postgres SQLSTATE for a foreign-key violation. ``message_events`` has
|
||||
# an FK to ``conversation_messages(id)``, so a 23503 means that parent
|
||||
# row was never committed — the event can't be journaled and a
|
||||
# ``sequence_no`` retry is pointless churn. A 23505 (PK collision), by
|
||||
# contrast, is recoverable by rewriting at a fresh seq.
|
||||
_FK_VIOLATION_SQLSTATE = "23503"
|
||||
|
||||
|
||||
def _is_foreign_key_violation(exc: IntegrityError) -> bool:
|
||||
"""Return ``True`` when ``exc`` wraps a Postgres FK violation (23503).
|
||||
|
||||
``IntegrityError.orig`` is the underlying driver error, which carries
|
||||
the SQLSTATE on ``.sqlstate`` (psycopg3) or ``.pgcode`` (psycopg2).
|
||||
When neither is present (an error we can't classify) this returns
|
||||
``False`` so the caller keeps its existing seq-retry behavior.
|
||||
"""
|
||||
orig = getattr(exc, "orig", None)
|
||||
sqlstate = getattr(orig, "sqlstate", None) or getattr(orig, "pgcode", None)
|
||||
return sqlstate == _FK_VIOLATION_SQLSTATE
|
||||
|
||||
|
||||
def record_event(
|
||||
message_id: str,
|
||||
sequence_no: int,
|
||||
@@ -108,59 +129,75 @@ def record_event(
|
||||
message_id, sequence_no, event_type, materialised_payload
|
||||
)
|
||||
journal_committed = True
|
||||
except IntegrityError:
|
||||
# Composite-PK collision on (message_id, sequence_no). Most
|
||||
# likely cause is a stale ``latest_sequence_no`` seed on a
|
||||
# continuation retry — the route read MAX(seq) from a separate
|
||||
# connection before another writer committed past it. Look up
|
||||
# the live latest and retry once with latest+1 so the event is
|
||||
# not silently lost. Bounded to a single retry — if two
|
||||
# writers keep racing in lockstep the route-level retry will
|
||||
# converge them across attempts.
|
||||
try:
|
||||
with db_readonly() as conn:
|
||||
latest = MessageEventsRepository(conn).latest_sequence_no(
|
||||
message_id
|
||||
)
|
||||
materialised_seq = (latest if latest is not None else -1) + 1
|
||||
with db_session() as conn:
|
||||
MessageEventsRepository(conn).record(
|
||||
except IntegrityError as exc:
|
||||
if _is_foreign_key_violation(exc):
|
||||
# No ``conversation_messages`` parent row — the event
|
||||
# references a message that was never committed. A
|
||||
# ``sequence_no`` retry can't fix that, so log the real
|
||||
# cause and drop instead of churning the readonly probe +
|
||||
# retry (which would just FK-fail again).
|
||||
logger.warning(
|
||||
"record_event: no conversation_messages row for "
|
||||
"message_id=%s (foreign-key violation); dropping "
|
||||
"(seq=%s type=%s). The parent message row must be "
|
||||
"committed before its events are journaled.",
|
||||
message_id,
|
||||
sequence_no,
|
||||
event_type,
|
||||
)
|
||||
else:
|
||||
# Composite-PK collision on (message_id, sequence_no). Most
|
||||
# likely cause is a stale ``latest_sequence_no`` seed on a
|
||||
# continuation retry — the route read MAX(seq) from a
|
||||
# separate connection before another writer committed past
|
||||
# it. Look up the live latest and retry once with latest+1
|
||||
# so the event is not silently lost. Bounded to a single
|
||||
# retry — if two writers keep racing in lockstep the
|
||||
# route-level retry will converge them across attempts.
|
||||
try:
|
||||
with db_readonly() as conn:
|
||||
latest = MessageEventsRepository(conn).latest_sequence_no(
|
||||
message_id
|
||||
)
|
||||
materialised_seq = (latest if latest is not None else -1) + 1
|
||||
with db_session() as conn:
|
||||
MessageEventsRepository(conn).record(
|
||||
message_id,
|
||||
materialised_seq,
|
||||
event_type,
|
||||
materialised_payload,
|
||||
)
|
||||
journal_committed = True
|
||||
logger.info(
|
||||
"record_event: collision at seq=%s recovered → wrote at "
|
||||
"seq=%s message_id=%s type=%s",
|
||||
sequence_no,
|
||||
materialised_seq,
|
||||
message_id,
|
||||
event_type,
|
||||
)
|
||||
except IntegrityError:
|
||||
# Second collision under the same retry — give up and
|
||||
# log. The route's nonlocal counter will continue at
|
||||
# ``sequence_no+1`` on the next emit; the next call may
|
||||
# land cleanly past the contended window.
|
||||
logger.warning(
|
||||
"record_event: IntegrityError persists after seq+1 "
|
||||
"retry; dropping. message_id=%s original_seq=%s "
|
||||
"retry_seq=%s type=%s",
|
||||
message_id,
|
||||
sequence_no,
|
||||
materialised_seq,
|
||||
event_type,
|
||||
materialised_payload,
|
||||
)
|
||||
journal_committed = True
|
||||
logger.info(
|
||||
"record_event: collision at seq=%s recovered → wrote at "
|
||||
"seq=%s message_id=%s type=%s",
|
||||
sequence_no,
|
||||
materialised_seq,
|
||||
message_id,
|
||||
event_type,
|
||||
)
|
||||
except IntegrityError:
|
||||
# Second collision under the same retry — give up and log.
|
||||
# The route's nonlocal counter will continue at
|
||||
# ``sequence_no+1`` on the next emit; the next call may
|
||||
# land cleanly past the contended window.
|
||||
logger.warning(
|
||||
"record_event: IntegrityError persists after seq+1 retry; "
|
||||
"dropping. message_id=%s original_seq=%s retry_seq=%s "
|
||||
"type=%s",
|
||||
message_id,
|
||||
sequence_no,
|
||||
materialised_seq,
|
||||
event_type,
|
||||
)
|
||||
except Exception:
|
||||
logger.exception(
|
||||
"record_event: retry path failed unexpectedly "
|
||||
"(message_id=%s seq=%s type=%s)",
|
||||
message_id,
|
||||
sequence_no,
|
||||
event_type,
|
||||
)
|
||||
except Exception:
|
||||
logger.exception(
|
||||
"record_event: retry path failed unexpectedly "
|
||||
"(message_id=%s seq=%s type=%s)",
|
||||
message_id,
|
||||
sequence_no,
|
||||
event_type,
|
||||
)
|
||||
except Exception:
|
||||
logger.exception(
|
||||
"message_events INSERT failed: message_id=%s seq=%s type=%s",
|
||||
@@ -297,7 +334,21 @@ class BatchedJournalWriter:
|
||||
MessageEventsRepository(conn).bulk_record(
|
||||
self._message_id, pending
|
||||
)
|
||||
except IntegrityError:
|
||||
except IntegrityError as exc:
|
||||
if _is_foreign_key_violation(exc):
|
||||
# The whole batch references a conversation_messages
|
||||
# parent that doesn't exist — every per-row retry would
|
||||
# FK-fail too, so drop the batch once instead of N
|
||||
# pointless re-attempts. The buffer is already cleared.
|
||||
logger.warning(
|
||||
"BatchedJournalWriter: no conversation_messages row "
|
||||
"for message_id=%s (foreign-key violation); dropping "
|
||||
"%d event(s). The parent message row must be committed "
|
||||
"before its events are journaled.",
|
||||
self._message_id,
|
||||
len(pending),
|
||||
)
|
||||
return
|
||||
logger.info(
|
||||
"BatchedJournalWriter: bulk INSERT collided for "
|
||||
"message_id=%s n=%d; falling back to per-row writes",
|
||||
|
||||
@@ -18,6 +18,21 @@ from application.streaming.message_journal import (
|
||||
)
|
||||
|
||||
|
||||
def _integrity_error(sqlstate: str):
|
||||
"""Build an ``IntegrityError`` whose ``.orig`` reports ``sqlstate``.
|
||||
|
||||
Mirrors how psycopg3 surfaces a Postgres error code on the DBAPI
|
||||
exception that SQLAlchemy wraps: ``23503`` = foreign-key violation
|
||||
(the ``conversation_messages`` parent row is missing), ``23505`` =
|
||||
unique/PK collision (a duplicate ``sequence_no``).
|
||||
"""
|
||||
from sqlalchemy.exc import IntegrityError
|
||||
|
||||
orig = Exception("pg error")
|
||||
orig.sqlstate = sqlstate
|
||||
return IntegrityError("stmt", {}, orig)
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestRecordEvent:
|
||||
def test_invalid_args_return_false(self):
|
||||
@@ -245,6 +260,85 @@ class TestRecordEvent:
|
||||
# serve it on a future reconnect anyway.
|
||||
mock_topic.publish.assert_called_once()
|
||||
|
||||
def test_foreign_key_violation_drops_without_seq_retry(self):
|
||||
"""A 23503 (foreign-key violation) means the
|
||||
``conversation_messages`` parent row was never committed — a
|
||||
``sequence_no`` retry can never land the row, so the write is
|
||||
dropped immediately without probing ``latest_sequence_no`` or
|
||||
re-attempting. The live publish still fires (best-effort tail).
|
||||
"""
|
||||
repo_first = MagicMock(name="repo_first")
|
||||
repo_first.record.side_effect = _integrity_error("23503")
|
||||
repo_readonly = MagicMock(name="repo_readonly")
|
||||
repo_retry = MagicMock(name="repo_retry")
|
||||
|
||||
repo_instances = iter([repo_first, repo_readonly, repo_retry])
|
||||
|
||||
with patch(
|
||||
"application.streaming.message_journal.db_session"
|
||||
) as mock_session, patch(
|
||||
"application.streaming.message_journal.db_readonly"
|
||||
) as mock_readonly, patch(
|
||||
"application.streaming.message_journal.MessageEventsRepository",
|
||||
side_effect=lambda conn: next(repo_instances),
|
||||
), patch(
|
||||
"application.streaming.message_journal.Topic"
|
||||
) as mock_topic_cls:
|
||||
mock_session.return_value.__enter__.return_value = MagicMock()
|
||||
mock_readonly.return_value.__enter__.return_value = MagicMock()
|
||||
mock_topic = MagicMock()
|
||||
mock_topic_cls.return_value = mock_topic
|
||||
|
||||
result = record_event("msg-1", 3, "answer", {"text": "hi"})
|
||||
|
||||
assert result is False
|
||||
repo_first.record.assert_called_once_with(
|
||||
"msg-1", 3, "answer", {"text": "hi"}
|
||||
)
|
||||
# No seq+1 retry: the readonly probe and retry insert never run.
|
||||
mock_readonly.assert_not_called()
|
||||
repo_readonly.latest_sequence_no.assert_not_called()
|
||||
repo_retry.record.assert_not_called()
|
||||
# Live tail still gets the frame at the original seq.
|
||||
mock_topic.publish.assert_called_once()
|
||||
wire = json.loads(mock_topic.publish.call_args[0][0])
|
||||
assert wire["sequence_no"] == 3
|
||||
|
||||
def test_pk_collision_still_retries_with_seq_plus_one(self):
|
||||
"""An explicit 23505 (unique/PK collision) keeps the existing
|
||||
recover-by-retry behavior: probe ``latest_sequence_no`` and
|
||||
rewrite at latest+1. Regression guard pinning the FK-vs-PK split.
|
||||
"""
|
||||
repo_first = MagicMock(name="repo_first")
|
||||
repo_first.record.side_effect = _integrity_error("23505")
|
||||
repo_readonly = MagicMock(name="repo_readonly")
|
||||
repo_readonly.latest_sequence_no.return_value = 7
|
||||
repo_retry = MagicMock(name="repo_retry")
|
||||
|
||||
repo_instances = iter([repo_first, repo_readonly, repo_retry])
|
||||
|
||||
with patch(
|
||||
"application.streaming.message_journal.db_session"
|
||||
) as mock_session, patch(
|
||||
"application.streaming.message_journal.db_readonly"
|
||||
) as mock_readonly, patch(
|
||||
"application.streaming.message_journal.MessageEventsRepository",
|
||||
side_effect=lambda conn: next(repo_instances),
|
||||
), patch(
|
||||
"application.streaming.message_journal.Topic"
|
||||
) as mock_topic_cls:
|
||||
mock_session.return_value.__enter__.return_value = MagicMock()
|
||||
mock_readonly.return_value.__enter__.return_value = MagicMock()
|
||||
mock_topic_cls.return_value = MagicMock()
|
||||
|
||||
result = record_event("msg-1", 3, "answer", {"text": "hi"})
|
||||
|
||||
assert result is True
|
||||
repo_readonly.latest_sequence_no.assert_called_once_with("msg-1")
|
||||
repo_retry.record.assert_called_once_with(
|
||||
"msg-1", 8, "answer", {"text": "hi"}
|
||||
)
|
||||
|
||||
def test_payload_with_null_byte_stripped_before_insert(self):
|
||||
"""``\\x00`` in a string value would crash JSONB INSERT — strip
|
||||
at the journal boundary so the row lands and reconnecting
|
||||
@@ -522,6 +616,62 @@ class TestBatchedJournalWriter:
|
||||
assert per_row_repo.record.call_args_list[0].args[1] == 0
|
||||
assert per_row_repo.record.call_args_list[1].args[1] == 1
|
||||
|
||||
def test_bulk_foreign_key_violation_drops_batch_without_per_row(self):
|
||||
"""A 23503 (foreign-key violation) on the bulk INSERT means the
|
||||
whole batch references a ``conversation_messages`` parent that
|
||||
doesn't exist — every per-row retry would FK-fail too. The
|
||||
batch is dropped once without the per-row fallback (no latest
|
||||
probe, no re-attempts, no publish), and the buffer is cleared.
|
||||
"""
|
||||
session_cm, readonly_cm, repo_cls_p, topic_cls_p = self._patch_io()
|
||||
with session_cm as mock_session, readonly_cm as mock_readonly, repo_cls_p as mock_repo_cls, topic_cls_p as mock_topic_cls:
|
||||
mock_session.return_value.__enter__.return_value = MagicMock()
|
||||
mock_readonly.return_value.__enter__.return_value = MagicMock()
|
||||
|
||||
bulk_repo = MagicMock(name="bulk_repo")
|
||||
bulk_repo.bulk_record.side_effect = _integrity_error("23503")
|
||||
per_row_repo = MagicMock(name="per_row_repo")
|
||||
mock_repo_cls.side_effect = [bulk_repo, per_row_repo]
|
||||
|
||||
mock_topic = MagicMock()
|
||||
mock_topic_cls.return_value = mock_topic
|
||||
|
||||
writer = BatchedJournalWriter("msg-1", batch_size=2)
|
||||
writer.record(0, "answer", {"text": "a"})
|
||||
writer.record(1, "answer", {"text": "b"})
|
||||
|
||||
bulk_repo.bulk_record.assert_called_once()
|
||||
# No per-row fallback: the parent is missing for the whole batch.
|
||||
mock_readonly.assert_not_called()
|
||||
per_row_repo.record.assert_not_called()
|
||||
mock_topic.publish.assert_not_called()
|
||||
# Buffer cleared so close() is a no-op and memory stays bounded.
|
||||
assert writer._buffer == []
|
||||
|
||||
def test_bulk_collision_23505_still_falls_back_to_per_row(self):
|
||||
"""An explicit 23505 (unique/PK collision) keeps the per-row
|
||||
fallback — the parent exists, only a ``sequence_no`` collided.
|
||||
Regression guard for the FK-vs-PK split at the bulk boundary.
|
||||
"""
|
||||
session_cm, readonly_cm, repo_cls_p, topic_cls_p = self._patch_io()
|
||||
with session_cm as mock_session, readonly_cm as mock_readonly, repo_cls_p as mock_repo_cls, topic_cls_p as mock_topic_cls:
|
||||
mock_session.return_value.__enter__.return_value = MagicMock()
|
||||
mock_readonly.return_value.__enter__.return_value = MagicMock()
|
||||
|
||||
bulk_repo = MagicMock(name="bulk_repo")
|
||||
bulk_repo.bulk_record.side_effect = _integrity_error("23505")
|
||||
per_row_repo = MagicMock(name="per_row_repo")
|
||||
mock_repo_cls.side_effect = [bulk_repo, per_row_repo, per_row_repo]
|
||||
|
||||
mock_topic_cls.return_value = MagicMock()
|
||||
|
||||
writer = BatchedJournalWriter("msg-1", batch_size=2)
|
||||
writer.record(0, "answer", {"text": "a"})
|
||||
writer.record(1, "answer", {"text": "b"})
|
||||
|
||||
bulk_repo.bulk_record.assert_called_once()
|
||||
assert per_row_repo.record.call_count == 2
|
||||
|
||||
def test_flush_clears_buffer_even_on_total_failure(self):
|
||||
"""A flush that fails with a non-IntegrityError exception must
|
||||
still clear the buffer — leaving rows would grow memory
|
||||
|
||||
Reference in new issue
Block a user