Aiven-Open / Aiven-Open/pghoard

Startup rescan can enqueue a duplicate UploadEvent for the same file, corrupting UploadEventProgressTracker and causing spurious upload failures/retries

Offen
#659 0 Kommentare 0 Reaktionen 0 zugewiesene Personen Auf GitHub ansehen
Vorherrschende Sprache
Python
Sterne
1.4k
Forks
111
Ø Merge
1 T. 5 Std.
Gemergte PRs (30 T.)
2

Beschreibung

## Related Issues/PRs

None found. Searched (via `gh issue list` / `gh pr list --search`, open and closed): `untrack tracking upload progress race`, `UploadEventProgressTracker`, `not being tracked`, `race condition`, `concurrent upload`, `duplicate transfer`. Only unrelated hits (one closed issue about `tar` during restoration, unrelated to upload tracking). `UploadEventProgressTracker` was introduced in #605 (merged 2024-01-23) with no follow-up reports against it since.

## Description

`UploadEventProgressTracker` (`pghoard/transfer.py`) assumes at most one in-flight upload attempt is ever tracked per `file_key` at a time. If two attempts for the *same* `file_key` end up tracked concurrently, one attempt's cleanup silently destroys the other's still-live tracking state, because none of the tracker's methods check *which* attempt owns the entry — only whether the key is present:

```python
# transfer.py:136-139
def track_upload_event(self, file_key: str, file_type: FileType, file_size: Optional[int]) -> None:
with self._tracked_events_lock:
self.log.debug("Tracking upload event for file %s", file_key)
self._tracked_events[file_key] = UploadEventProgress(key=file_key, file_type=file_type, file_size=file_size)

# transfer.py:141-146
def untrack_upload_event(self, file_key: str) -> None:
if file_key not in self._tracked_events:
return
with self._tracked_events_lock:
self._tracked_events.pop(file_key)
```

`track_upload_event` unconditionally **overwrites** whatever is currently stored under `file_key`, and `untrack_upload_event` unconditionally **removes** whatever is currently stored under `file_key` — neither has any per-attempt identity/token, so if attempt A's cleanup runs while attempt B's tracking for the same `file_key` is live, A's cleanup deletes B's entry.

### Concrete downstream failure

`TransferAgent.handle_upload` uses the tracker via a context manager (`track_upload_event`, `transfer.py:235-244`) around the actual upload:

```python
# transfer.py:480-493 (abridged)
with file_ctx as f, \
track_upload_event(
progress_tracker=self.upload_tracker, file_key=key, upload_event=file_to_transfer,
):
...
upload_progress_fn = partial(self.upload_tracker.increment, file_key=key)
storage.store_file_object(
key, f, metadata=metadata,
upload_progress_fn=lambda n_bytes: upload_progress_fn(total_bytes_uploaded=n_bytes),
)
```

If a second, unrelated attempt for the same `file_key` untracks (context-manager exit) while this upload is still genuinely in progress, the very next `upload_progress_fn` callback from `storage.store_file_object` calls `self.upload_tracker.increment(file_key=key, ...)`, which raises:

```python
# transfer.py:151-153
with self._tracked_events_lock:
if file_key not in self._tracked_events:
raise Exception(f"UploadEvent for {file_key} is not being tracked.")
```

That exception propagates out of `storage.store_file_object` and is caught by `handle_upload`'s broad `except Exception` — a perfectly healthy, still-succeeding upload is treated as a failure and needlessly retried (`retry_number` incremented, backoff sleep, re-enqueued).

### How two attempts for the same `file_key` end up tracked concurrently — a concrete, ordinary-conditions trigger, not a contrived edge case

1. `PGHoard.run()` starts thread pools *before* the startup rescan runs:
```python
# pghoard.py:978-980
def run(self):
self.start_threads_on_startup()
self.startup_walk_for_missed_files()
```
`start_threads_on_startup()` starts the compressor thread pool and the `TransferAgent` thread pool (size `config["transfer"]["thread_count"]`, all sharing one `UploadEventProgressTracker` instance), so both are live and draining their queues before the rescan below even begins.

2. `startup_walk_for_missed_files()` (`pghoard.py:672`) has two loops per backup site. Loop 1 handles leftover *uncompressed* WAL files (an ordinary occurrence after any restart while WAL streaming was active) and enqueues a `CompressionEvent` for each into `compression_queue`:
```python
# pghoard.py:749
self.compression_queue.put(compression_event)
```
This queue is already being actively drained by the live compressor threads from step 1. A compressor thread can pick this up, finish compressing, and — as a normal, real-time side effect of finishing any compression job, not just at startup — immediately enqueue the resulting `UploadEvent`:
```python
# compressor.py:289
self.transfer_queue.put(transfer_object)
```

3. Loop 2 of the *same* rescan function then unconditionally re-scans the compressed-file directory and builds a **second**, independent `UploadEvent` for anything it finds there:
```python
# pghoard.py:752-753, 780
for compressed_file_dir in [compressed_xlog_path, compressed_timeline_path]:
for filename in os.listdir(compressed_file_dir):
...
self.transfer_queue.put(transfer_event)
```
with **no check** for whether the file it just found was already picked up and forwarded by the compressor in step 2 — even though loop 1's own enqueue is what triggered that compression in the first place. Notably, the surrounding code shows the team is aware of exactly this failure mode and guarded against it for the *uncompressed→compressed* transition:
```python
# pghoard.py:711-712
# verify if file was already compressed, otherwise the transfer agent will encounter
# duplicated UploadEvents. In case it was compressed, we should just add it to the deletion queue
```
(`is_already_compressed` check at `pghoard.py:717`) — but the equivalent guard is missing for loop 2 against the compressor's real-time output racing the rescan.

4. `TransferAgent.run_safe()` (`transfer.py:330-337`) has no deduplication before dispatch — it is a plain `while self.running: file_to_transfer = self.transfer_queue.get(...); ...; handle_upload(...)` loop per thread, pulling from the shared queue.

The only condition needed for two `UploadEvent`s for the same file to reach two different `TransferAgent` threads is: **compression of a leftover WAL file (loop 1) finishes before loop 2 reaches `os.listdir` on the same directory.** With any leftover WAL backlog after a restart, this is an ordinary ordering, not a narrow adversarial race.

## Reproduction Steps / Example Code (Python)

The tracker-level corruption (item under "Concrete downstream failure" above) reproduces deterministically through the real public `track_upload_event` context manager — the same one `TransferAgent.handle_upload` uses — with two threads simulating two concurrent attempts for the same `file_key` (as would arise from the duplicate-enqueue path described above):

```python
import threading
import time
from pghoard.transfer import UploadEventProgressTracker, track_upload_event
from pghoard.common import FileType

class FakeMetrics:
def gauge(self, *a, **k): pass
def increase(self, *a, **k): pass
def unexpected_exception(self, *a, **k): pass

class FakeUploadEvent:
def __init__(self, file_type, file_size):
self.file_type = file_type
self.file_size = file_size

tracker = UploadEventProgressTracker(metrics=FakeMetrics())
file_key = "site/xlog/000000010000000000000042"
barrier = threading.Barrier(2)
result = {}

def thread_a():
# Older attempt for this file_key (e.g. the compressor's real-time enqueue).
with track_upload_event(tracker, file_key, FakeUploadEvent(FileType.Wal, 100)):
barrier.wait()
time.sleep(0.05)
# __exit__ -> untrack_upload_event(file_key) fires here

def thread_b():
# Newer, still-genuinely-in-progress attempt for the SAME file_key
# (e.g. the startup rescan's independent re-discovery of the same file).
barrier.wait()
time.sleep(0.01)
with track_upload_event(tracker, file_key, FakeUploadEvent(FileType.Wal, 200)):
time.sleep(0.08) # A's untrack fires here, mid-upload for B
with tracker._tracked_events_lock:
result["b_entry_survived"] = file_key in tracker._tracked_events

ta, tb = threading.Thread(target=thread_a), threading.Thread(target=thread_b)
ta.start(); tb.start(); ta.join(); tb.join()
print("B's live tracking entry survived A's unrelated untrack:", result["b_entry_survived"])
```

## Actual output (current `main`, 30/30 runs identical)

```
B's live tracking entry survived A's unrelated untrack: False
```

A's `untrack_upload_event` (for its own, unrelated, already-finished attempt) removes B's entry while B's upload is still genuinely in progress. Extending the repro so B calls `tracker.increment(file_key=file_key, total_bytes_uploaded=1024)` immediately after (exactly what `storage.store_file_object`'s `upload_progress_fn` callback does inside `handle_upload`) reproduces the concrete downstream failure 15/15 runs:

```
Exception: UploadEvent for site/xlog/000000010000000000000042 is not being tracked.
```

Note: `increment()` calls `PersistedProgress.read()` unconditionally at its start (`transfer.py:149`); see the second, separate finding below for why that call currently raises on its own, unrelated to this bug. The repro above patches `PersistedProgress` out (via `unittest.mock.patch`) to isolate this specific claim from that one.

**Expected:** two concurrent attempts for the same `file_key` should not be able to interfere with each other's tracked progress — each attempt's `untrack` should only ever remove *its own* entry, and a still-in-progress upload's `increment()` call should never fail because of an unrelated attempt's cleanup.

**Actual:** any attempt's `untrack_upload_event(file_key)` removes whichever entry currently exists under that key, regardless of which attempt created it, causing the still-live attempt to fail its own `increment()` call and, in real usage via `handle_upload`, be treated as a failed upload and needlessly retried.

---

## Second, separate finding: `PersistedProgress` is uninstantiable under the currently-installable `pydantic`

Found while isolating the repro above. Unrelated root cause to the tracking bug — documenting here rather than as a separate issue since both were found investigating the same file pair (`transfer.py`/`common.py`) in the same pass, but they are two different bugs, not one.

`pghoard/common.py:141-143`:
```python
class PersistedProgress(BaseModel):
progress: Dict[str, ProgressData] = Field(default_factory=dict)
_lock: threading.Lock = threading.Lock()
```

`pghoard/common.py:28` imports plain `from pydantic import BaseModel, Field` — pydantic v2's native `BaseModel`, under which a leading-underscore class attribute like `_lock` is auto-converted to a `PrivateAttr`. Pydantic v2 resolves/copies private-attribute defaults on every instantiation, and `threading.Lock` objects cannot be pickled or deep-copied, so this fails with a `TypeError` on **every** instantiation path, regardless of the public `progress` field's data:

```python
>>> from pghoard.common import PersistedProgress
>>> PersistedProgress()
TypeError: cannot pickle '_thread.lock' object
>>> PersistedProgress.parse_raw('{"progress": {}}')
TypeError: cannot pickle '_thread.lock' object
```

Verified both routes independently, isolated (no threading involved — this is a plain, deterministic instantiation failure, reproduces every time):
- `PersistedProgress()` — the no-`PROGRESS_FILE` path (`common.py:154`)
- `PersistedProgress.parse_raw()` — the file-exists path (`common.py:150`)

`parse_raw` also emits, before the `TypeError`:
```
PydanticDeprecatedSince20: The `parse_raw` method is deprecated; if your data is JSON use `model_validate_json`, otherwise load the data then use `model_validate` instead. Deprecated in Pydantic V2.0 to be removed in V3.0.
PydanticDeprecatedSince20: `load_str_bytes` is deprecated. Deprecated in Pydantic V2.0 to be removed in V3.0.
```
— further evidence this class was written for pydantic v1 and not updated for v2.

**Why this is live now:** `pyproject.toml:31` leaves the runtime `pydantic` dependency unconstrained (`"pydantic",`), while `pyproject.toml:112` pins `"pydantic==1.10.14"` in the `constraints` list. But `rohmu` (a required runtime dependency) imports `pydantic.v1` in several modules (e.g. `rohmu/delta/common.py:10`, `rohmu/common/models.py:7`, `rohmu/object_storage/config.py:10`) — the `pydantic.v1` compatibility shim only exists inside actual pydantic v2 packages, so installing `rohmu` requires pydantic v2 regardless of what the `constraints` pin says. Installing this repo fresh today (`pip install -e .`) resolves `pydantic==2.13.4`, and under that version `PersistedProgress` is broken by the mechanism above.

**Severity, stated carefully:** `PersistedProgress.read()` is the unconditional first line of `increment()` (`transfer.py:148-149`), before any file-type check. If this dependency resolution holds in real production installs too, this would mean every `increment()` call — i.e. every upload-progress callback, for every file type — fails here before reaching any of the logic the first finding is about. I want to be careful not to overstate this: I don't know how pghoard is actually deployed in practice (e.g. whether some installs still pin an old pydantic 1.x environment some other way) — I only know what the currently-committed `pyproject.toml` resolves to on a fresh install today.

### Reproduction (isolated, no threading)

```python
from pghoard.common import PersistedProgress

# Route 1: no PROGRESS_FILE on disk
PersistedProgress() # TypeError: cannot pickle '_thread.lock' object

# Route 2: PROGRESS_FILE exists with valid content
PersistedProgress.parse_raw('{"progress": {}}') # same TypeError, plus two
# PydanticDeprecatedSince20 warnings
```

## System Info

```
Aiven-Open/pghoard, commit 8565adbcd360007ca2c0f818e8aee749584013fc (main, verified 2026-07-10)
python: 3.14
pydantic: 2.13.4 (resolved by a fresh `pip install -e .` today)
rohmu: 2.8.2
```

Beitragsleitfaden

Für dieses Repository ist kein Beitragsleitfaden indexiert

Bewertung

Dieses Issue wurde noch nicht bewertet.

Neue Issues direkt in Ihr Postfach

Eine kurze Übersicht über anfängerfreundliche GitHub-Issues.