fix(plugins): stop delivering telemetry events twice, and stop losing parked ones (#7324)
This commit is contained in:
@@ -100,6 +100,16 @@ BATCH_SIZE = 100
|
||||
SEND_TIMEOUT = 5
|
||||
CLAIM_STALE_SECONDS = 120
|
||||
CLAIM_EXPIRY_SECONDS = 7 * 24 * 60 * 60
|
||||
# A batch is only discarded once it has genuinely been retried this many times.
|
||||
MAX_CLAIM_ATTEMPTS = 3
|
||||
# Parked claims drained per run, after the live spool. Bounded so a long backlog
|
||||
# cannot turn one flush into an unbounded send loop.
|
||||
MAX_PARKED_PER_RUN = 3
|
||||
# Added to the wait before a released claim becomes reclaimable, per attempt
|
||||
# already spent. Releasing straight to "reclaimable now" let two senders burn the
|
||||
# whole budget within seconds of one another on a single momentary failure, and
|
||||
# discard a batch a retry a minute later would have delivered.
|
||||
RETRY_COOLDOWN_SECONDS = 60
|
||||
|
||||
|
||||
def is_enabled() -> bool:
|
||||
@@ -399,38 +409,201 @@ def spawn_flush() -> bool:
|
||||
return False
|
||||
|
||||
|
||||
def _claim_name(attempt: int = 0) -> str:
|
||||
"""Claim filename. The attempt count rides in the name so the 7-day expiry
|
||||
only ever discards a batch that was actually retried and failed."""
|
||||
return f"telemetry-{os.getpid()}-{uuid.uuid4().hex[:8]}-a{attempt}.sending"
|
||||
|
||||
|
||||
def _claim_attempt(claim: Path) -> int:
|
||||
"""Attempts recorded in a claim filename; 0 for the pre-attempt-count shape.
|
||||
|
||||
Anchored on field position, not on a leading "a": the legacy shape is
|
||||
``telemetry-<pid>-<hex>.sending`` and a hex id such as ``a1234567`` would
|
||||
otherwise parse as attempt 1234567 and be discarded unsent on the first
|
||||
flush after an upgrade.
|
||||
"""
|
||||
stem = claim.name[: -len(".sending")] if claim.name.endswith(".sending") else claim.name
|
||||
parts = stem.split("-")
|
||||
if len(parts) != 4:
|
||||
return 0
|
||||
tail = parts[3]
|
||||
if tail.startswith("a") and tail[1:].isdigit():
|
||||
return int(tail[1:])
|
||||
return 0
|
||||
|
||||
|
||||
def _touch(path: Path) -> None:
|
||||
"""Refresh mtime so a claim's age measures time since it was claimed.
|
||||
|
||||
``Path.replace`` is ``os.rename``, which preserves mtime — so a claim created
|
||||
after a quiet minute inherited the spool's last-write time and looked
|
||||
abandoned the instant it was made. A second sender would then take it over
|
||||
while the first was still posting, and both would deliver the batch.
|
||||
"""
|
||||
try:
|
||||
os.utime(path, None)
|
||||
except OSError:
|
||||
pass
|
||||
|
||||
|
||||
def _claim_spool() -> Path | None:
|
||||
"""Rename the spool aside so exactly one sender owns each batch."""
|
||||
directory = memory_core.data_dir()
|
||||
claim = directory / f"telemetry-{os.getpid()}-{uuid.uuid4().hex[:8]}.sending"
|
||||
claim = directory / _claim_name()
|
||||
spool = _spool_path()
|
||||
try:
|
||||
spool.replace(claim)
|
||||
_touch(claim)
|
||||
return claim
|
||||
except OSError:
|
||||
pass
|
||||
return _claim_parked(directory)
|
||||
|
||||
|
||||
def _sweep_debris(directory: Path) -> None:
|
||||
"""Remove files nothing else will ever pick up again.
|
||||
|
||||
*.partial is a temp file orphaned by a crash between write and rename.
|
||||
*.corrupt is a batch quarantined for undecodable content. No glob in this
|
||||
module matches either, so without this they accumulate on disk for the life
|
||||
of the install.
|
||||
|
||||
Quarantined batches are kept far longer than debris: they are the only
|
||||
evidence left of events that could not be delivered, and someone diagnosing
|
||||
a report of missing telemetry has to be able to find one.
|
||||
"""
|
||||
now = time.time()
|
||||
for orphan in sorted(directory.glob("telemetry-*.sending")):
|
||||
for debris in directory.glob("telemetry-*.partial"):
|
||||
try:
|
||||
if now - debris.stat().st_mtime > CLAIM_STALE_SECONDS:
|
||||
debris.unlink()
|
||||
except OSError:
|
||||
continue
|
||||
for quarantined in directory.glob("telemetry-*.corrupt"):
|
||||
try:
|
||||
if now - quarantined.stat().st_mtime > CLAIM_EXPIRY_SECONDS:
|
||||
quarantined.unlink()
|
||||
except OSError:
|
||||
continue
|
||||
# The same reasoning covers *.tmp. _write_identity and _install_salt both
|
||||
# create one and unlink it in a finally, which a SIGKILL skips, and no glob
|
||||
# in this module matches the leftovers either.
|
||||
for temporary in directory.glob("telemetry-*.tmp"):
|
||||
try:
|
||||
if now - temporary.stat().st_mtime > CLAIM_STALE_SECONDS:
|
||||
temporary.unlink()
|
||||
except OSError:
|
||||
continue
|
||||
|
||||
|
||||
def _claim_parked(directory: Path) -> Path | None:
|
||||
"""Take the oldest abandoned claim, if any lease has actually expired.
|
||||
|
||||
Kept separate from the live spool so flush() can drain both in one run.
|
||||
Previously parked batches were only reachable when no spool existed at all,
|
||||
and because sessions keep recording there usually was one — so a batch
|
||||
parked by a failed send waited until the 7-day expiry deleted it unsent,
|
||||
even though its own presence is what started the sender.
|
||||
"""
|
||||
now = time.time()
|
||||
for orphan in sorted(directory.glob("telemetry-*.sending"), key=_safe_mtime):
|
||||
try:
|
||||
age = now - orphan.stat().st_mtime
|
||||
except OSError:
|
||||
continue
|
||||
if age > CLAIM_EXPIRY_SECONDS:
|
||||
if age < CLAIM_STALE_SECONDS:
|
||||
# Someone else holds a live lease on it. This check has to come
|
||||
# first. Claiming a file bumps its attempt count and refreshes its
|
||||
# mtime, so a sender that has just taken the final attempt looks
|
||||
# exhausted to everyone else while it is actively draining. Judging
|
||||
# exhaustion before liveness let a second sender unlink a batch out
|
||||
# from under its owner, losing every event in it.
|
||||
continue
|
||||
# Attempts, not age. Every re-claim touches the mtime and every release
|
||||
# backdates it by a fixed amount, so age is pinned near the stale
|
||||
# threshold and never reaches the expiry. Age stays only as a backstop
|
||||
# for files that never carried an attempt marker.
|
||||
if _claim_attempt(orphan) >= MAX_CLAIM_ATTEMPTS or age > CLAIM_EXPIRY_SECONDS:
|
||||
try:
|
||||
orphan.unlink()
|
||||
except OSError:
|
||||
pass
|
||||
continue
|
||||
if age < CLAIM_STALE_SECONDS:
|
||||
continue
|
||||
claim = orphan.parent / _claim_name(_claim_attempt(orphan) + 1)
|
||||
try:
|
||||
orphan.replace(claim)
|
||||
_touch(claim)
|
||||
return claim
|
||||
except OSError:
|
||||
continue
|
||||
return None
|
||||
|
||||
|
||||
def _safe_mtime(path: Path) -> float:
|
||||
try:
|
||||
return path.stat().st_mtime
|
||||
except OSError:
|
||||
return 0.0
|
||||
|
||||
|
||||
def _rewrite_claim(claim: Path, remaining: list[dict[str, Any]]) -> bool:
|
||||
"""Persist the unsent remainder, atomically, and refresh the lease.
|
||||
|
||||
Called after every successful batch. Two jobs: a retry resumes where the
|
||||
send stopped instead of re-posting from the top, and the rewrite doubles as
|
||||
the lease heartbeat, so a slow sender does not have its claim stolen
|
||||
mid-flight. Interval is one batch, well inside CLAIM_STALE_SECONDS.
|
||||
"""
|
||||
if not remaining:
|
||||
try:
|
||||
claim.unlink()
|
||||
except OSError:
|
||||
pass
|
||||
return True
|
||||
temporary = claim.with_suffix(f".{os.getpid()}.partial")
|
||||
try:
|
||||
payload = "".join(json.dumps(event, separators=(",", ":"), default=str) + "\n" for event in remaining)
|
||||
# fsync before the rename: without it the rename can land while the
|
||||
# bytes have not, and the claim comes back empty or truncated after a
|
||||
# crash. _drain then reads zero events and unlinks it.
|
||||
with open(temporary, "w", encoding="utf-8") as handle:
|
||||
handle.write(payload)
|
||||
handle.flush()
|
||||
os.fsync(handle.fileno())
|
||||
temporary.replace(claim)
|
||||
_touch(claim)
|
||||
return True
|
||||
except OSError:
|
||||
try:
|
||||
temporary.unlink()
|
||||
except OSError:
|
||||
pass
|
||||
return False
|
||||
|
||||
|
||||
def _release_claim(claim: Path, remaining: list[dict[str, Any]]) -> None:
|
||||
"""Persist the remainder and drop the lease, because this sender has given up.
|
||||
|
||||
Distinct from the per-batch heartbeat: heartbeating on the way out would
|
||||
make an abandoned batch look actively owned for a further
|
||||
CLAIM_STALE_SECONDS, delaying the retry for no reason. Ageing it past the
|
||||
threshold lets the next flush pick it up immediately, while the attempt
|
||||
count in the filename still bounds how many times that can happen.
|
||||
"""
|
||||
if not _rewrite_claim(claim, remaining):
|
||||
return
|
||||
try:
|
||||
# Backdate past the stale threshold so the next flush can pick it up,
|
||||
# minus a cooldown that grows with the attempts already spent. Clamped so
|
||||
# the mtime never lands in the future, which would read as a live lease.
|
||||
cooldown = min(_claim_attempt(claim) * RETRY_COOLDOWN_SECONDS, CLAIM_STALE_SECONDS)
|
||||
released = time.time() - CLAIM_STALE_SECONDS - 1 + cooldown
|
||||
os.utime(claim, (released, released))
|
||||
except OSError:
|
||||
pass
|
||||
|
||||
|
||||
def _resolve_email(key: str) -> str:
|
||||
"""Trade the API key for the account email so events join other Mem0 surfaces."""
|
||||
url = os.environ.get("MEM0_API_URL", memory_core.DEFAULT_API_URL).rstrip("/") + "/v1/ping/"
|
||||
@@ -478,16 +651,61 @@ def resolve_distinct_id() -> tuple[str, str]:
|
||||
|
||||
|
||||
def flush() -> int:
|
||||
"""Drain claimed spools to PostHog and return the number of events sent."""
|
||||
"""Drain the live spool, then any parked claims, and return events sent."""
|
||||
if not is_enabled():
|
||||
return 0
|
||||
claim = _claim_spool()
|
||||
sent, delivered = _drain(_claim_spool())
|
||||
if not delivered:
|
||||
# The network is failing. Retrying other batches now would only burn
|
||||
# their attempt budget against the same broken connection.
|
||||
return sent
|
||||
|
||||
# Parked batches used to starve behind the live spool indefinitely. Bounded
|
||||
# per run so a long backlog cannot turn one flush into an unbounded loop.
|
||||
directory = memory_core.data_dir()
|
||||
_sweep_debris(directory)
|
||||
for _ in range(MAX_PARKED_PER_RUN):
|
||||
parked = _claim_parked(directory)
|
||||
if parked is None:
|
||||
break
|
||||
count, delivered = _drain(parked)
|
||||
sent += count
|
||||
if not delivered:
|
||||
break
|
||||
return sent
|
||||
|
||||
|
||||
def _drain(claim: Path | None) -> tuple[int, bool]:
|
||||
"""Post one claimed batch file, recording progress after every batch.
|
||||
|
||||
Returns (events sent, whether everything was delivered).
|
||||
"""
|
||||
if claim is None:
|
||||
return 0
|
||||
return 0, True
|
||||
try:
|
||||
lines = claim.read_text(encoding="utf-8").splitlines()
|
||||
except ValueError:
|
||||
# UnicodeDecodeError from a torn write: the content is unrecoverable, so
|
||||
# quarantine rather than retry. flush() runs from a bare `finally:` in
|
||||
# flush_worker, so raising here also skips the handoff cleanup, and an
|
||||
# undecodable file would otherwise be re-read on every flush forever.
|
||||
# Reported as delivered because there is nothing left to deliver and the
|
||||
# rest of the run should continue.
|
||||
try:
|
||||
claim.replace(claim.with_suffix(".corrupt"))
|
||||
except OSError:
|
||||
try:
|
||||
claim.unlink()
|
||||
except OSError:
|
||||
pass
|
||||
return 0, True
|
||||
except OSError:
|
||||
return 0
|
||||
# Could not read it, which is not the same as having nothing to send.
|
||||
# The file is left exactly where it is: a vanished or briefly unreadable
|
||||
# claim is retryable, and quarantining it here would discard events over
|
||||
# a transient filesystem error. Reported as undelivered so the run stops
|
||||
# instead of counting a batch nothing was posted from as delivered.
|
||||
return 0, False
|
||||
events = []
|
||||
for line in lines:
|
||||
try:
|
||||
@@ -497,11 +715,18 @@ def flush() -> int:
|
||||
if isinstance(value, dict) and value.get("event"):
|
||||
events.append(value)
|
||||
if not events:
|
||||
# Only delete when the file really is empty. A non-empty file that
|
||||
# parses to nothing is a torn write, and its contents are the unsent
|
||||
# remainder — deleting it is the data loss this PR exists to prevent.
|
||||
try:
|
||||
claim.unlink()
|
||||
empty = claim.stat().st_size == 0
|
||||
except OSError:
|
||||
empty = True
|
||||
try:
|
||||
claim.replace(claim.with_suffix(".corrupt")) if not empty else claim.unlink()
|
||||
except OSError:
|
||||
pass
|
||||
return 0
|
||||
return 0, True
|
||||
|
||||
distinct_id, aliased_anonymous_id = resolve_distinct_id()
|
||||
if aliased_anonymous_id:
|
||||
@@ -520,10 +745,13 @@ def flush() -> int:
|
||||
|
||||
sent = 0
|
||||
for start in range(0, len(events), BATCH_SIZE):
|
||||
chunk = events[start : start + BATCH_SIZE]
|
||||
batch = [
|
||||
{
|
||||
"event": event["event"],
|
||||
"distinct_id": distinct_id,
|
||||
# Carried through from record() so a resend can be collapsed.
|
||||
"uuid": event.get("uuid"),
|
||||
"timestamp": event.get("timestamp"),
|
||||
"properties": {
|
||||
# Fallback only: events recorded by a build before source
|
||||
@@ -535,16 +763,24 @@ def flush() -> int:
|
||||
**(event.get("properties") or {}),
|
||||
},
|
||||
}
|
||||
for event in events[start : start + BATCH_SIZE]
|
||||
for event in chunk
|
||||
]
|
||||
if not _post({"api_key": POSTHOG_API_KEY, "batch": batch}, POSTHOG_BATCH_URL):
|
||||
return sent
|
||||
sent += len(batch)
|
||||
try:
|
||||
claim.unlink()
|
||||
except OSError:
|
||||
pass
|
||||
return sent
|
||||
# Keep only what has not been delivered, and release the lease.
|
||||
# Previously the whole file was kept and the retry re-posted every
|
||||
# batch, including the ones that had already arrived.
|
||||
_release_claim(claim, events[start:])
|
||||
return sent, False
|
||||
sent += len(chunk)
|
||||
# Record progress and refresh the lease after each successful batch, so
|
||||
# a crash repeats at most one batch instead of the entire file. If the
|
||||
# rewrite fails the claim still holds delivered events, so stop rather
|
||||
# than carry on as though progress were recorded — continuing is how the
|
||||
# duplicate delivery this PR fixes would come back.
|
||||
if not _rewrite_claim(claim, events[start + len(chunk) :]):
|
||||
_release_claim(claim, events[start + len(chunk) :])
|
||||
return sent, False
|
||||
return sent, True
|
||||
|
||||
|
||||
def main() -> int:
|
||||
|
||||
@@ -199,9 +199,11 @@ def test_a_stale_claim_is_reclaimed(isolated_env, monkeypatch):
|
||||
telemetry.record("search")
|
||||
orphan = telemetry._claim_spool()
|
||||
assert orphan is not None
|
||||
monkeypatch.setattr(
|
||||
telemetry.time, "time", lambda: orphan.stat().st_mtime + telemetry.CLAIM_STALE_SECONDS + 1
|
||||
)
|
||||
# Frozen rather than re-stat'd per call: flush() drains the live spool and
|
||||
# then looks for parked claims in the same run, so by the second look this
|
||||
# file no longer exists.
|
||||
stale_now = orphan.stat().st_mtime + telemetry.CLAIM_STALE_SECONDS + 1
|
||||
monkeypatch.setattr(telemetry.time, "time", lambda: stale_now)
|
||||
|
||||
with patch.object(telemetry, "_post", lambda payload, url: True):
|
||||
assert telemetry.flush() == 1
|
||||
@@ -211,9 +213,13 @@ def test_an_expired_claim_is_dropped(isolated_env, monkeypatch):
|
||||
telemetry.record("search")
|
||||
orphan = telemetry._claim_spool()
|
||||
assert orphan is not None
|
||||
monkeypatch.setattr(
|
||||
telemetry.time, "time", lambda: orphan.stat().st_mtime + telemetry.CLAIM_EXPIRY_SECONDS + 1
|
||||
)
|
||||
expired_now = orphan.stat().st_mtime + telemetry.CLAIM_EXPIRY_SECONDS + 1
|
||||
monkeypatch.setattr(telemetry.time, "time", lambda: expired_now)
|
||||
# Expiry now only discards a batch that was genuinely retried and failed,
|
||||
# so age alone is not enough — age it past the attempt budget too.
|
||||
retried = orphan.parent / orphan.name.replace("-a0.", f"-a{telemetry.MAX_CLAIM_ATTEMPTS}.")
|
||||
orphan.replace(retried)
|
||||
os.utime(retried, (expired_now, expired_now - telemetry.CLAIM_EXPIRY_SECONDS - 1))
|
||||
assert telemetry._claim_spool() is None
|
||||
assert not list(memory_core.data_dir().glob("telemetry-*.sending"))
|
||||
|
||||
|
||||
Reference in New Issue
Block a user