diff --git a/.github/workflows/slo.yml b/.github/workflows/slo.yml index 8c066a3d4..b5a1fe8fb 100644 --- a/.github/workflows/slo.yml +++ b/.github/workflows/slo.yml @@ -10,37 +10,132 @@ permissions: checks: write jobs: - ydb-slo-action: + select-scenarios: # On `labeled` events run only when the `SLO` label itself was just added — # otherwise unrelated label changes (e.g. the AI-review bot toggling # `ai_review_in_process` / `ai_reviewed`) would spawn a fresh run that # cancels the in-progress one via `cancel-in-progress`. For the other # trigger types keep gating on the `SLO` label being present. + # + # The gate lives only here: the workload job `needs` this one, so a skipped + # gate skips the whole run exactly as before. if: >- (github.event.action == 'labeled' && github.event.label.name == 'SLO') || (github.event.action != 'labeled' && contains(github.event.pull_request.labels.*.name, 'SLO')) + name: Select SLO scenarios + runs-on: ubuntu-latest + outputs: + sdk: ${{ steps.select.outputs.sdk }} + + steps: + # A full SLO run is long and burns external runners, so only the scenarios that can + # actually be affected by the diff are started. The scenario list lives here (and only + # here) so the map and the definitions cannot drift apart. + - name: Pick scenarios affected by the changed files + id: select + env: + GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + PR_NUMBER: ${{ github.event.pull_request.number }} + REPO: ${{ github.repository }} + run: | + set -euo pipefail + + ALL_SCENARIOS='[ + {"name":"sync-table","command":"--read-rps 1000 --write-rps 100"}, + {"name":"sync-query","command":"--read-rps 1000 --write-rps 100"}, + {"name":"async-query","command":"--read-rps 1000 --write-rps 100"}, + {"name":"sync-topic","command":"--write-rps 200 --write-threads 8 --read-threads 8", + "metrics_yaml_path":"sdk-current/tests/slo/metrics-topic.yaml", + "thresholds_yaml_path":"sdk-current/tests/slo/thresholds-topic.yaml"}, + {"name":"async-topic","command":"--write-rps 200 --write-threads 8 --read-threads 8", + "metrics_yaml_path":"sdk-current/tests/slo/metrics-topic.yaml", + "thresholds_yaml_path":"sdk-current/tests/slo/thresholds-topic.yaml"}, + {"name":"sync-topic-multiwriter", + "command":"--write-rps 200 --write-threads 2 --keys-per-writer 16 --read-threads 8", + "metrics_yaml_path":"sdk-current/tests/slo/metrics-topic.yaml", + "thresholds_yaml_path":"sdk-current/tests/slo/thresholds-topic.yaml"} + ]' + + CHANGED=$(gh api "repos/${REPO}/pulls/${PR_NUMBER}/files" --paginate --jq '.[].filename') + echo "Changed files:" + echo "${CHANGED}" | sed 's/^/ /' + + TOPIC=0; QUERY=0; TABLE=0; EVERYTHING=0 + while IFS= read -r file; do + [ -n "${file}" ] || continue + case "${file}" in + # Order matters: the first matching pattern wins, so the specific service + # paths have to be tested before the `ydb/*` catch-all below. + ydb/_topic_reader/*|ydb/_topic_writer/*|ydb/_topic_common/*| \ + ydb/topic.py|ydb/aio/topic.py|ydb/_grpc/grpcwrapper/ydb_topic*.py) + TOPIC=1 ;; + ydb/query/*|ydb/aio/query/*) + QUERY=1 ;; + ydb/table.py|ydb/aio/table.py|ydb/_session_impl.py) + TABLE=1 ;; + # Harness parts that belong to one service only. + tests/slo/src/jobs/*topic*|tests/slo/src/runners/topic_runner.py| \ + tests/slo/metrics-topic.yaml|tests/slo/thresholds-topic.yaml) + TOPIC=1 ;; + # One runner drives sync-table, sync-query and async-query alike. + tests/slo/src/jobs/*table*|tests/slo/src/runners/table_runner.py) + TABLE=1; QUERY=1 ;; + # Prose about the harness changes no behaviour. + tests/slo/*.md) + ;; + # Shared harness (options, runners entry point, metrics, image) or this workflow: + # every scenario is affected. + tests/slo/*|.github/workflows/slo.yml) + EVERYTHING=1 ;; + # Anything else inside the SDK is shared machinery (driver, pool, connection, + # retries, credentials, generated stubs, convert/types used by more than one + # service) — assume it can move any scenario. + ydb/*) + EVERYTHING=1 ;; + # Docs, examples, packaging, other CI: no SLO impact. + *) ;; + esac + done <> "$GITHUB_OUTPUT" + + ydb-slo-action: + needs: select-scenarios name: Run YDB SLO Tests runs-on: "large-runner-python-sdk" strategy: fail-fast: false matrix: - sdk: - - name: sync-table - command: "--read-rps 1000 --write-rps 100" - - name: sync-query - command: "--read-rps 1000 --write-rps 100" - - name: async-query - command: "--read-rps 1000 --write-rps 100" - - name: sync-topic - command: "--write-rps 200 --write-threads 8 --read-threads 8" - metrics_yaml_path: sdk-current/tests/slo/metrics-topic.yaml - thresholds_yaml_path: sdk-current/tests/slo/thresholds-topic.yaml - - name: async-topic - command: "--write-rps 200 --write-threads 8 --read-threads 8" - metrics_yaml_path: sdk-current/tests/slo/metrics-topic.yaml - thresholds_yaml_path: sdk-current/tests/slo/thresholds-topic.yaml + sdk: ${{ fromJSON(needs.select-scenarios.outputs.sdk) }} concurrency: group: slo-${{ github.ref }}-${{ matrix.sdk.name }} diff --git a/CHANGELOG.md b/CHANGELOG.md index fc442de70..ea286477d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,3 +1,5 @@ +* Add a topic multi-partition writer (`topic_client.multiwriter(...)`) that routes messages across partitions by their `key`, with Kafka-hash and key-range partition choosers, and transparently resends in-flight messages to child partitions on an auto-partition split (no loss, no duplicates); expose partition `key_range` on `describe_topic` results + ## 3.31.4 ## * Fixed async `QuerySessionPool` permanently losing a pool slot when `acquire()` was cancelled while a new session was being created: `asyncio.CancelledError` no longer leaks the pool size counter, so a pool under deadline-driven cancellations can no longer end up exhausted and blocking forever. A cancelled or interrupted session attach now also closes the session instead of orphaning it server-side diff --git a/MULTIWRITER_ARCHITECTURE.md b/MULTIWRITER_ARCHITECTURE.md new file mode 100644 index 000000000..cf3a87874 --- /dev/null +++ b/MULTIWRITER_ARCHITECTURE.md @@ -0,0 +1,791 @@ +# Topic multi-partition writer — architecture + +How write-by-key is built, why each piece is the way it is, and what breaks if it is built +differently. Written from the Python implementation, but deliberately language-agnostic: it +describes the protocol interactions and the state machine, not the API. + +Cross-checked against the two existing implementations — C++ `TProducer` +(`src/client/topic/impl/producer.{h,cpp}`) and Go `topicmultiwriter` +(`internal/topic/topicmultiwriter`). Differences are called out where they exist. + +--- + +## 1. What it is + +A **multi-writer** is one logical writer that spreads messages over all partitions of a topic +while keeping per-key ordering. The caller attaches a **key** to each message; the writer decides +which partition owns that key and writes there. + +The reason it exists is server-side auto-partitioning: the server splits a partition when it gets +hot. A plain writer targets one partition and simply dies when that partition goes away. The +multi-writer follows the topology instead. + +``` + write(key="user-42", data=...) + │ + ▼ + ┌───────────────────────┐ + │ orchestrator │ routing, seqno, in-flight ownership, + │ │ repartition handling + └───────────┬───────────┘ + route by key │ one sub-writer per partition + ┌──────────────┬──────┴───────┬──────────────┐ + ▼ ▼ ▼ ▼ + sub-writer sub-writer sub-writer sub-writer + partition 0 partition 1 partition 2 partition N + producer producer producer producer + "pfx-0" "pfx-1" "pfx-2" "pfx-N" + │ │ │ │ + ▼ ▼ ▼ ▼ + independent write stream per partition + (own buffering, encoding, reconnect, token refresh) +``` + +The sub-writer is an **ordinary single-partition writer**, unmodified. The orchestrator adds +routing, sequence numbering and repartition handling on top. This is the single most important +structural decision: buffering, compression, reconnection and auth refresh are not reimplemented. + +> Why this lives in the client at all. Routing by key was originally meant to be a server feature: +> the client would attach a `message_group_id` to each message and the server would spread the +> groups across partitions itself. The protocol still carries that plan — `message_group_id` exists +> both as a session-level setting and, in `MessageData.partitioning`, as a per-message one, which +> only makes sense if the server was going to do the routing. It never shipped, so every SDK does +> the routing on the client instead, and the multi-writer is that work. +> +> The practical consequence: **treat `message_group_id` as another name for the producer id.** Every +> implementation sets the two to the same value, nothing routes by it, and the key you attach to a +> message has nothing to do with it. It is not a second identity to reason about. + +--- + +## 2. The contract + +What a caller may rely on: + +| # | Guarantee | +|---|---| +| 1 | All messages with a given key go to one partition, or to its descendants after a split | +| 2 | Order is preserved per key | +| 3 | Every accepted message reaches a terminal state: acknowledged or failed. Never silently dropped, never left pending | +| 4 | A split is invisible: no loss, and no duplicates for messages the server had already persisted | +| 5 | A split does not require the caller to recreate the writer | +| 6 | Transient errors are retried internally | +| 7 | Closing with flush delivers everything already accepted | + +What it deliberately does **not** guarantee: + +- No global order across different keys. Only per-key. + +--- + +## 3. Internal invariants + +These are the properties every part of the implementation is protecting. If you are porting this, +these are the things to write tests against. + +1. **One key → one branch of the partition tree.** A key may travel `p0 → p2 → p5`, but must never + appear in two siblings. Violating this breaks per-key ordering for readers. +2. **The routing view never has a gap.** Every key must map to exactly one live partition at every + moment, including mid-split. +3. **A seqno identifies a message for its whole life.** It is assigned once and never changes, even + when the message moves to another partition. +4. **A message is in exactly one place:** in flight under exactly one partition, or resolved. +5. **The dedup cut only ever comes from the server.** See §9.4 — this is the subtlest part. + +--- + +## 4. State + +Everything below is owned by the orchestrator and guarded by a single lock (§11). + +| State | Shape | Purpose | +|---|---|---| +| routing view | `partition_id → partition info` | live leaves only; what the chooser routes to | +| sub-writers | `partition_id → writer` | lazily created, idle-evicted | +| in-flight | `partition_id → {seqno → message}` | messages the orchestrator may still have to resend | +| seqno cursor | single integer | one sequence for the whole writer (§7) | +| init seqno | `partition_id → int` | server's persisted seqno when the current sub-writer opened | +| retired seqno | `partition_id → int` | server's final persisted seqno for a retired producer; cacheable | +| max acked | `partition_id → int` | highest ack actually observed | +| retiring | set of partition ids | suppresses expected ack failures during teardown | +| repartition tasks | `partition_id → task` | coalescing + shutdown ownership | +| last write time | `partition_id → timestamp` | idle eviction | + +The routing view holds live leaves only: a partition is dropped from it the moment it is retired, +or keys would keep being sent to it. Nothing here remembers retired partitions, and nothing needs +to — see §9.4 for why the dedup cut never looks past the partition being retired. + +--- + +## 5. Routing + +### 5.1 Key ranges + +**The server owns the ranges; the client only reads them.** They arrive in the `DescribeTopic` +response, one per partition: + +``` +DescribeTopicResult.PartitionInfo.key_range : PartitionKeyRange + optional bytes from_bound // inclusive left border, empty = -inf + optional bytes to_bound // exclusive right border, empty = +inf +``` + +Nothing on the client side derives, splits or interpolates a range. It reads what describe reports, +hashes the key, and picks the partition whose range contains the hash. A client that computed +ranges itself would immediately disagree with the server about which key belongs where. + +Note the shape: **one interval per partition, not a set of them.** There is no way to express a +partition that owns two disjoint stretches of the key space. That single fact answers most +questions about what topology changes are even possible — see §9.6 for what it means for merges. + +Bounds are opaque byte strings compared **bytewise, lexicographically**. No locale, no encoding +awareness, no numeric interpretation. Observed live: bounds are short — a real split produced +2-byte bounds like `efea`, `f7f5`. + +``` +before split p0: ["", "") ← one partition owns everything + +after split p0: ["", "") inactive, children [1,2] + p1: ["", m) active + p2: [m, "") active + +key hashes to "apple" → p1 +key hashes to "zebra" → p2 +key hashes to "m" → p2 ← left bound inclusive, right exclusive +``` + +### 5.2 The key is hashed before comparison + +> Two different hashes appear in this document, for two different routing schemes. They are not +> alternatives and they are not a contradiction: +> +> | Scheme | Hash | Used for | +> |---|---|---| +> | range-based (§5.4) | **MurmurHash64A**, 64-bit | placing a key inside `[from_bound, to_bound)` | +> | hash-modulo (§5.5) | **murmur2, 32-bit** | `% partition_count` on a fixed partition count | +> +> The range-based one is the YDB path and is the one that has to match other writers. The 32-bit +> one exists only to reproduce Apache Kafka's partitioner and never touches key ranges. + +For the range-based scheme the client does not compare the raw key against the bounds. It hashes +it: + +``` +partition_key = MurmurHash64A(key_bytes, seed=0) → 8 bytes, big-endian +``` + +and puts that value into **message metadata under the key `__partition_key`**. + +This is a real, supported contract, not a client-side convention — the server reads that exact +metadata key. All three SDKs use the same name and the same hash. A custom hasher is allowed but +must agree with whatever produced the topic's bounds. + +> Historical trap: an older C++ implementation used a **16-byte** hasher (two hashes, the second +> with seed `0x9E3779B97F4A7C15`, high half first). The current one is 8 bytes. The server's bound +> encoding supports both widths — it just writes an integer big-endian — so a mismatched hasher +> does not fail loudly, it silently routes keys to the wrong partitions. + +### 5.3 What the server does with `__partition_key` + +**It does not validate routing.** The value is consumed only by the auto-partitioning logic, as +input to deciding *where to cut* when the partition splits. A write goes to whatever partition the +client addressed, whether or not the key belongs there. + +Consequence for implementers: a routing bug produces no error. It shows up much later, as one key +smeared across two branches of the partition tree. Client-side range checking is the only defence. + +### 5.4 Lookup + +The ranges **tile** the key space: they do not overlap, and together they cover all of it. So a +hashed key falls into exactly one of them, and that partition owns it. That is the whole model. + +In code the lookup is a lower-bound search — index the partitions by `from_bound` and take the +greatest one at or below the hashed key: + +``` +partition = greatest from_bound <= hashed_key +``` + +On a complete set that is already the answer: the tiling guarantees the key is below that +partition's `to_bound`. + +Which is exactly why comparing against `to_bound` anyway is worth doing. It is not part of the +lookup, it is a check on *our own view of the topology*: + +``` +if partition.to_bound is not empty and hashed_key >= partition.to_bound: + refuse — our partition set is not a tiling; it has a hole +``` + +A lower-bound search always returns *something*. If our set is incomplete — a split whose second +child is not visible yet — the something it returns is the left neighbour, a sibling. Sending the +key there breaks invariant 1, and per §5.3 the server will not catch it. The check turns a silent +misroute into a refusal the caller can retry. + +The same applies when a partition set is accepted: validate it **after sorting**, not in arrival +order. Describe responses carry no ordering promise, and only the leftmost partition may have an +open lower bound. + +### 5.5 Two choosers + +| Chooser | Rule | Use | +|---|---|---| +| range-based | as above | auto-partitioned topics (the only correct one there) | +| hash-modulo | `murmur2_32(key) & 0x7FFFFFFF, then % partition_count` | fixed partition count; Kafka-compatible | + +The masking before the modulo matters for Kafka compatibility — Kafka applies `toPositive()` to a +*signed* 32-bit hash before the modulo. + +> Trap if you port from Go: its hash chooser carries a comment saying "Same as Kafka Partitioner" +> with a link to `BuiltInPartitioner`, but computes `hash % uint32(n)` with no mask. Unsigned +> modulo and masked signed modulo disagree for any hash with the high bit set, so keys land on +> different partitions than a Kafka producer would choose. For key `"a"` over 3 partitions Kafka +> picks index 2; the unmasked form picks 1. + +Selection is **adaptive**: pick range-based if the topic has auto-partitioning enabled *or* any +partition reports a key range, else hash-modulo. + +> Trap: a freshly created single-partition auto-partitioned topic can report **no key range at +> all**. Keying the decision on "does any partition have a range" alone therefore picks the +> hash-modulo chooser, which then rejects the bounded children the first split produces — the +> writer breaks precisely when the feature it exists for finally triggers. Prefer the topic's +> auto-partitioning strategy as the signal. + +--- + +## 6. Sub-writers + +### 6.1 Identity + +``` +producer_id = + separator + +``` + +The producer id is **per partition**, not per key. All three implementations agree. + +This has a large consequence: server-side deduplication is scoped to +`(partition, producer_id, seqno)`, so it **cannot span a split** — the child writes under a +different producer id. Everything in §9.4 exists because of this. + +Per-key producer ids would give cross-split dedup for free, but a write session carries exactly +one producer id, so that would mean one session per key. Nobody does this. + +### 6.2 Creation + +Sub-writers are created **lazily, on the first write to a partition**, never up front. An +auto-partitioned topic can have far more partitions than the caller's keys touch. + +Creation is: build settings → open the session → **wait for the init handshake** → register. + +Two details that are easy to get wrong: + +- **Bound the init wait.** A session pinned to an inactive partition never completes its handshake. + Since creation happens under the orchestrator lock, an unbounded wait there freezes every write, + flush and repartition in the whole writer. This is reachable in normal operation: splits cascade, + so the child you just routed to may itself have split already. +- **Register only after init succeeds.** Otherwise a failed creation leaves a permanently broken + writer in the pool. + +Sub-writers are configured with **automatic sequence numbering disabled** — the orchestrator +supplies the numbers (§7). + +### 6.3 Sessions always name their partition + +A session that writes always sets `partition_id`. There is no second mode to choose from. + +The one place the partition id is deliberately left out is reading the persisted seqno of a +partition that has already split (§9.4). That session never writes anything: it is opened only for +its init response and closed again. Leaving the id out is what makes it work at all — pinned to an +inactive partition, the session would never finish its init, and the number being read belongs to +the producer rather than to any one partition anyway. + +### 6.4 Idle eviction + +A sub-writer with **no in-flight messages** and no writes for the idle timeout is closed and +recreated on demand. Both conditions matter: closing a session with unacked writes would strand +them. + +Eviction must not be observable. In particular the **seqno cursor survives it** — a re-opened +partition continues the sequence rather than restarting. + +Recommended idle timeout is minutes, not seconds; too short and a bursty workload pays for constant +reconnects. + +### 6.5 Lifecycle + +``` + first write to partition + │ + ▼ + ┌─────────────┐ idle, no in-flight ┌──────────┐ + │ ACTIVE │ ──────────────────────► │ EVICTED │ + │ │ ◄────────────────────── │ │ + └──────┬──────┘ write arrives └──────────┘ + │ + │ OVERLOADED on the write stream + ▼ + ┌─────────────┐ + │ RETIRING │ session closed, ack failures suppressed, + │ │ successes still honoured + └──────┬──────┘ + │ + ┌────────┴────────┐ + ▼ ▼ + topology changed no children found + → RETIRED → recovered in place, back to ACTIVE +``` + +--- + +## 7. Sequence numbers + +### 7.1 One cursor for the whole writer + +There is **one counter for the entire multi-writer** — not one per partition, and not one per key. + +``` +write(key=A) → seqno 1 → partition 0 +write(key=B) → seqno 2 → partition 1 +write(key=A) → seqno 3 → partition 0 +write(key=C) → seqno 4 → partition 0 +``` + +Per-partition numbering looks simpler and is wrong in a specific way: when a split moves a message +to a child, a number drawn from the parent's sequence means nothing in the child's, so the message +must be **renumbered** — and a message that changes identity mid-flight can no longer be matched +against the attempt that may already have been persisted. A global cursor keeps the number valid in +every partition, which is what makes resend-preserving-seqno possible at all. + +Both reference implementations use a single global cursor for exactly this reason. + +Sequences per partition therefore have gaps. That is fine and expected. + +### 7.2 Seeding + +Each partition's producer has its own persisted history on the server. When a sub-writer opens, its +handshake reports that producer's last persisted seqno, and the cursor is lifted: + +``` +cursor = max(cursor, last_seqno_reported_by_this_session) +``` + +This is what makes a **stable producer id prefix** resume numbering across restarts instead of +colliding with what is already stored. Without it, a restarted writer re-uses low numbers and the +server rejects them as duplicates. + +### 7.3 Caller-supplied seqno + +When the caller numbers messages itself, the writer validates and tracks rather than generates: + +- a missing number is a **validation error**, not a writer-stopped error — the writer is healthy, + the input is not; +- the cursor is still lifted to the highest number seen, so a later switch or an internal use stays + consistent; +- uniqueness is enforced **per partition**, matching Go. Global uniqueness is *not* required. + +Merge is the one case where per-partition scope is not enough: two partitions may each hold seqno 5, +and a merge brings both into one child. This is detected and reported on one of the two messages +rather than silently overwriting state. Note both reference implementations ignore this case +entirely — they assume a single parent per partition. + +--- + +## 8. The normal write path + +``` +caller orchestrator sub-writer server + │ │ │ │ + │ write(key, data) │ │ │ + ├───────────────────────►│ │ │ + │ │ 1. choose partition by key │ │ + │ │ (stamps __partition_key)│ │ + │ │ │ │ + │ │ 2. get or create sub-writer│ │ + │ ├───────────────────────────►│ │ + │ │ │ init handshake │ + │ │ ├──────────────────►│ + │ │ │◄──────────────────┤ + │ │ cursor = max(cursor, last_seqno) │ + │ │ │ │ + │ │ 3. seqno = ++cursor │ │ + │ │ 4. hand message to writer │ │ + │ ├───────────────────────────►│ │ + │ │ ├──────────────────►│ + │ │ 5. record in-flight │ │ + │◄───────────────────────┤ return caller's future │ │ + │ │ │ │ + │ │ │◄─── ack ──────────┤ + │ │◄───────────────────────────┤ │ + │ │ 6. drop from in-flight │ │ + │ │ raise max-acked │ │ + │◄───────────────────────┤ resolve caller's future │ │ +``` + +Order matters in two places: + +- **Choose the partition before assigning the number.** The range-based chooser stamps + `__partition_key` into the message during routing; anything that freezes the message earlier + loses it. +- **Record in-flight only after the sub-writer accepts the message.** If admission fails — + backpressure timeout, stopped writer — an entry recorded beforehand leaks a future nobody will + ever resolve. + +--- + +## 9. Split + +### 9.1 The signal + +There is **no split notification**. A split is discovered as an ordinary error on the write stream: +the partition goes inactive, and the next write to it is rejected as `OVERLOADED`. + +``` +status = OVERLOADED +message = "Write to inactive partition N" +``` + +All three SDKs key on plain `OVERLOADED`. The server does carry a more precise internal code +(`WRITE_ERROR_PARTITION_INACTIVE`, surfaced as an issue code inside the overload status), which can +distinguish a split from ordinary load — but relying on it alone risks missing splits if it is not +always present, so it is best used as a fast positive signal with describe-based confirmation as the +fallback. + +Because ordinary overload is indistinguishable by type, **the topology change must be confirmed by +a describe** before anything is committed. + +### 9.2 Discovery + +Children are the **active leaf partitions that list the retiring partition as a parent**. Re-describe +with backoff until they appear; if none ever do, this was ordinary overload (§10). + +Discovery must also verify **coverage**: the children's ranges must span the parent's range without +a gap. + +``` +parent p0: ["", "") + +partial describe complete describe + p1: ["", 0x80) p1: ["", 0x80) + (p2 not yet active) p2: [0x80, "") + + ▲ retiring p0 here ▲ safe to retire p0 + leaves [0x80, "") + unowned → keys in + that range route to + p1, a sibling +``` + +A mid-split describe genuinely can show one child before its sibling becomes active. Committing to +that view breaks invariants 1 and 2. Treat an incomplete graph as *retry later*, never as a +successful split. + +Coverage means **"at least the parent's range"**, not an exact tiling: a merge child owns the ranges +of both its parents, so it legitimately covers more. + +### 9.3 Ordering of operations + +``` +1. discover children, verify coverage ← may abort; nothing committed yet +2. update the routing view: + add children, then remove every retired parent ← atomic, under the lock +3. quiesce every retired parent: + close its session, let pending acks settle +4. for each retired parent: + read the dedup cut ← only now + migrate its in-flight messages +``` + +Step 2 before step 4 means migration re-routes only to surviving partitions. Step 3 before step 4 is +what stops a still-open sibling from persisting a message *after* its cut was read — which would +duplicate that message on resend. + +All of this happens **under the orchestrator lock**, so no write can interleave and see a +half-updated topology. + +### 9.4 The dedup cut — the subtle part + +The question migration has to answer for every in-flight message is: *was this already persisted to +the partition we are retiring?* Below the cut → report it written, do not resend. Above → resend. + +**The cut must come from the server, not from acks the client observed.** + +``` +1. server persists seqno 42 +2. the session dies in the split, the ack never arrives +3. client's highest observed ack is 41 +4. client concludes 42 was not written and resends it to the child +5. the child writes under a different producer id, so nothing on the + server can recognise the duplicate +6. the reader sees the message twice +``` + +The observed-ack cut is wrong in exactly the window the mechanism exists to protect. + +Getting the server's answer for a partition that is already inactive is the trick worth copying: + +``` +open a session with the retiring partition's producer id + WITHOUT a partition id + → init handshake reports that producer's persisted seqno + → close it +``` + +Why that works is in the protocol itself. `InitResponse.last_seq_no` is documented as *"last +persisted message's sequence number for this producer"*, and `get_last_seq_no` warns that it *"may +be expensive, if producer wrote to many partitions before"* — the number belongs to the producer, +not to the partition the session happened to land on. So the session does not need to reach the +retired partition; it only needs to name its producer. + +Which is fortunate, because it cannot reach it: a session pinned to an inactive partition never +finishes its handshake — a real hang, not a slow path. Dropping the partition id is what lets the +session start at all. Both reference implementations do exactly this, and only for the split case. + +The cut is **the retiring partition's own high-water mark, and nothing else**. It is tempting to +walk up the partition tree and take the maximum over the ancestors as well -- both reference +implementations do -- but that is unnecessary here, and with merges it is actively wrong. + +Unnecessary, because a message sitting in this partition cannot have been persisted under any +producer it used earlier: + +``` +every move is gated by a cut of at least that producer's server seqno + -> anything already stored there was resolved as written on the spot + -> a message that travelled has a number strictly above it + -> and a retired producer never grows + +so asking an ancestor can only repeat an answer already known to be too low to matter +``` + +Wrong, because a merge child has two parents, and those branches numbered independently: + +``` +branch p1 numbered up to 5 branch p2 numbered up to 100 + \ / + → merged into p3 ← + +max over the parents = 100 +a message that came down p1 with seqno 6 is "already written" -- and is dropped +``` + +The sibling's history says nothing about messages that came down this branch, and can be +arbitrarily higher than theirs. Neither reference hits this: both read only the first parent and +state that a partition is assumed to have exactly one, so for them the walk is merely redundant. +Combining their walk with real merge support is what makes it unsafe. + +The same reasoning is why the recovery path (§10) has always used the partition's own value only. +With the walk gone the two paths compute the cut the same way. + +### 9.5 Migration + +``` +for each in-flight message of the retiring partition, in seqno order: + + seqno <= cut ? + └─ yes → resolve as written (offset unknown), do not resend + └─ no → re-route through the chooser + open the child's sub-writer + re-check the message is still in flight ← it may have been + acked during that await + keep the seqno, change only the partition + hand it to the child, re-wire the ack +``` + +Three things that look optional and are not: + +- **Iterate in seqno order.** Migration is the one place order can be lost. +- **Re-check after every await.** Opening the child's session yields; an ack can land in that window + and resolve the message. Resending it then duplicates it. +- **On a placement failure, fail this message and every message after it.** Dropping loses data + silently; skipping ahead reorders the key. Failing the tail is the only option that preserves both + invariants. + +### 9.6 Merge + +Structurally the same path: a merge child lists **two** parents. Discovering children of one parent +finds it, and every parent of that child that we still hold is retired **together** — otherwise the +sibling parent lingers in the routing view with a range that now overlaps the child's, and routing +becomes ambiguous. + +Both reference implementations ignore merge (they read only the first parent). At the time of +writing the server does not implement it either, so this path is defensive. + +**Only adjacent partitions can merge, and this is a protocol constraint, not a client one.** A +partition's range is a single `[from_bound, to_bound)` (§5.1). Merging two partitions that are not +neighbours in the key space would give the child two disjoint stretches, and there is no field in +which to report that — so the situation cannot arise no matter how the server chooses to implement +merging. Adjacency is what makes the result expressible: + +``` +p1 [a, m) + p2 [m, z) → p3 [a, z) one interval, reportable +p1 [a, m) + p3 [t, z) → [a,m) ∪ [t,z) no representation for this +``` + +This also means the client needs no special rule to reject a gapped child: it cannot be described +in the first place. What the client does have to handle is the *transition* — the moment when both +parents are still in its view alongside the child, which is why every parent of a discovered child +is retired in the same step. + +### 9.7 Full sequence + +``` +sub-writer(p0) orchestrator describe probe child(p2) + │ │ │ │ │ + │ OVERLOADED │ │ │ │ + ├───────────────────►│ │ │ │ + │ │ stop p0's writer with │ │ │ + │◄───────────────────┤ a distinguishable error│ │ │ + │ │ │ │ │ + │ │ schedule repartition │ │ │ + │ │ (coalesced per partition) │ │ + │ │ │ │ │ + │ ├───── describe ────────►│ │ │ + │ │◄──── children [2,3] ───┤ │ │ + │ │ verify coverage │ │ │ + │ │ │ │ │ + │ │ routing: +2 +3, -0 │ │ │ + │ │ │ │ │ + │ │ quiesce p0 (close, settle acks) │ │ + │ │ │ │ │ + │ ├─── unpinned session, producer "pfx-0" ►│ │ + │ │◄── last_seqno = 2593 ──────────────────┤ │ + │ │ cut = 2593 │ │ + │ │ │ │ │ + │ │ seqno <= 2593 → resolve as written │ │ + │ │ seqno > 2593 → resend, same seqno ───────────────────►│ + │ │ │ │ │ +``` + +--- + +## 10. Ordinary overload + +If describe never shows children, the overload was not a topology change. The partition is kept and +**recovered in place**: drop the stopped session, open a fresh one for the same partition, resend +the in-flight messages with their original seqnos. + +Here the cut is exact and cheap — the partition is still active, so the fresh session's handshake +reports the current persisted seqno directly. No probe needed. + +It is worth spelling out why a cut is needed here at all, because the obvious argument says it is +not: the partition and the producer id have not changed, so if the server already stored one of +these messages it will recognise the repeat and deduplicate it. Resending everything should be +harmless. + +It is not, and the reason has nothing to do with the server: + +``` +ack for seqno 42 is lost when the session breaks +recovery opens a fresh sub-writer for the same partition + its handshake reports last_seq_no = 42 ← the server does have it + the writer takes 42 as its high-water mark +resend 42 → rejected by our own writer, "seqno is duplicated" + the message never leaves the process +resend 43, 44, ... → never attempted: the loop stopped at the error +``` + +Sub-writers run with automatic numbering disabled (§6.2), and such a writer refuses any seqno at or +below the high-water mark it learned at init. So the message is stopped on the client, before the +server ever gets the chance to deduplicate it. Because the resend is a single ordered pass, that +one rejection also takes down every message queued behind it: none of them are sent, and none of +them are resolved. + +Applying the cut removes the problem at the source — seqno 42 is at or below it, so it is reported +as written instead of resent, and the pass continues from 43. + +--- + +## 11. Concurrency + +A **single lock** serialises: routing decisions, sub-writer creation, seqno assignment, in-flight +bookkeeping, topology updates, migration, idle eviction. + +That is coarse, and deliberately so: routing and topology changes must not interleave. The cost is +that anything slow taken under it stalls the whole writer, which is why every network wait under the +lock is bounded (§6.2, §9.4). + +Repartition runs as a background task, and it must be **owned**: + +- **coalesced per partition** — a burst of `OVERLOADED` must not start several concurrent recoveries + of the same partition; +- **cancelled and awaited on close** — otherwise it keeps describing topics and opening sessions for + a writer the caller believes is shut down; +- **blocked after close** — a late signal must not resurrect anything. + +--- + +## 12. How a message can end + +Every accepted message reaches exactly one of these: + +| Outcome | When | +|---|---| +| written, with offset | normal ack | +| written, offset unknown | at or below a dedup cut — the server had it, the ack was lost | +| failed, placement | no ready partition owns the key, or the target could not be opened | +| failed, seqno conflict | a caller-supplied number collided in the target partition | +| failed, partition unusable | repartition and recovery both failed | +| failed, writer stopped | closed with messages still in flight | + +The "offset unknown" outcome is a real part of the contract: the write happened, but the offset +came back on a stream that died. Both reference implementations do the same — C++ synthesises an +acknowledgement for exactly this case. + +--- + +## 13. Traps, ranked by how much they cost to find + +1. **A pinned session to an inactive partition hangs forever.** Costs: a frozen writer, if it is + under the lock. Found only against a real split. +2. **The dedup cut from observed acks is wrong.** Costs: duplicates, rarely, unreproducibly. +3. **Closing a hook-stopped session re-raises the stop reason.** Closing a session you are + discarding must swallow errors, or the repartition aborts on cleanup. +4. **Suppressing ack failures during teardown must not suppress successes.** A success landing while + the session closes is real; dropping it lowers the cut and duplicates the message. +5. **A mid-split describe can show one child of two.** Committing to it smears a key across + siblings. +6. **Greatest-lower-bound routing never fails, it just answers wrongly** when the partition set has + a hole. And the server will not catch it. +7. **A single-partition auto-partitioned topic may report no key range**, which fools chooser + selection that keys on ranges alone. +8. **You may not be able to make this writer split a topic at all.** On the cluster we tested, + load driven through the multi-writer never triggered auto-partitioning, while separate + load from ordinary writers, with no explicit partition id, split the topic in seconds. Whether that is general or a property of that + build, we did not establish -- but plan test scenarios for it: drive splits with a separate + producer rather than expecting the writer under test to cause its own. +9. **Test doubles hide all of the above.** Fakes that always accept a seqno, always initialise, and + never raise on close will keep every one of these bugs green. Model the real contracts: the + seqno high-water rejection, the init that never completes, the close that re-raises, and state + keyed by producer id rather than by partition. + +--- + +## 14. Implementation cross-reference + +| Aspect | Here | C++ `TProducer` | Go `topicmultiwriter` | +|---|---|---|---| +| producer id | per partition | per partition | per partition | +| seqno cursor | one, writer-wide | one, writer-wide | one, writer-wide | +| resend keeps seqno | yes | yes | yes | +| cut source | server, retiring partition only | server, whole lineage | server, whole lineage | +| cut probe | unpinned session | unpinned session | unpinned session | +| split signal | `OVERLOADED` | `OVERLOADED` | `OVERLOADED` | +| key hash | murmur64a, 8B big-endian | same | same | +| Kafka positive mask | yes | n/a | no | +| merge | handled | single parent assumed | single parent assumed | +| transient overload | recovered in place | — | stops with error | +| coverage check before retiring | yes | no | no | +| upper-bound check on routing | yes | no | no | + +The last four rows are places where this implementation is deliberately stricter than the +references. + +--- + +## 15. Open items + +- Batch writes are admitted message by message; a mid-batch failure can leave earlier messages + accepted while the call raises, with no handle returned for them. +- Buffer limits are per sub-writer, so the effective budget scales with the number of open + partitions rather than belonging to the logical writer. +- Split coverage in automated tests depends on the environment producing a real split; an + environment that cannot must be distinguished from a split that failed to happen. diff --git a/MULTIWRITER_ARCHITECTURE.ru.md b/MULTIWRITER_ARCHITECTURE.ru.md new file mode 100644 index 000000000..2df23f101 --- /dev/null +++ b/MULTIWRITER_ARCHITECTURE.ru.md @@ -0,0 +1,828 @@ +# Мультипартиционный писатель топиков — архитектура + +Как устроена запись по ключу, почему каждая часть сделана именно так и что ломается, если сделать +иначе. Написано по реализации в Python SDK, но без привязки к языку: описываются взаимодействия по +протоколу и логика состояний, а не API. + +Сверено с двумя существующими реализациями — C++ `TProducer` +(`src/client/topic/impl/producer.{h,cpp}`) и Go `topicmultiwriter` +(`internal/topic/topicmultiwriter`). Расхождения отмечены отдельно. + +Термины протокола (`seqno`, `producer_id`, `partition_id`, `from_bound`, +`to_bound`, `describe`, `OVERLOADED`, `ack`, `flush`, `future`) оставлены как есть — так их называют +в коде и в обсуждениях. + +--- + +## 1. Что это такое + +**Мультиписатель** — один логический писатель, который раскладывает сообщения по всем партициям +топика, сохраняя порядок в пределах ключа. Вызывающий прикладывает к каждому сообщению **ключ**, +писатель решает, какая партиция этим ключом владеет, и пишет туда. + +Нужен он из-за серверного автопартиционирования: сервер делит партицию, когда она становится +горячей. Обычный писатель привязан к одной партиции и просто умирает, когда та исчезает. +Мультиписатель вместо этого следует за изменениями топологии. + +``` + write(key="user-42", data=...) + │ + ▼ + ┌───────────────────────┐ + │ оркестратор │ выбор партиции, выдача seqno, + │ │ учёт неподтверждённых, обработка сплита + └───────────┬───────────┘ + выбор по ключу │ один под-писатель на партицию + ┌──────────────┬──────┴───────┬──────────────┐ + ▼ ▼ ▼ ▼ + под-писатель под-писатель под-писатель под-писатель + партиция 0 партиция 1 партиция 2 партиция N + producer_id producer_id producer_id producer_id + "pfx-0" "pfx-1" "pfx-2" "pfx-N" + │ │ │ │ + ▼ ▼ ▼ ▼ + отдельный поток записи на каждую партицию + (свои буферизация, сжатие, переподключение, обновление токена) +``` + +Под-писатель — это **обычный однопартиционный писатель**, без изменений. Оркестратор добавляет +поверх только выбор партиции, выдачу seqno и обработку сплита. Это самое важное структурное +решение: буферизация, сжатие, переподключение и обновление авторизации не переписываются заново. + +> Почему это вообще живёт на клиенте. Маршрутизацию по ключу изначально задумывали серверной: +> клиент прикладывает к каждому сообщению `message_group_id`, а сервер сам раскладывает группы по +> партициям. В протоколе этот замысел так и остался — `message_group_id` присутствует и как +> настройка сессии, и, в `MessageData.partitioning`, как поле **каждого сообщения**, что имеет +> смысл только если раскладывать собирался сервер. До реализации не дошло, поэтому маршрутизацию +> делает каждый SDK у себя, и мультиписатель — это она и есть. +> +> Практическое следствие: **считайте `message_group_id` вторым именем для `producer_id`.** Все +> реализации кладут туда одно и то же значение, ничего по нему не маршрутизируется, и к ключу +> сообщения оно отношения не имеет. Это не вторая сущность, о которой надо думать. + +--- + +## 2. Контракт + +На что может рассчитывать вызывающий: + +| # | Гарантия | +|---|---| +| 1 | Все сообщения с данным ключом идут в одну партицию либо в её потомков после сплита | +| 2 | Порядок в пределах ключа сохраняется | +| 3 | Каждое принятое сообщение получает окончательный исход: либо ack, либо ошибка. Не теряется молча и не остаётся висеть навсегда | +| 4 | Сплит незаметен: ни потерь, ни дублей для сообщений, которые сервер уже записал | +| 5 | Сплит не требует пересоздавать писателя | +| 6 | Временные ошибки ретраятся внутри | +| 7 | Закрытие с flush доставляет всё, что уже было принято | + +Чего он сознательно **не** гарантирует: + +- Общего порядка между разными ключами нет. Только внутри ключа. + +--- + +## 3. Внутренние инварианты + +Это свойства, которые защищает каждая часть реализации. Если вы портируете — тесты надо писать +именно против них. + +1. **Один ключ → одна ветка дерева партиций.** Ключ может пройти путь `p0 → p2 → p5`, но не должен + оказаться сразу в двух детях одного родителя. Иначе читатель потеряет порядок по этому ключу. +2. **В наборе партиций не бывает дыр.** Каждый ключ в каждый момент отображается ровно в одну живую + партицию, в том числе пока сплит ещё идёт. +3. **seqno опознаёт сообщение всю его жизнь.** Выдаётся один раз и не меняется, даже когда + сообщение переезжает в другую партицию. +4. **Сообщение находится ровно в одном месте:** либо ждёт ack ровно у одной партиции, либо уже + завершено. +5. **Срез дедупликации берётся только с сервера.** См. §9.4 — самая тонкая часть. + +--- + +## 4. Состояние + +Всё перечисленное принадлежит оркестратору и защищено одним общим локом (§11). + +| Состояние | Форма | Назначение | +|---|---|---| +| набор партиций | `partition_id → сведения о партиции` | только живые листья; среди них и идёт выбор | +| под-писатели | `partition_id → писатель` | создаются лениво, закрываются по простою | +| неподтверждённые | `partition_id → {seqno → сообщение}` | сообщения, которые, возможно, придётся переотправить | +| курсор seqno | одно целое | одна последовательность на весь писатель (§7) | +| seqno на момент открытия | `partition_id → int` | что сервер сообщил при инициализации текущей сессии | +| seqno снятой партиции | `partition_id → int` | финальное значение для снятого с маршрутизации producer_id; кешируется | +| максимум подтверждённого | `partition_id → int` | наибольший seqno, по которому реально пришёл ack | +| снимаемые партиции | множество `partition_id` | подавляет ожидаемые ошибки ack во время закрытия сессии | +| задачи обработки сплита | `partition_id → задача` | схлопывание повторных сигналов и корректное завершение | +| время последней записи | `partition_id → отметка времени` | закрытие по простою | + +Набор партиций содержит только живые листья: партицию убирают оттуда сразу, как её сняли, иначе +ключи продолжали бы в неё уезжать. Ничего о снятых партициях здесь не хранится, и не нужно — +почему срез дедупликации никогда не заглядывает дальше снимаемой партиции, см. §9.4. + +--- + +## 5. Выбор партиции + +### 5.1 Диапазоны ключей + +**Диапазонами владеет сервер, клиент их только читает.** Они приходят в ответе `DescribeTopic`, по +одному на партицию: + +``` +DescribeTopicResult.PartitionInfo.key_range : PartitionKeyRange + optional bytes from_bound // включаемая левая граница, пусто = -inf + optional bytes to_bound // исключаемая правая граница, пусто = +inf +``` + +Ничто на стороне клиента диапазоны не вычисляет, не делит и не достраивает. Клиент читает то, что +сообщил `describe`, хеширует ключ и выбирает партицию, в чей диапазон хеш попал. Клиент, который +считал бы диапазоны сам, тут же разошёлся бы с сервером в том, какой ключ кому принадлежит. + +Обрати внимание на форму: **один интервал на партицию, а не набор интервалов.** Выразить партицию, +владеющую двумя несвязными кусками пространства ключей, попросту нечем. Из этого одного факта +следует ответ на большинство вопросов о том, какие изменения топологии вообще возможны, — что это +значит для merge, см. §9.6. + +Границы — это байтовые строки. Сравниваются **побайтово**, как обычные последовательности байт. +Никаких локалей, кодировок и числовой интерпретации. На живом кластере границы оказались короткими: +реальный сплит дал двухбайтовые значения `efea` и `f7f5`. + +Важно: с границами сравнивается **не сам ключ, а его хеш** (§5.2). Поэтому в примере ниже слева +стоят байты хеша, а не текст ключа. + +``` +до сплита + p0 from_bound = нет to_bound = нет активна, владеет всем пространством + +после сплита + p0 from_bound = нет to_bound = нет НЕактивна, дети [1, 2] + p1 from_bound = нет to_bound = 0x80 активна, владеет хешами < 0x80 + p2 from_bound = 0x80 to_bound = нет активна, владеет хешами >= 0x80 + +хеш ключа 0x1a... < 0x80 → p1 +хеш ключа 0xff... >= 0x80 → p2 +хеш ключа 0x80... ровно граница → p2 левая граница включена, правая нет +``` + +### 5.2 Ключ хешируется перед сравнением + +> В документе встречаются два разных хеша — по одному на каждый способ выбора партиции. Это не +> альтернативы и не противоречие: +> +> | Способ | Хеш | Для чего | +> |---|---|---| +> | по диапазонам (§5.4) | **MurmurHash64A**, 64 бита | попадание ключа в `[from_bound, to_bound)` | +> | хеш по модулю (§5.5) | **murmur2, 32 бита** | `% число_партиций` при фиксированном их числе | +> +> Первый — это путь YDB, и именно он обязан совпадать с другими писателями топика. Второй нужен +> только чтобы повторить партиционирование Apache Kafka, и до диапазонов ключей он не доходит. + +Для выбора по диапазонам клиент не сравнивает сырой ключ с границами. Он его хеширует: + +``` +partition_key = MurmurHash64A(байты_ключа, seed = 0) → 8 байт, big-endian +``` + +и кладёт результат в метаданные сообщения под ключом `__partition_key`. + +Это не самодеятельность клиента, а поддерживаемый контракт: сервер читает именно это имя поля. Все +три SDK используют одно имя и один хеш. Свой хешер задать можно, но он обязан совпадать с тем, чем +построены границы топика. + +> Историческая ловушка: старая реализация в C++ считала **16 байт** (два хеша, второй с сидом +> `0x9E3779B97F4A7C15`, старшая половина впереди). Актуальная считает 8. Сервер принимает обе +> ширины — он просто пишет целое в big-endian, — поэтому несовпадающий хешер не даёт ошибки, а +> молча отправляет ключи не в те партиции. + +### 5.3 Что сервер делает с `__partition_key` + +**Он не проверяет, туда ли вы пишете.** Значение читает только логика автопартиционирования, и +нужно оно ей для одного — решить, *в каком месте резать* партицию при сплите. Сама запись уходит в +ту партицию, которую указал клиент, независимо от того, принадлежит ей этот ключ или нет. + +Отсюда следствие: ошибка выбора партиции не приводит ни к какой ошибке. Она всплывёт сильно позже, +когда окажется, что один ключ размазан по двум веткам дерева. Проверка на стороне клиента — +единственная защита. + +### 5.4 Как ищется партиция + +Диапазоны **разбивают** пространство ключей: они не пересекаются и вместе покрывают его целиком. +Значит хеш ключа попадает ровно в один из них, и владеет ключом соответствующая партиция. Вот и вся +модель. + +В коде поиск делается по нижней границе: партиции индексируются по `from_bound`, берётся +наибольшая, не превышающая хеш ключа. + +``` +партиция = наибольший from_bound <= хеш_ключа +``` + +На полном наборе это уже ответ: раз диапазоны разбивают пространство без дыр, ключ гарантированно +меньше `to_bound` найденной партиции. + +Именно поэтому сравнение с `to_bound` всё равно стоит делать. Это не часть поиска, а проверка +**нашего собственного представления о топологии**: + +``` +если у партиции есть to_bound и хеш_ключа >= to_bound: + отказ — значит наш набор партиций не разбиение, в нём дыра +``` + +Поиск по нижней границе всегда возвращает *какую-нибудь* партицию. Если наш набор неполон — +например, второй ребёнок сплита ещё не виден, — вернётся соседняя слева, то есть другой ребёнок +того же родителя. Отправить туда ключ значит нарушить инвариант 1, и по §5.3 сервер этого не +заметит. Проверка превращает молчаливую ошибку маршрутизации в отказ, который вызывающий может +повторить. + +То же касается приёма набора партиций: проверять его надо **после сортировки**, а не в порядке +поступления. Ответ `describe` не обещает никакого порядка, а отсутствовать `from_bound` может +только у самой левой партиции. + +### 5.5 Два способа выбора + +| Способ | Правило | Когда применим | +|---|---|---| +| по диапазонам | как выше | автопартиционированные топики — там верен только он | +| хеш по модулю | `murmur2_32(key) & 0x7FFFFFFF`, затем `% число_партиций` | фиксированное число партиций; совместимость с Kafka | + +Маска перед взятием модуля нужна именно для совместимости с Kafka: там `toPositive()` применяется к +**знаковому** 32-битному хешу. + +> Ловушка при портировании с Go: у их хеш-варианта стоит комментарий «Same as Kafka Partitioner» со +> ссылкой на `BuiltInPartitioner`, но считается `hash % uint32(n)`, без маски. Беззнаковый модуль и +> маскированный знаковый дают разный результат для любого хеша со старшим установленным битом, так +> что ключи уедут не туда, куда отправил бы продюсер Kafka. Для ключа `"a"` на трёх партициях Kafka +> даёт 2, вариант без маски — 1. + +Способ выбирается автоматически: по диапазонам, если у топика включено автопартиционирование **или** +хотя бы одна партиция сообщила диапазон; иначе хеш по модулю. + +> Ловушка: только что созданный автопартиционированный топик с одной партицией может не сообщить +> диапазон **вообще**. Если смотреть только на «есть ли у кого-то диапазон», выберется хеш по +> модулю, а он потом откажется принимать детей первого же сплита — писатель сломается ровно тогда, +> когда наконец сработает то, ради чего он нужен. Смотреть надо на настройку автопартиционирования +> топика. + +--- + +## 6. Под-писатели + +### 6.1 Идентификатор + +``` +producer_id = <префикс> + разделитель + +``` + +`producer_id` привязан **к партиции**, а не к ключу. Все три реализации сходятся. + +Отсюда важное следствие: серверная дедупликация работает в пределах тройки +`(партиция, producer_id, seqno)` и потому **не может пересечь границу сплита** — ребёнок пишет под +другим `producer_id`. Весь §9.4 существует из-за этого. + +Отдельный `producer_id` на каждый ключ дал бы дедупликацию через сплит даром, но сессия записи +несёт ровно один `producer_id`, то есть это означало бы отдельную сессию на каждый ключ. Так не +делает никто. + +### 6.2 Создание + +Под-писатели создаются **лениво, при первой записи в партицию**, и никогда заранее. У +автопартиционированного топика партиций может быть намного больше, чем задевают ключи клиента. + +Порядок: собрать настройки → открыть сессию → **дождаться ответа на инициализацию** → +зарегистрировать в пуле. + +Две детали, в которых легко ошибиться: + +- **Ограничьте ожидание инициализации по времени.** Сессия с явным `partition_id` неактивной + партиции ответа не дождётся никогда. А так как создание идёт под общим локом, бесконечное + ожидание там замораживает все записи, `flush` и обработку сплита во всём писателе. Это не + экзотика: сплиты идут каскадом, и ребёнок, который вы только что выбрали, может уже сам + расщепиться. +- **Регистрируйте писателя только после успешной инициализации.** Иначе неудачная попытка оставит в + пуле навсегда нерабочую сессию. + +Под-писателям выключают автоматическую выдачу seqno — номера выдаёт оркестратор (§7). + +### 6.3 Сессия всегда указывает свою партицию + +Сессия, которая пишет, всегда задаёт `partition_id`. Никакого второго режима на выбор нет. + +Единственное место, где `partition_id` сознательно не указывается, — чтение записанного seqno у +партиции, которая уже расщепилась (§9.4). Такая сессия ничего не пишет: она открывается только +ради ответа на инициализацию и сразу закрывается. И работает она именно потому, что партиция не +указана: пришпиленная к неактивной партиции, она не дождалась бы инициализации никогда. + +### 6.4 Закрытие по простою + +Под-писателя, у которого **нет неподтверждённых сообщений** и не было записей в течение таймаута, +закрывают, а при следующей записи открывают заново. Важны оба условия: закрыть сессию с +неподтверждёнными записями значит бросить их без владельца. + +Закрытие по простою не должно быть заметно снаружи. В частности, **курсор seqno его переживает** — +заново открытая партиция продолжает нумерацию, а не начинает с нуля. + +Разумный таймаут — минуты, а не секунды. Слишком короткий заставит неравномерную нагрузку постоянно +переподключаться. + +### 6.5 Жизненный цикл + +``` + первая запись в партицию + │ + ▼ + ┌─────────────┐ простой, ничего не ждёт ack ┌─────────────┐ + │ РАБОТАЕТ │ ────────────────────────────► │ ЗАКРЫТ │ + │ │ ◄──────────────────────────── │ по простою │ + └──────┬──────┘ пришла запись └─────────────┘ + │ + │ OVERLOADED в потоке записи + ▼ + ┌─────────────┐ сессия закрыта; ошибки ack подавлены, + │ СНИМАЕТСЯ │ но успешные ack по-прежнему учитываются + └──────┬──────┘ + │ + ┌────────┴─────────┐ + ▼ ▼ + топология изменилась детей не нашлось + → СНЯТ → восстановлен на месте, + с маршрутизации снова РАБОТАЕТ +``` + +--- + +## 7. Выдача seqno + +### 7.1 Один курсор на весь писатель + +Счётчик **один на весь мультиписатель** — не на партицию и не на ключ. + +``` +write(key=A) → seqno 1 → партиция 0 +write(key=B) → seqno 2 → партиция 1 +write(key=A) → seqno 3 → партиция 0 +write(key=C) → seqno 4 → партиция 0 +``` + +Отдельный счётчик на каждую партицию выглядит проще, но неверен вполне конкретно. Когда сплит +переносит сообщение в ребёнка, номер из последовательности родителя в последовательности ребёнка +ничего не значит, поэтому сообщение приходится перенумеровывать. А сообщение, сменившее номер на +лету, уже не сопоставить с попыткой, которая, возможно, уже записана на сервере. Общий счётчик +оставляет номер осмысленным в любой партиции — только поэтому переотправка и может сохранять seqno. + +Обе эталонные реализации держат один общий счётчик ровно по этой причине. + +Побочный эффект: внутри отдельной партиции номера идут с пропусками. Это нормально. + +### 7.2 Откуда берётся стартовое значение + +У `producer_id` каждой партиции своя история на сервере. Когда открывается под-писатель, сервер в +ответе на инициализацию сообщает последний записанный seqno этого `producer_id`, и курсор +подтягивается вверх: + +``` +курсор = max(курсор, last_seqno из ответа сервера) +``` + +Именно это позволяет писателю со **стабильным префиксом `producer_id`** продолжить нумерацию после +перезапуска, а не столкнуться с тем, что уже записано. Без этого перезапущенный писатель начнёт с +низких номеров, и сервер отвергнет их как дубли. + +### 7.3 Когда seqno задаёт вызывающий + +Если номера проставляет сам клиент, писатель их не генерирует, а проверяет: + +- отсутствующий номер — это **ошибка входных данных**, а не «писатель остановлен». Писатель здоров, + проблема в сообщении, и тип ошибки не должен вводить в заблуждение обработку ретраев; +- курсор всё равно подтягивается до наибольшего увиденного номера, чтобы дальнейшая работа осталась + согласованной; +- уникальность проверяется **в пределах партиции**, как в Go. Глобальная уникальность **не** + требуется. + +Единственное место, где области видимости партиции не хватает, — это merge: две партиции могут +держать seqno 5 каждая, а merge сводит обе в одного ребёнка. Такой конфликт обнаруживается и +возвращается ошибкой на одном из двух сообщений, вместо того чтобы молча затереть состояние. Обе +эталонные реализации этот случай вообще не рассматривают: они считают, что у партиции ровно один +родитель. + +--- + +## 8. Штатный путь записи + +``` +вызывающий оркестратор под-писатель сервер + │ │ │ │ + │ write(key, data) │ │ │ + ├───────────────────────►│ │ │ + │ │ 1. выбрать партицию по ключу │ + │ │ (заодно проставляется __partition_key) │ + │ │ │ │ + │ │ 2. взять из пула или создать под-писателя │ + │ ├───────────────────────────►│ │ + │ │ │ инициализация │ + │ │ ├──────────────────►│ + │ │ │◄──────────────────┤ + │ │ курсор = max(курсор, last_seqno) │ + │ │ │ │ + │ │ 3. seqno = ++курсор │ │ + │ │ 4. отдать сообщение под-писателю │ + │ ├───────────────────────────►│ │ + │ │ ├──────────────────►│ + │ │ 5. запомнить как неподтверждённое │ + │◄───────────────────────┤ вернуть future вызывающему │ + │ │ │ │ + │ │ │◄────── ack ───────┤ + │ │◄───────────────────────────┤ │ + │ │ 6. убрать из неподтверждённых │ + │ │ поднять максимум подтверждённого │ + │◄───────────────────────┤ завершить future │ │ +``` + +Порядок важен в двух местах: + +- **Партиция выбирается до выдачи seqno.** Выбор по диапазонам попутно проставляет + `__partition_key` в метаданные сообщения. Если сериализовать или «заморозить» сообщение раньше, + это поле потеряется. +- **Сообщение попадает в список неподтверждённых только после того, как под-писатель его принял.** + Если приём сорвётся — истёк таймаут ожидания места в буфере, писатель остановлен, — то запись, + сделанная заранее, оставит future, который никто никогда не завершит. + +--- + +## 9. Сплит + +### 9.1 Сигнал + +**Отдельного уведомления о сплите нет.** Сплит виден как обычная ошибка в потоке записи: партиция +становится неактивной, и следующая запись в неё отвергается с `OVERLOADED`. + +``` +status = OVERLOADED +message = "Write to inactive partition N" +``` + +Все три SDK опираются именно на голый `OVERLOADED`. У сервера есть и точный внутренний код +(`WRITE_ERROR_PARTITION_INACTIVE`, он приходит как issue-код внутри статуса перегрузки), по которому +сплит отличается от обычной нагрузки. Но опираться только на него рискованно: если он придёт не во +всех случаях, мы вообще перестанем замечать сплиты. Разумнее использовать его как быстрый +положительный признак, оставив подтверждение через `describe` запасным путём. + +Раз обычная перегрузка по типу ошибки неотличима, **изменение топологии обязательно подтверждается +через `describe`** до того, как что-то будет зафиксировано. + +### 9.2 Поиск детей + +Дети — это **активные листовые партиции, у которых снимаемая партиция указана родителем**. Повторяем +`describe` с паузами, пока они не появятся. Если так и не появились — это была обычная перегрузка +(§10). + +Кроме самого факта появления надо проверить **покрытие**: диапазоны детей должны перекрывать +диапазон родителя без дыр. + +``` +родитель p0 владеет всем пространством ключей + +неполный ответ describe полный ответ describe + p1: до 0x80 p1: до 0x80 + (p2 ещё не активна) p2: от 0x80 и дальше + + ▲ если снять p0 сейчас, ▲ здесь снимать p0 безопасно: + диапазон от 0x80 и дальше весь диапазон родителя + останется без владельца, покрыт детьми + и ключи оттуда уедут в p1 — + то есть в соседнюю ветку +``` + +`describe` посреди сплита действительно может показать одного ребёнка раньше, чем второй станет +активным. Зафиксировать такую картину — нарушить инварианты 1 и 2. Неполный граф надо считать +состоянием «повторить позже», а не успешным сплитом. + +Покрытие означает «**не меньше** диапазона родителя», а не точное совпадение: ребёнок merge владеет +диапазонами обоих родителей, поэтому законно покрывает больше. + +### 9.3 Порядок операций + +``` +1. найти детей, проверить покрытие ← может прерваться; ничего ещё не зафиксировано +2. обновить набор партиций: + добавить детей, затем убрать всех снимаемых родителей ← атомарно, под локом +3. для каждого снимаемого родителя: + закрыть его сессию и дать осесть уже пришедшим ack +4. для каждого снимаемого родителя: + прочитать срез дедупликации ← только теперь + переотправить его неподтверждённые сообщения +``` + +Шаг 2 раньше шага 4 — чтобы переотправка выбирала только среди выживших партиций. Шаг 3 раньше +шага 4 — чтобы ещё открытая сессия соседнего родителя не записала сообщение уже *после* того, как +его срез прочитан: такое сообщение продублировалось бы при переотправке. + +Всё это происходит **под общим локом**, поэтому ни одна запись не вклинится и не увидит +полуобновлённую топологию. + +### 9.4 Срез дедупликации — самая тонкая часть + +Про каждое неподтверждённое сообщение переотправка обязана ответить на вопрос: *было ли оно уже +записано в ту партицию, которую мы снимаем?* Не выше среза → считаем записанным и не отправляем +повторно. Выше среза → отправляем. + +**Срез обязан приходить с сервера, а не из тех ack, которые успел увидеть клиент.** + +``` +1. сервер записал сообщение с seqno 42 +2. сессия умерла из-за сплита, ack не дошёл +3. наибольший увиденный клиентом ack — 41 +4. клиент решает, что 42 не записано, и отправляет его ребёнку +5. ребёнок пишет под другим producer_id, поэтому на сервере + ничто не может распознать дубль +6. читатель видит сообщение дважды +``` + +То есть срез по увиденным ack неверен ровно в том окне, ради которого весь механизм и существует. + +Как спросить сервер про уже неактивную партицию — тот приём, который стоит перенять: + +``` +открыть сессию под producer_id снимаемой партиции, + НЕ указывая partition_id + → в ответе на инициализацию придёт last_seqno этого producer_id + → сессию сразу закрыть +``` + +Почему это работает — написано в самом протоколе. `InitResponse.last_seq_no` описан как «last +persisted message's sequence number for **this producer**», а у флага `get_last_seq_no` стоит +предупреждение «may be expensive, if producer wrote to **many partitions** before». То есть число +принадлежит продюсеру, а не той партиции, куда сессию отнесло. Дотягиваться до снятой партиции +сессии не нужно — достаточно назвать её продюсера. + +Сессия с явным `partition_id` неактивной партиции ответа не дождётся — это настоящее зависание, а +не медленный путь. Сессия без `partition_id` отвечает нормально. Обе эталонные реализации делают +именно так, и только для случая сплита. + +Срез — это **верхняя отметка самой снимаемой партиции, и ничего больше**. Соблазнительно +подняться по дереву и взять максимум ещё и по предкам — обе эталонные реализации так и делают, — +но здесь это лишнее, а при слиянии ещё и неверно. + +Лишнее, потому что сообщение, лежащее в этой партиции, не могло быть записано ни под одним из +прежних своих продюсеров: + +``` +каждый переезд проходит через срез не ниже серверного значения того продюсера + -> всё, что там уже лежало, тут же завершилось как записанное + -> значит у доехавшего сообщения номер строго выше + -> а снятый продюсер больше не растёт + +спрашивать предка — значит повторно получить ответ, заведомо слишком низкий, чтобы на что-то влиять +``` + +Неверно, потому что у ребёнка слияния два родителя, и ветки нумеровались независимо: + +``` +ветка p1 дошла до 5 ветка p2 дошла до 100 + \ / + → слились в p3 ← + +максимум по родителям = 100 +сообщение, пришедшее по ветке p1 с номером 6, «уже записано» — и теряется +``` + +История соседней ветки ничего не говорит о сообщениях, пришедших по этой, и может быть сколь +угодно выше. Ни один из эталонов на это не натыкается: оба читают только первого родителя и прямо +заявляют, что у партиции предполагается ровно один. Для них обход просто избыточен. Опасным его +делает сочетание их обхода с настоящей поддержкой merge. + +По той же причине путь восстановления (§10) всегда считал срез только по самой партиции. После +удаления обхода оба пути считают его одинаково. + +### 9.5 Переотправка + +``` +для каждого неподтверждённого сообщения снимаемой партиции, по возрастанию seqno: + + seqno не выше среза ? + └─ да → завершить future как «записано» (офсет неизвестен), не отправлять + └─ нет → заново выбрать партицию по ключу + открыть под-писателя выбранного ребёнка + перепроверить, что сообщение всё ещё неподтверждено ← за время + ожидания мог прийти ack + сохранить прежний seqno, сменить только партицию + отдать ребёнку и перевесить обработку ack +``` + +Три вещи, которые выглядят необязательными, но таковыми не являются: + +- **Идти по возрастанию seqno.** Переотправка — единственное место, где можно потерять порядок. +- **Перепроверять после каждого ожидания.** Открытие сессии ребёнка отдаёт управление; в это окно + может прийти ack и завершить сообщение. Отправив его после этого, вы получите дубль. +- **Если разместить сообщение не удалось — завершить ошибкой и его, и все последующие.** Выбросить + значит молча потерять данные; пропустить и продолжить дальше значит переставить сообщения местами. + Завалить хвост — единственный вариант, сохраняющий оба инварианта. + +### 9.6 Merge + +Устроен так же, только у ребёнка merge **два** родителя. Поиск детей по одному родителю его +находит, и дальше все родители этого ребёнка, которых мы ещё держим, снимаются с маршрутизации +**вместе**. Иначе второй родитель останется в наборе с диапазоном, который теперь пересекается с +диапазоном ребёнка, и выбор партиции станет неоднозначным. + +Обе эталонные реализации merge игнорируют — читают только первого родителя. На момент написания +сервер его тоже не реализует, так что этот путь у нас защитный. + +**Сливаться могут только соседние партиции, и это ограничение протокола, а не клиента.** Диапазон +партиции — один `[from_bound, to_bound)` (§5.1). Слияние двух партиций, не соседних в пространстве +ключей, дало бы ребёнку два несвязных куска, а сообщить об этом попросту негде — значит такая +ситуация не возникнет, как бы сервер ни реализовал merge. Именно соседство делает результат +выразимым: + +``` +p1 [a, m) + p2 [m, z) → p3 [a, z) один интервал, его можно сообщить +p1 [a, m) + p3 [t, z) → [a,m) ∪ [t,z) представить это нечем +``` + +Отсюда же следует, что клиенту не нужно отдельное правило, отвергающее ребёнка с дырой: такого +ребёнка нельзя описать в принципе. Обрабатывать клиенту приходится **переход** — момент, когда оба +родителя ещё лежат в его наборе рядом с ребёнком; ради этого все родители найденного ребёнка и +снимаются одним шагом. + +### 9.7 Как выглядит сплит целиком + +Главное, что стоит вынести из этого раздела, — **порядок четырёх обращений наружу**. Их легко +переставить местами, и каждая перестановка ломает что-то своё. + +Ниже — сквозной пример: партиция 0 расщепилась на 2 и 3, у нас в полёте сообщения с seqno до 2600, +сервер успел записать до 2593. + +``` +1. ошибка записи + p0 ──OVERLOADED──► оркестратор + оркестратор останавливает сессию p0 и ставит задачу на обработку сплита + (повторные OVERLOADED по этой же партиции схлопываются в одну задачу) + +2. подтверждение топологии ← обязательно ДО любых изменений + оркестратор ──describe──► сервер + сервер ──► активные листья 2 и 3, у обоих родитель 0 + проверяем, что диапазоны 2 и 3 покрывают диапазон 0 без дыр + +3. правка набора партиций ← атомарно, под локом + добавить 2 и 3, убрать 0 + теперь выбор по ключу не может вернуть 0 + +4. закрытие старой сессии ← ДО чтения среза + оркестратор закрывает сессию p0 и даёт осесть уже пришедшим ack + иначе сообщение может записаться уже после того, как срез прочитан + +5. чтение среза ← отдельная короткоживущая сессия + оркестратор ──► открыть сессию под producer_id "pfx-0" БЕЗ partition_id + сервер ──► last_seqno = 2593 + сессия сразу закрывается; срез = 2593 + +6. разбор неподтверждённых сообщений + seqno <= 2593 → сервер их уже записал → завершить future как «записано» + seqno > 2593 → заново выбрать партицию → отправить ребёнку с ТЕМ ЖЕ seqno + (в примере ключи ведут в p2) +``` + +Что сломается при перестановке: + +| Если сделать | Что произойдёт | +|---|---| +| шаг 3 после шага 6 | переотправка сможет снова выбрать снимаемую партицию | +| шаг 5 до шага 4 | сессия ещё жива и допишет сообщение после того, как срез прочитан → дубль | +| шаг 5 через сессию с `partition_id` | партиция уже неактивна, ответа не будет никогда → зависание под локом | +| шаг 6 со сменой seqno | сообщение нельзя сопоставить с уже записанной попыткой → дубль | + +--- + +## 10. Обычная перегрузка + +Если `describe` так и не показал детей, значит топология не менялась. Партиция остаётся, а сессия +**восстанавливается на месте**: остановленная выбрасывается, открывается новая на ту же партицию, и +неподтверждённые сообщения отправляются повторно со своими прежними seqno. + +Здесь срез получается точным и бесплатно: партиция всё ещё активна, поэтому новая сессия в ответе на +инициализацию сразу сообщает актуальный записанный seqno. Отдельная сессия для среза не нужна. + +Стоит проговорить, зачем здесь вообще нужен срез, потому что очевидное рассуждение говорит, что не +нужен: партиция та же, `producer_id` тот же, значит если сервер что-то из этих сообщений уже +сохранил, он распознает повтор и отсеет его. Казалось бы, переотправлять всё подряд безопасно. + +Небезопасно, и причина к серверу отношения не имеет: + +``` +ack на seqno 42 потерялся вместе с оборванной сессией +восстановление открывает новую сессию на ту же партицию + в ответе на инициализацию приходит last_seq_no = 42 ← у сервера оно есть + писатель берёт 42 как свою верхнюю отметку +отправляем 42 → отвергнуто нашим же писателем, «seqno is duplicated» + сообщение не покидает процесс +отправляем 43, 44, ... → до них дело не доходит: цикл оборвался на ошибке +``` + +Под-писатели работают с выключенной автоматической выдачей номеров (§6.2), а такой писатель +отвергает любой seqno не выше верхней отметки, которую узнал при инициализации. То есть сообщение +останавливают на клиенте, и сервер не получает шанса его дедуплицировать. А поскольку переотправка +идёт одним упорядоченным проходом, этот единственный отказ роняет и все сообщения, стоящие за ним: +их не отправят и не завершат. + +Срез убирает проблему в корне: seqno 42 не выше него, поэтому сообщение сразу помечается как +записанное, а проход продолжается с 43. + +--- + +## 11. Многопоточность + +**Один общий лок** выстраивает в очередь: выбор партиции, создание под-писателей, выдачу seqno, учёт +неподтверждённых, обновление топологии, переотправку и закрытие по простою. + +Это грубо и сделано намеренно: выбор партиции и изменение топологии не должны переплетаться. Плата в +том, что любое долгое действие под этим локом останавливает писатель целиком, — поэтому каждое +сетевое ожидание под локом ограничено по времени (§6.2, §9.4). + +Обработка сплита идёт фоновой задачей, и этой задачей надо **владеть**: + +- **схлопывать по партиции** — серия `OVERLOADED` не должна запускать несколько параллельных + обработок одной и той же партиции; +- **отменять и дожидаться при закрытии** — иначе она продолжит дёргать `describe` и открывать сессии + для писателя, который вызывающий считает закрытым; +- **не запускать после закрытия** — запоздалый сигнал не должен ничего воскрешать. + +--- + +## 12. Чем может закончиться сообщение + +Каждое принятое сообщение приходит ровно к одному исходу: + +| Исход | Когда | +|---|---| +| записано, офсет известен | обычный ack | +| записано, офсет неизвестен | сообщение не выше среза дедупликации — сервер его записал, а ack потерялся | +| ошибка: некуда положить | ни одна готовая партиция не владеет ключом, либо её не удалось открыть | +| ошибка: конфликт seqno | номер, заданный клиентом, столкнулся с другим в целевой партиции | +| ошибка: партиция непригодна | и обработка сплита, и восстановление на месте провалились | +| ошибка: писатель остановлен | закрытие, когда сообщения ещё не подтверждены | + +Исход «офсет неизвестен» — настоящая часть контракта, а не заглушка: запись состоялась, но офсет +вернулся бы по потоку, который к тому моменту умер. Обе эталонные реализации поступают так же — C++ +специально формирует для этого случая искусственный ack. + +--- + +## 13. Ловушки, по цене обнаружения + +1. **Сессия с явным `partition_id` неактивной партиции висит вечно.** Цена: замороженный писатель, + если это происходит под локом. Обнаруживается только на настоящем сплите. +2. **Срез по увиденным ack неверен.** Цена: дубли, редко и невоспроизводимо. +3. **Закрытие сессии, остановленной хуком, заново бросает причину остановки.** Закрывая сессию, + которую вы и так выбрасываете, ошибки надо глотать — иначе обработка сплита сорвётся на уборке. +4. **Подавляя ошибки ack при закрытии сессии, не подавите успешные ack.** Успех, пришедший в момент + закрытия, настоящий; потеряв его, вы занизите срез и продублируете сообщение. +5. **`describe` посреди сплита может показать одного ребёнка из двух.** Зафиксируете такую картину — + размажете ключ по соседним веткам. +6. **Поиск по нижней границе никогда не падает, он просто отвечает неправильно**, если в наборе + партиций дыра. И сервер этого не поймает. +7. **Автопартиционированный топик с одной партицией может не сообщить диапазон**, из-за чего + автоматический выбор способа маршрутизации ошибётся. +8. **Возможно, этим писателем вообще не получится расщепить топик.** На кластере, где мы + проверяли, нагрузка через мультиписатель автопартиционирование не запускала ни разу, а + отдельная нагрузка обычными писателями, без явного `partition_id`, расщепляла топик за секунды. Общее это свойство или + особенность той сборки — мы не выясняли, но при планировании тестов на это стоит закладываться: + вызывать сплит отдельным продюсером, а не ждать, что его вызовет проверяемый писатель. +9. **Заглушки в тестах прячут всё перечисленное.** Мок, который всегда принимает любой seqno, всегда + успешно инициализируется и никогда не бросает при закрытии, оставит каждую из этих ошибок + зелёной. Моделируйте настоящее поведение: отказ по верхней отметке seqno, инициализацию, которая + не завершается никогда, закрытие, которое заново бросает ошибку, и состояние, привязанное к + `producer_id`, а не к партиции. + +--- + +## 14. Сверка реализаций + +| Аспект | Здесь | C++ `TProducer` | Go `topicmultiwriter` | +|---|---|---|---| +| `producer_id` | на партицию | на партицию | на партицию | +| счётчик seqno | один на писатель | один на писатель | один на писатель | +| переотправка сохраняет seqno | да | да | да | +| откуда берётся срез | с сервера, только снимаемая партиция | по всей цепочке предков | по всей цепочке предков | +| чем читается срез | сессия без `partition_id` | так же | так же | +| сигнал сплита | `OVERLOADED` | `OVERLOADED` | `OVERLOADED` | +| хеш ключа | murmur64a, 8 байт big-endian | так же | так же | +| маска Kafka `toPositive` | да | не применимо | нет | +| merge | обрабатывается | считается, что родитель один | считается, что родитель один | +| обычная перегрузка | восстановление на месте | — | останавливает писатель ошибкой | +| проверка покрытия перед снятием | да | нет | нет | +| проверка верхней границы при выборе | да | нет | нет | + +Последние четыре строки — места, где эта реализация сознательно строже эталонов. + +--- + +## 15. Что осталось незакрытым + +- Пакетная запись принимает сообщения по одному. Падение в середине пакета может оставить более + ранние сообщения уже принятыми, при этом сам вызов завершится исключением и не вернёт на них ни + одного future. +- Ограничения буфера заданы для каждого под-писателя по отдельности, поэтому фактический объём + растёт вместе с числом открытых партиций, вместо того чтобы принадлежать писателю целиком. +- Покрытие сплита в автотестах зависит от того, сумеет ли окружение устроить настоящий сплит. + Окружение, которое этого не умеет, надо отличать от сплита, который должен был случиться и не + случился. diff --git a/docs/topic.rst b/docs/topic.rst index bd5455691..965f98782 100644 --- a/docs/topic.rst +++ b/docs/topic.rst @@ -259,6 +259,71 @@ For high-throughput pipelines, buffer writes and gather futures: raise f.exception() +Writing by Key (Multiple Partitions) +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +A regular writer targets a single partition for its whole lifetime. To spread load across +all partitions of a topic from one logical writer — while keeping every message with the same +key on the same partition (so per-key ordering is preserved) — use ``topic_client.multiwriter()``. + +Each message carries a ``key``; the writer hashes it and routes the message to the owning +partition, maintaining a separate underlying writer per partition. This is the client-side +companion to auto-partitioning. + +**Synchronous:** + +.. code-block:: python + + with driver.topic_client.multiwriter("/local/my-topic") as writer: + writer.write(ydb.TopicWriterMessage(data="event", key="user-42")) + writer.write(ydb.TopicWriterMessage(data="event", key="user-7")) + +**Asynchronous:** + +.. code-block:: python + + async with driver.topic_client.multiwriter("/local/my-topic") as writer: + await writer.write(ydb.TopicWriterMessage(data="event", key="user-42")) + +**Partition choosers** decide how a key maps to a partition. By default the writer picks one +automatically after describing the topic — the key-range chooser for auto-partitioned topics, +the Kafka-hash chooser otherwise — so ``multiwriter(topic)`` works out of the box. You can also +set one explicitly: + +* :class:`~ydb.TopicWriterPartitionByKeyBound` — hashes the key and selects the partition whose + server-side key range owns it. Mirrors YDB auto-partitioning, so a key lands where the server + expects it. Used automatically for auto-partitioned topics. +* :class:`~ydb.TopicWriterPartitionByKeyKafka` — ``murmur2(key) % partitions_count``, + Kafka-compatible. Best for topics with a fixed partition count. + +.. code-block:: python + + writer = driver.topic_client.multiwriter( + "/local/my-topic", + partition_chooser=ydb.TopicWriterPartitionByKeyBound(), + producer_id_prefix="my-app", # each partition writer uses "-" + ) + +The multi-writer accepts the same ``codec``, ``encoders``, ``auto_seqno``, ``auto_created_at`` +and buffer-limit parameters as :meth:`writer`, and exposes ``write``, ``write_with_ack``, +``flush`` and ``close`` with the same semantics. ``wait_init()`` differs: it waits until the +topic has been described and the partition set is known, and (unlike the single-partition writer) +returns nothing, because the multi-writer manages a stream per partition rather than one stream. + +.. note:: + + When an auto-partitioned partition is split (one into two) or merged (two into one), the + multi-writer re-describes the topic, routes subsequent keys to the new child partition(s), and + transparently resends the messages that were still in flight to the retired partition(s). + + Which of those messages to resend is decided by asking the server how far the retired + partition's producer actually got, not by looking at the acknowledgements the client happened + to receive: a message can be persisted and its acknowledgement lost together with the session + that a split tears down. Anything at or below that point is reported as written instead of + being sent again, so a split produces neither loss nor duplicates, and per-key ordering is + preserved throughout. + + Writer Backpressure ^^^^^^^^^^^^^^^^^^^ diff --git a/examples/topic/multiwriter_example.py b/examples/topic/multiwriter_example.py new file mode 100644 index 000000000..9363d8d93 --- /dev/null +++ b/examples/topic/multiwriter_example.py @@ -0,0 +1,59 @@ +"""Examples for the topic multi-partition writer (write-by-key). + +The multi-writer routes each message to a partition by its ``key``, keeping every +message with the same key on the same partition (per-key ordering). By default it +picks a partition chooser automatically: the key-range chooser for +auto-partitioned topics, the Kafka-hash chooser otherwise. +""" + +import asyncio + +import ydb + + +def write_by_key_sync(db: ydb.Driver, topic_path: str): + with db.topic_client.multiwriter(topic_path, producer_id_prefix="my-app") as writer: + writer.write(ydb.TopicWriterMessage(data="event-a", key="user-42")) + writer.write(ydb.TopicWriterMessage(data="event-b", key="user-7")) + # messages with the same key always land on the same partition, in order + writer.write(ydb.TopicWriterMessage(data="event-c", key="user-42")) + writer.flush() + + +async def write_by_key_async(db: ydb.aio.Driver, topic_path: str): + async with db.topic_client.multiwriter(topic_path, producer_id_prefix="my-app") as writer: + await writer.write(ydb.TopicWriterMessage(data="event-a", key="user-42")) + # wait for the server ack of a specific message + await writer.write_with_ack(ydb.TopicWriterMessage(data="event-b", key="user-7")) + + +def write_by_key_with_explicit_chooser(db: ydb.Driver, topic_path: str): + # Force the Kafka-compatible hash chooser (murmur2(key) % partitions_count). + with db.topic_client.multiwriter( + topic_path, + partition_chooser=ydb.TopicWriterPartitionByKeyKafka(), + ) as writer: + writer.write(ydb.TopicWriterMessage(data="event", key="user-42")) + + +def run_sync(): + with ydb.Driver( + connection_string="grpc://localhost:2135?database=/local", + credentials=ydb.credentials.AnonymousCredentials(), + ) as db: + db.wait(timeout=5, fail_fast=True) + write_by_key_sync(db, "/local/topic") + + +async def run_async(): + async with ydb.aio.Driver( + connection_string="grpc://localhost:2135?database=/local", + credentials=ydb.credentials.AnonymousCredentials(), + ) as db: + await db.wait(timeout=5, fail_fast=True) + await write_by_key_async(db, "/local/topic") + + +if __name__ == "__main__": + run_sync() + asyncio.run(run_async()) diff --git a/tests/slo/README.md b/tests/slo/README.md index 88d46e724..cee02e912 100644 --- a/tests/slo/README.md +++ b/tests/slo/README.md @@ -24,10 +24,33 @@ async (`ydb.aio`) path: | `async-query` | Query service | async | | `sync-topic` | Topic service | sync | | `async-topic` | Topic service | async | +| `sync-topic-multiwriter` | Topic service, writes by key | sync | > The `--async` CLI flag is kept as a manual override for `*-run` commands. > The bare `topic` label is still accepted as an alias for `sync-topic`. +### Which scenarios a CI run starts + +The `SLO` label on a PR starts the workflow, but not necessarily every scenario: a full run is +long and occupies external runners, so `.github/workflows/slo.yml` first computes the matrix from +the changed files. + +| Changed | Scenarios | +|---|---| +| `ydb/_topic_*/**`, `ydb/topic.py`, `ydb/aio/topic.py`, `ydb/_grpc/grpcwrapper/ydb_topic*.py` | the three topic ones | +| `ydb/query/**`, `ydb/aio/query/**` | `sync-query`, `async-query` | +| `ydb/table.py`, `ydb/aio/table.py`, `ydb/_session_impl.py` | `sync-table` | +| anything else under `ydb/` — driver, pool, connection, retries, credentials, generated stubs, and shared helpers such as `convert.py` | all | +| `tests/slo/**` or `.github/workflows/slo.yml` | all | +| nothing SLO-relevant (docs, examples, packaging, other CI) | all | + +The last row is deliberate rather than a fallback for its own sake: putting the label on such a PR +is an explicit request, so it is honoured — and it doubles as the way to force a full run. Touch +any file outside the SDK, add the label, and every scenario starts. + +The job logs both the file list it saw and the scenarios it picked, so a surprising selection can +be diagnosed without re-running anything. + ### Usage: Each workload type has 3 commands: @@ -267,7 +290,18 @@ When running `topic-run` (`sync-topic` / `async-topic`), the program creates `re - a **backward** seqno (already seen) is a **duplicate** — reconnect redelivery; with producer-id dedup it should stay near zero (informational); - **end-to-end latency** is `read_ts − write_ts` for the first delivery of each message (writer and reader share the process, so the timestamps are comparable). -Each message carries `writer_id:seqno:write_ts_ns:` followed by padding to the configured size. Topics are scoped per ref so the current and baseline containers (same cluster, run in parallel) don't share a topic. +Each message carries `writer_id:seqno:write_ts_ns:` followed by padding to the configured size. Topics are scoped per workload and per ref, so neither two workloads nor the current and baseline containers (same cluster, run in parallel) share a topic — sharing one would mix their producers into each other's delivery and ordering accounting. + +### Topic multi-writer workload + +`sync-topic-multiwriter` runs the same read side and the same accounting, but writes through `topic_client.multiwriter(...)`: each message carries a `key` and the writer picks the partition itself. That covers what the pinned-writer workload cannot — key routing, the per-partition sub-writer pool, and the recovery of both when nodes disappear under chaos. + +- `writeJob` — each thread owns one multi-writer and cycles through `--keys-per-writer` distinct keys. One payload stream per **key** (not per partition), because a key is what the writer keeps ordered, and two keys may share a partition. Use fewer write threads than the pinned workload: every thread opens a sub-writer per partition it touches. +- `readJob` — unchanged. A forward gap is still loss and a backward seqno is still a duplicate; the streams it validates are keys rather than partition-pinned producers. + +Enabled by `--use-multiwriter` on `topic-run`; the label sets it automatically. There is no async variant yet — the sync facade drives the same async implementation underneath, so the code under test is the same. + +> Local re-runs against an existing topic will report duplicates: per-key seqno counters restart at 1 with the process, while the reader still expects the sequence left by the previous run. Recreate the topic (`topic-cleanup` + `topic-create`) between local runs. CI is unaffected — every run gets its own topic. ## Collected metrics - `oks` - amount of OK requests diff --git a/tests/slo/docker-entrypoint.sh b/tests/slo/docker-entrypoint.sh index 55d933ec3..edb5b6b1a 100755 --- a/tests/slo/docker-entrypoint.sh +++ b/tests/slo/docker-entrypoint.sh @@ -19,9 +19,13 @@ set -e +WORKLOAD_ARGS="" case "${WORKLOAD_NAME:-sync-query}" in sync-table|sync-query|async-query) PREFIX=table ;; topic|sync-topic|async-topic) PREFIX=topic ;; + # Same topic workload, but writes go through the multi-partition writer (routing by key + # plus its sub-writer pool) instead of one writer pinned per partition. + sync-topic-multiwriter) PREFIX=topic; WORKLOAD_ARGS="--use-multiwriter" ;; *) echo "Unknown WORKLOAD_NAME: ${WORKLOAD_NAME}" >&2 exit 1 @@ -37,11 +41,15 @@ DURATION="${WORKLOAD_DURATION:-600}" # Scope the topic by ref so the current and baseline containers (same cluster, # run in parallel) don't share a topic — otherwise their readers/producers would # cross-contaminate delivery/ordering validation. +# The workload name is part of the path too: topic workloads differ in what they write +# (producer ids, key layout), so two of them sharing a topic would corrupt each other's +# delivery and ordering accounting. EXTRA_ARGS="" if [ "$PREFIX" = "topic" ]; then REF_RAW="${WORKLOAD_REF:-${REF:-main}}" SAFE_REF=$(printf '%s' "$REF_RAW" | tr -c 'a-zA-Z0-9_' '_') - EXTRA_ARGS="--path ${DATABASE%/}/slo_topic_${SAFE_REF}" + SAFE_WORKLOAD=$(printf '%s' "${WORKLOAD_NAME:-topic}" | tr -c 'a-zA-Z0-9_' '_') + EXTRA_ARGS="--path ${DATABASE%/}/slo_topic_${SAFE_WORKLOAD}_${SAFE_REF}" fi # Schema prep is idempotent at the SDK level for topics; for tables, a parallel @@ -53,4 +61,5 @@ exec python ./tests/slo/src \ "${PREFIX}-run" "$ENDPOINT" "$DATABASE" \ --time "$DURATION" \ $EXTRA_ARGS \ + $WORKLOAD_ARGS \ "$@" diff --git a/tests/slo/src/jobs/topic_multiwriter_jobs.py b/tests/slo/src/jobs/topic_multiwriter_jobs.py new file mode 100644 index 000000000..3818d4008 --- /dev/null +++ b/tests/slo/src/jobs/topic_multiwriter_jobs.py @@ -0,0 +1,112 @@ +"""Topic SLO workload variant that writes through the multi-partition writer. + +The regular topic workload pins one writer per partition and drives it directly. This one +writes by key instead and lets the multi-writer decide the partition, so the chaos run +exercises the parts that only exist there: routing, the per-partition sub-writer pool, and +recovery of both when nodes go away underneath. + +Everything else is inherited: the reader side, the metrics, and the delivery accounting. +That works unchanged because the accounting only needs a stream whose sequence numbers grow +monotonically, and per-key ordering gives exactly that -- one payload stream per key rather +than per partition. +""" + +import logging +import threading +import time + +from core.metrics import OP_TYPE_WRITE, REF + +import ydb + +from .base import SyncRateLimiter +from .topic_jobs import TopicJobManager +from .topic_payload import encode_payload + +logger = logging.getLogger(__name__) + + +class TopicMultiWriterJobManager(TopicJobManager): + def __init__(self, driver, args, metrics): + super().__init__(driver, args, metrics) + self.keys_per_writer = max(1, int(getattr(self.args, "keys_per_writer", 8))) + + def _run_topic_write_jobs(self): + logger.info( + "Start topic multi-writer jobs: %d writers x %d keys", + self.args.write_threads, + self.keys_per_writer, + ) + + write_rps = int(getattr(self.args, "write_rps", 0)) + write_limiter = SyncRateLimiter(min_interval_s=0.0 if write_rps <= 0 else 1.0 / write_rps) + + futures = [] + for i in range(self.args.write_threads): + future = threading.Thread( + name=f"slo_topic_multiwrite_{i}", + target=self._run_topic_writes, + args=(i, write_limiter), + ) + future.start() + futures.append(future) + return futures + + def _stream_id(self, writer_id: int, key_index: int) -> int: + """Payload stream id for one key. + + The reader validates ordering per stream id, so it has to be one id per key: a key is + what the multi-writer keeps ordered, and two keys may share a partition. + """ + return writer_id * self.keys_per_writer + key_index + + def _run_topic_writes(self, writer_id, limiter): + start_time = time.time() + producer_id_prefix = f"{REF}-mw{writer_id}" + write_timeout = self.args.write_timeout / 1000 + keys = [f"slo-key-{self._stream_id(writer_id, j)}" for j in range(self.keys_per_writer)] + + logger.info("Start topic multi-writer %s (prefix %s, keys %s)", writer_id, producer_id_prefix, len(keys)) + + # Sequence numbers live across writer recreations so the reader never sees a stream + # restart, exactly as in the single-partition workload. + seqno = {key: 1 for key in keys} + next_key = 0 + + while time.time() - start_time < self.args.time: + try: + with self.driver.topic_client.multiwriter( + self.args.path, + producer_id_prefix=producer_id_prefix, + codec=ydb.TopicCodec.RAW, + ) as writer: + while time.time() - start_time < self.args.time: + with limiter: + key_index = next_key % len(keys) + key = keys[key_index] + next_key += 1 + + payload = encode_payload( + self._stream_id(writer_id, key_index), + seqno[key], + time.monotonic_ns(), + self.args.message_size, + ) + message = ydb.TopicWriterMessage(data=payload, key=key) + + ts = self.metrics.start((OP_TYPE_WRITE,)) + try: + writer.write_with_ack(message, timeout=write_timeout) + self.metrics.stop((OP_TYPE_WRITE,), ts) + # Advance only on success: a failed write retries the same + # seqno, so at worst it duplicates, never fakes a loss. + seqno[key] += 1 + except Exception as e: + self.metrics.stop((OP_TYPE_WRITE,), ts, error=e) + logger.error("Multi-writer write error (recreating writer): %s", e) + break # drop the possibly wedged writer and remake it + except Exception as e: + logger.error("Topic multi-writer %s recreate: %s", writer_id, e) + time.sleep(0.2) + + logger.info("Stop topic multi-writer %s", writer_id) diff --git a/tests/slo/src/options.py b/tests/slo/src/options.py index 200c3c556..2ef5eeaed 100644 --- a/tests/slo/src/options.py +++ b/tests/slo/src/options.py @@ -130,6 +130,17 @@ def make_topic_run_parser(subparsers): help="Number of threads for topic writing", ) topic_parser.add_argument("--message-size", default=100, type=int, help="Topic message size in bytes") + topic_parser.add_argument( + "--use-multiwriter", + action="store_true", + help="Write by key through the multi-partition writer instead of one writer per partition", + ) + topic_parser.add_argument( + "--keys-per-writer", + default=8, + type=int, + help="Distinct routing keys each multi-writer thread cycles through (--use-multiwriter only)", + ) topic_parser.add_argument("--time", default=10, type=int, help="Time to run in seconds") topic_parser.add_argument( diff --git a/tests/slo/src/runners/topic_runner.py b/tests/slo/src/runners/topic_runner.py index c9a8bdaa0..4ff0f127d 100644 --- a/tests/slo/src/runners/topic_runner.py +++ b/tests/slo/src/runners/topic_runner.py @@ -3,6 +3,7 @@ from core.metrics import create_metrics from jobs.async_topic_jobs import AsyncTopicJobManager from jobs.topic_jobs import TopicJobManager +from jobs.topic_multiwriter_jobs import TopicMultiWriterJobManager import ydb import ydb.aio @@ -72,9 +73,11 @@ def run(self, args): assert self.driver is not None, "Driver is not initialized. Call set_driver() before run()." metrics = create_metrics(args.otlp_endpoint) - self.logger.info("Starting topic SLO tests") + use_multiwriter = getattr(args, "use_multiwriter", False) + self.logger.info("Starting topic SLO tests (multiwriter=%s)", use_multiwriter) - job_manager = TopicJobManager(self.driver, args, metrics) + manager_cls = TopicMultiWriterJobManager if use_multiwriter else TopicJobManager + job_manager = manager_cls(self.driver, args, metrics) job_manager.run_tests() self.logger.info("Topic SLO tests completed") @@ -85,6 +88,10 @@ def run(self, args): async def run_async(self, args): """Async version of topic SLO tests using ydb.aio.Driver""" assert self.driver is not None, "Driver is not initialized. Call set_driver() before run_async()." + if getattr(args, "use_multiwriter", False): + # Silently running the single-partition workload under a multiwriter label would + # make the SLO report compare something other than what it claims. + raise NotImplementedError("--use-multiwriter has no async workload yet; run it as a sync-* workload") metrics = create_metrics(args.otlp_endpoint) self.logger.info("Starting async topic SLO tests") diff --git a/tests/slo/thresholds-topic.yaml b/tests/slo/thresholds-topic.yaml index c90edaead..2b9a6ba3d 100644 --- a/tests/slo/thresholds-topic.yaml +++ b/tests/slo/thresholds-topic.yaml @@ -30,3 +30,25 @@ metrics: direction: neutral - name: topic_e2e_latency_p99_ms direction: neutral + + # Retry attempts: keep the absolute gate, drop the relative comparison. + # + # The metric is zero almost everywhere and spikes only in the seconds around a chaos fault. + # Comparing two such series relatively divides one accidental spike by another: across two runs + # this flagged sync-table (▲50%, sums 12 vs 14) and then async-topic (▲450%, sums 75 vs 12), + # each time from 2 non-zero samples out of 120, with the report's own `concordance` at 0.008 -- + # i.e. the two series almost never have data at the same moment. Whichever workload draws the + # bigger spike gets marked red, which is noise, not a regression signal. + # + # `direction: neutral` disables exactly the relative check (evaluateRelativeThreshold only ever + # flags `lower_is_better` / `higher_is_better`), while the absolute check still runs for both + # bounds. `warning_max` is restated on purpose: findMatchingThreshold returns the FIRST matching + # entry whole and an exact `name` beats a `pattern`, so this entry replaces the action's + # `*_attempts` default rather than merging with it -- omit it and the "retries must stay at + # zero" gate silently disappears along with the noise. + - name: read_retry_attempts + direction: neutral + warning_max: 0.0 + - name: write_retry_attempts + direction: neutral + warning_max: 0.0 diff --git a/tests/topics/test_topic_writer.py b/tests/topics/test_topic_writer.py index 035c3b801..475d62306 100644 --- a/tests/topics/test_topic_writer.py +++ b/tests/topics/test_topic_writer.py @@ -1,6 +1,7 @@ from __future__ import annotations import asyncio +import datetime from typing import List # noqa: F401 import pytest @@ -324,3 +325,182 @@ class TestException(Exception): writer.write_with_ack("123") raise TestException() + + +@pytest.mark.asyncio +class TestTopicMultiWriterAsyncIO: + async def _recreate(self, driver, path, consumer, **kwargs): + try: + await driver.topic_client.drop_topic(path) + except ydb.SchemeError: + pass + await driver.topic_client.create_topic(path=path, consumers=[consumer], **kwargs) + + def _auto_partitioning(self): + return ydb.TopicAutoPartitioningSettings( + strategy=ydb.TopicAutoPartitioningStrategy.SCALE_UP, + up_utilization_percent=1, + down_utilization_percent=1, + stabilization_window=datetime.timedelta(seconds=1), + ) + + async def test_key_range_exposed_for_autopartitioned_topic(self, driver, database, topic_consumer): + path = database + "/mw-keyrange" + await self._recreate( + driver, + path, + topic_consumer, + min_active_partitions=2, + max_active_partitions=50, + auto_partitioning_settings=self._auto_partitioning(), + ) + desc = await driver.topic_client.describe_topic(path) + assert any(p.key_range is not None for p in desc.partitions) + + async def test_write_by_key_preserves_per_key_order(self, driver, database, topic_consumer): + path = database + "/mw-plain" + await self._recreate(driver, path, topic_consumer, min_active_partitions=3) + + keys = ["user-1", "user-2", "user-3", "user-4", "user-5"] + per_key = 8 + async with driver.topic_client.multiwriter(path, producer_id_prefix="mw") as writer: + await writer.wait_init() + assert isinstance(writer._chooser, ydb.TopicWriterPartitionByKeyKafka) + for i in range(per_key): + for key in keys: + await writer.write(ydb.TopicWriterMessage(data=("%s:%d" % (key, i)).encode(), key=key)) + await writer.flush() + + total = per_key * len(keys) + received = {key: [] for key in keys} + async with driver.topic_client.reader(path, consumer=topic_consumer) as reader: + for _ in range(total): + message = await asyncio.wait_for(reader.receive_message(), timeout=30) + key, index = message.data.decode().split(":") + received[key].append(int(index)) + reader.commit(message) + + for key in keys: + assert received[key] == list(range(per_key)), (key, received[key]) + + async def test_write_by_key_on_autopartitioned_topic(self, driver, database, topic_consumer): + path = database + "/mw-auto" + await self._recreate( + driver, + path, + topic_consumer, + min_active_partitions=2, + max_active_partitions=50, + auto_partitioning_settings=self._auto_partitioning(), + ) + + keys = ["alpha", "beta", "gamma", "delta"] + per_key = 5 + async with driver.topic_client.multiwriter(path, producer_id_prefix="mw") as writer: + await writer.wait_init() + # auto-partitioned topics report key ranges -> adaptive default picks the bound chooser + assert isinstance(writer._chooser, ydb.TopicWriterPartitionByKeyBound) + for i in range(per_key): + for key in keys: + # write_with_ack verifies the server accepts bound-routed writes + await writer.write_with_ack(ydb.TopicWriterMessage(data=("%s:%d" % (key, i)).encode(), key=key)) + + total = per_key * len(keys) + seen = 0 + async with driver.topic_client.reader(path, consumer=topic_consumer) as reader: + for _ in range(total): + message = await asyncio.wait_for(reader.receive_message(), timeout=30) + seen += 1 + reader.commit(message) + assert seen == total + + async def test_write_by_key_survives_partition_split(self, driver, database, topic_consumer): + # Aggressive auto-partitioning + a low write-speed limit force the topic to split + # under load, exercising the resend path. The test keeps writing until a split is + # observed, then asserts exactly-once delivery (no loss, no duplicates). + path = database + "/mw-split" + await self._recreate( + driver, + path, + topic_consumer, + min_active_partitions=1, + max_active_partitions=100, + partition_write_speed_bytes_per_second=1024, + auto_partitioning_settings=self._auto_partitioning(), + ) + + partitions_before = len((await driver.topic_client.describe_topic(path)).partitions) + payload = b"x" * 512 + written = 0 + max_batches = 15 + async with driver.topic_client.multiwriter(path, producer_id_prefix="mw-split") as writer: + for _ in range(max_batches): + for _ in range(100): + await writer.write( + ydb.TopicWriterMessage(data=b"%d:%s" % (written, payload), key="k%d" % (written % 32)) + ) + written += 1 + await writer.flush() + # give the auto-partitioning actuator time to measure and split + await asyncio.sleep(1.5) + if len((await driver.topic_client.describe_topic(path)).partitions) > partitions_before: + break + await writer.flush() + + total = written + partitions_after = len((await driver.topic_client.describe_topic(path)).partitions) + if partitions_after <= partitions_before: + # Auto-partitioning did not split the topic in this environment (single-node + # clusters often don't actuate). The split/resend path is covered deterministically + # by the unit tests; here we only assert exactly-once when a split actually happened. + pytest.skip("topic did not split under load; resend path covered by unit tests") + + seen = set() + duplicates = 0 + async with driver.topic_client.reader(path, consumer=topic_consumer) as reader: + while len(seen) < total: + try: + message = await asyncio.wait_for(reader.receive_message(), timeout=30) + except asyncio.TimeoutError: + break + index = int(message.data.split(b":", 1)[0]) + if index in seen: + duplicates += 1 + seen.add(index) + reader.commit(message) + + assert duplicates == 0, "resend produced duplicate messages" + assert seen == set(range(total)), "some messages were lost (partitions %d->%d)" % ( + partitions_before, + partitions_after, + ) + + +class TestTopicMultiWriterSync: + def test_write_by_key_preserves_per_key_order(self, driver_sync, database, topic_consumer): + path = database + "/mw-sync" + try: + driver_sync.topic_client.drop_topic(path) + except ydb.SchemeError: + pass + driver_sync.topic_client.create_topic(path=path, consumers=[topic_consumer], min_active_partitions=3) + + keys = ["a", "b", "c"] + per_key = 6 + with driver_sync.topic_client.multiwriter(path, producer_id_prefix="mw-sync") as writer: + for i in range(per_key): + for key in keys: + writer.write(ydb.TopicWriterMessage(data=("%s:%d" % (key, i)).encode(), key=key)) + writer.flush() + + total = per_key * len(keys) + received = {key: [] for key in keys} + with driver_sync.topic_client.reader(path, consumer=topic_consumer) as reader: + for _ in range(total): + message = reader.receive_message(timeout=30) + key, index = message.data.decode().split(":") + received[key].append(int(index)) + reader.commit(message) + + for key in keys: + assert received[key] == list(range(per_key)), (key, received[key]) diff --git a/ydb/_grpc/grpcwrapper/ydb_topic.py b/ydb/_grpc/grpcwrapper/ydb_topic.py index 19b38bf01..3609225f1 100644 --- a/ydb/_grpc/grpcwrapper/ydb_topic.py +++ b/ydb/_grpc/grpcwrapper/ydb_topic.py @@ -1657,6 +1657,38 @@ def to_public(self) -> ydb_topic_public_types.PublicDescribeTopicResult: topic_stats=topic_stats, ) + @dataclass + class PartitionKeyRange( + IFromProto[ + Optional["ydb_topic_pb2.PartitionKeyRange"], + Optional["DescribeTopicResult.PartitionKeyRange"], + ], + IToPublic, + ): + # Empty bytes mean an open bound: from_bound == b"" is the start of the + # key space (first partition), to_bound == b"" is the end (last partition). + from_bound: bytes + to_bound: bytes + + @staticmethod + def from_proto( + msg: Optional[ydb_topic_pb2.PartitionKeyRange], + ) -> Optional["DescribeTopicResult.PartitionKeyRange"]: + if msg is None: + return None + return DescribeTopicResult.PartitionKeyRange( + from_bound=msg.from_bound, + to_bound=msg.to_bound, + ) + + def to_public( + self, + ) -> ydb_topic_public_types.PublicDescribeTopicResult.PartitionKeyRange: + return ydb_topic_public_types.PublicDescribeTopicResult.PartitionKeyRange( + from_bound=self.from_bound, + to_bound=self.to_bound, + ) + @dataclass class PartitionInfo( IFromProto[ @@ -1670,6 +1702,7 @@ class PartitionInfo( child_partition_ids: List[int] parent_partition_ids: List[int] partition_stats: Optional["PartitionStats"] + key_range: Optional["DescribeTopicResult.PartitionKeyRange"] @staticmethod def from_proto( @@ -1678,12 +1711,17 @@ def from_proto( if msg is None: return None + key_range = None + if msg.HasField("key_range"): + key_range = DescribeTopicResult.PartitionKeyRange.from_proto(msg.key_range) + return DescribeTopicResult.PartitionInfo( partition_id=msg.partition_id, active=msg.active, child_partition_ids=list(msg.child_partition_ids), parent_partition_ids=list(msg.parent_partition_ids), partition_stats=PartitionStats.from_proto(msg.partition_stats), + key_range=key_range, ) def to_public( @@ -1692,12 +1730,16 @@ def to_public( partition_stats = None if self.partition_stats is not None: partition_stats = self.partition_stats.to_public() + key_range = None + if self.key_range is not None: + key_range = self.key_range.to_public() return ydb_topic_public_types.PublicDescribeTopicResult.PartitionInfo( partition_id=self.partition_id, active=self.active, child_partition_ids=self.child_partition_ids, parent_partition_ids=self.parent_partition_ids, partition_stats=partition_stats, + key_range=key_range, ) @dataclass diff --git a/ydb/_grpc/grpcwrapper/ydb_topic_public_types.py b/ydb/_grpc/grpcwrapper/ydb_topic_public_types.py index afb031d91..b6bda90c7 100644 --- a/ydb/_grpc/grpcwrapper/ydb_topic_public_types.py +++ b/ydb/_grpc/grpcwrapper/ydb_topic_public_types.py @@ -242,6 +242,14 @@ class PublicDescribeTopicResult: auto_partitioning_settings: Optional["PublicAutoPartitioningSettings"] + @dataclass + class PartitionKeyRange: + from_bound: bytes + "Inclusive lower bound of the partition key range; empty bytes mean an open (leftmost) bound" + + to_bound: bytes + "Exclusive upper bound of the partition key range; empty bytes mean an open (rightmost) bound" + @dataclass class PartitionInfo: partition_id: int @@ -259,6 +267,9 @@ class PartitionInfo: partition_stats: Optional["PublicPartitionStats"] "Stats for partition, filled only when include_stats in request is true" + key_range: Optional["PublicDescribeTopicResult.PartitionKeyRange"] = None + "Key range owned by the partition; filled for auto-partitioned topics" + @dataclass class TopicStats: store_size_bytes: int diff --git a/ydb/_grpc/grpcwrapper/ydb_topic_test.py b/ydb/_grpc/grpcwrapper/ydb_topic_test.py index b9e306034..a9765b847 100644 --- a/ydb/_grpc/grpcwrapper/ydb_topic_test.py +++ b/ydb/_grpc/grpcwrapper/ydb_topic_test.py @@ -2,8 +2,11 @@ from google.protobuf.json_format import MessageToDict +# Same version dispatch the module under test uses: the CI matrix runs protobuf v3..v6. +from ydb._grpc.common.protos import ydb_topic_pb2 + from ydb._grpc.grpcwrapper.ydb_topic import OffsetsRange -from .ydb_topic import AlterTopicRequest +from .ydb_topic import AlterTopicRequest, DescribeTopicResult from .ydb_topic_public_types import ( AlterTopicRequestParams, PublicAlterConsumer, @@ -96,3 +99,44 @@ def test_alter_topic_request_from_public_to_proto(): } assert msg_dict == expected_dict + + +def test_partition_key_range_round_trip(): + """A bounded partition must survive proto -> internal -> public unchanged. + + The multi-partition writer routes by these bounds, and an empty bound is meaningful (it + marks an open end of the key space), so it has to stay distinguishable from an absent one. + """ + msg = ydb_topic_pb2.DescribeTopicResult.PartitionInfo( + partition_id=7, + active=True, + child_partition_ids=[8, 9], + parent_partition_ids=[3], + key_range=ydb_topic_pb2.PartitionKeyRange(from_bound=b"\x10", to_bound=b"\x80"), + ) + + internal = DescribeTopicResult.PartitionInfo.from_proto(msg) + assert internal.key_range.from_bound == b"\x10" + assert internal.key_range.to_bound == b"\x80" + + public = internal.to_public() + assert public.partition_id == 7 + assert public.child_partition_ids == [8, 9] + assert public.parent_partition_ids == [3] + assert public.key_range.from_bound == b"\x10" + assert public.key_range.to_bound == b"\x80" + + +def test_partition_without_key_range_stays_none(): + """Topics that are not auto-partitioned report no range at all. + + That is not the same as an open range: it is what tells the writer to route by hash + instead of by bounds, so it must not be turned into empty bounds along the way. + """ + msg = ydb_topic_pb2.DescribeTopicResult.PartitionInfo(partition_id=0, active=True) + + internal = DescribeTopicResult.PartitionInfo.from_proto(msg) + assert internal.key_range is None + assert internal.to_public().key_range is None + + assert DescribeTopicResult.PartitionKeyRange.from_proto(None) is None diff --git a/ydb/_topic_writer/topic_writer.py b/ydb/_topic_writer/topic_writer.py index 23e4cd5a2..5bb409042 100644 --- a/ydb/_topic_writer/topic_writer.py +++ b/ydb/_topic_writer/topic_writer.py @@ -42,6 +42,10 @@ class PublicWriterSettings: # Backpressure is enabled when at least one of the limits above is set. # None = wait indefinitely for buffer space; positive value = raise TopicWriterBufferFullError on timeout. buffer_wait_timeout_sec: Optional[float] = None + # Internal hook used by the multi-partition writer. Called with a connection error + # before the default retry classification; returning True force-stops the writer + # (used to catch OVERLOADED on a split partition). Not part of the public writer() API. + _on_check_retriable_error: Optional[typing.Callable[[BaseException], bool]] = None def __post_init__(self): if self.producer_id is None: @@ -122,6 +126,7 @@ class PublicMessage: created_at: Optional[datetime.datetime] data: "PublicMessage.SimpleSourceType" metadata_items: Optional[Dict[str, "PublicMessage.SimpleSourceType"]] + key: Optional[str] SimpleSourceType = Union[str, bytes] # Will be extend @@ -132,11 +137,15 @@ def __init__( metadata_items: Optional[Dict[str, "PublicMessage.SimpleSourceType"]] = None, seqno: Optional[int] = None, created_at: Optional[datetime.datetime] = None, + key: Optional[str] = None, ): self.seqno = seqno self.created_at = created_at self.data = data self.metadata_items = metadata_items + # Partitioning key: used only by the multi-partition writer to route the + # message to a partition. Ignored by the single-partition writer. + self.key = key @staticmethod def _create_message(data: Message) -> "PublicMessage": @@ -240,6 +249,17 @@ class TopicWriterBufferFullError(TopicWriterError): pass +class TopicWriterPartitionSplitError(TopicWriterRepeatableError): + """Raised internally when the partition targeted by the writer has split. + + Used by the multi-partition writer to stop a per-partition sub-writer so its + messages can be re-routed to the child partitions. + """ + + def __init__(self): + super().__init__("topic writer partition was split") + + def default_serializer_message_content(data: Any) -> bytes: if data is None: return bytes() diff --git a/ydb/_topic_writer/topic_writer_asyncio.py b/ydb/_topic_writer/topic_writer_asyncio.py index 81b9e7191..ca943a71e 100644 --- a/ydb/_topic_writer/topic_writer_asyncio.py +++ b/ydb/_topic_writer/topic_writer_asyncio.py @@ -19,6 +19,7 @@ InternalMessage, TopicWriterStopped, TopicWriterError, + TopicWriterPartitionSplitError, TopicWriterBufferFullError, internal_message_size_bytes, messages_to_proto_requests, @@ -537,6 +538,17 @@ async def _connection_loop(self): return err = issues.ConnectionLost("gRPC stream cancelled") + if self._settings._on_check_retriable_error is not None and self._settings._on_check_retriable_error( + err + ): + logger.debug( + "writer reconnector %s stop connection loop by on_check_retriable_error hook due to %s", + self._id, + err, + ) + self._stop(TopicWriterPartitionSplitError()) + return + err_info = check_retriable_error(err, retry_settings, attempt) if not err_info.is_retriable or self._tx is not None: # no retries in tx writer logger.debug("writer reconnector %s stop connection loop due to %s", self._id, err) diff --git a/ydb/_topic_writer/topic_writer_asyncio_test.py b/ydb/_topic_writer/topic_writer_asyncio_test.py index 5c77fc2f1..2f207f495 100644 --- a/ydb/_topic_writer/topic_writer_asyncio_test.py +++ b/ydb/_topic_writer/topic_writer_asyncio_test.py @@ -8,6 +8,7 @@ import gzip import sys import typing +import weakref from concurrent.futures import ThreadPoolExecutor from queue import Queue, Empty from typing import List, Callable, Optional @@ -36,6 +37,9 @@ PublicWriteResult, TopicWriterError, TopicWriterBufferFullError, + TopicWriterClosedError, + TopicWriterPartitionSplitError, + TopicWriterStopped, ) from .._grpc.grpcwrapper.ydb_topic_public_types import PublicCodec from .._topic_common.test_helpers import StreamMock, wait_for_fast @@ -45,6 +49,16 @@ WriterAsyncIOReconnector, WriterAsyncIO, ) +from . import topic_writer_multi_asyncio +from .topic_writer_multi_asyncio import TopicWriterMultiAsyncIO, MultiWriterSettings +from .topic_writer_partition_chooser import ( + PublicPartitionByKeyKafka, + PublicPartitionByKeyBound, + PublicPartitionChooser, + PARTITION_KEY_METADATA_KEY, + murmur2_32, +) +from .._grpc.grpcwrapper.ydb_topic_public_types import PublicDescribeTopicResult from ..credentials import AnonymousCredentials @@ -522,6 +536,56 @@ async def wait_stop(): with pytest.raises(TestException): await reconnector.close(flush=False) + async def test_retriable_error_hook_stops_the_writer(self, default_driver, default_settings, get_stream_writer): + """The hook must be able to end the connection loop on an otherwise retriable error. + + OVERLOADED normally means "back off and retry", but for the multi-partition writer it is + also how a split announces itself: the partition went inactive and retrying against it + would spin forever. The hook lets the owner claim such an error and stop the writer with a + reason it can recognise, instead of the generic retry path swallowing it. + """ + seen = [] + + def hook(err): + seen.append(err) + return True + + settings = copy.deepcopy(default_settings) + settings._on_check_retriable_error = hook + reconnector = WriterAsyncIOReconnector(default_driver, settings) + + get_stream_writer().from_server.put_nowait(issues.Overloaded("Write to inactive partition 0")) + + with pytest.raises(TopicWriterPartitionSplitError): + + async def wait_stop(): + while True: + await reconnector.write_with_ack_future([PublicMessage(data="123", seqno=3)]) + await asyncio.sleep(0.01) + + await asyncio.wait_for(wait_stop(), 2) + + assert seen and isinstance(seen[0], issues.Overloaded) + + with pytest.raises(TopicWriterPartitionSplitError): + await reconnector.close(flush=False) + + async def test_retriable_error_hook_declining_keeps_default_retry( + self, default_driver, default_settings, get_stream_writer + ): + """A hook that declines must leave the normal retry policy untouched.""" + settings = copy.deepcopy(default_settings) + settings._on_check_retriable_error = lambda err: False + reconnector = WriterAsyncIOReconnector(default_driver, settings) + + get_stream_writer().from_server.put_nowait(issues.Overloaded("ordinary overload")) + await reconnector.write_with_ack_future([PublicMessage(data="123", seqno=3)]) + await asyncio.sleep(0.1) + + # Retriable error + declining hook -> the writer reconnected instead of stopping. + assert not reconnector._stop_reason.done() + await reconnector.close(flush=False) + async def test_wait_init(self, default_driver, default_settings, get_stream_writer): init_seqno = 100 expected_init_info = PublicWriterInitInfo(last_seqno=init_seqno, supported_codecs=[]) @@ -1146,3 +1210,1752 @@ async def test_writer_create_failure_does_not_leak_grpc_thread(): finally: channel.close() server.stop() + + +class _PublicDescription: + def __init__(self, partitions): + self.partitions = partitions + + +class _MultiFakeDescribeDriver: + """Fake driver that answers DescribeTopic with a sequence of descriptions.""" + + _credentials = AnonymousCredentials() + + def __init__(self, descriptions): + self._descriptions = list(descriptions) + self.describe_calls = 0 + + async def __call__(self, request, stub, method, wrapper=None, *args, **kwargs): + idx = min(self.describe_calls, len(self._descriptions) - 1) + self.describe_calls += 1 + description = _PublicDescription(self._descriptions[idx]) + + class _Result: + def to_public(self): + return description + + return _Result() + + +_real_sleep = asyncio.sleep + + +async def _no_sleep(_delay): + """Run timer-driven loops at full speed. Holds the real sleep so patching + asyncio.sleep with this does not make it call itself.""" + await _real_sleep(0) + + +# Per-partition last persisted seqno seen by the fakes' wait_init(); tests mutate it. +_FAKE_LAST_SEQNO: dict = {} + + +class _FakeSubWriter: + """Stand-in for a per-partition WriterAsyncIO used by the multi-writer. + + Acks every write immediately. + """ + + def __init__(self, driver, settings): + self.settings = settings + self.partition_id = settings.partition_id + self.producer_id = settings.producer_id + self.split_hook = settings._on_check_retriable_error + self.messages: List = [] + self.closed = False + # The server keys persisted state by producer id, not by the partition a session happens + # to be pinned to: that is why an unpinned probe session can still report the last_seqno + # of a partition that has already gone inactive. Keying the fake on partition_id instead + # would make every probe read 0 and hide whether the dedup cut works at all. + self.producer_partition_id = int(settings.producer_id.rsplit("-", 1)[-1]) + + async def wait_init(self): + last_seqno = _FAKE_LAST_SEQNO.get(self.producer_partition_id, 0) + return PublicWriterInitInfo(last_seqno=last_seqno, supported_codecs=[]) + + async def write_with_ack_future(self, message): + self.messages.append(message) + future = asyncio.get_running_loop().create_future() + future.set_result(PublicWriteResult.Written(offset=len(self.messages))) + return future + + async def flush(self): + pass + + async def close(self, flush=True): + self.closed = True + + +class _ControllableSubWriter(_FakeSubWriter): + """Sub-writer whose acks are resolved manually, to test split-resend.""" + + def __init__(self, driver, settings): + super().__init__(driver, settings) + self.pending: List = [] + + async def write_with_ack_future(self, message): + self.messages.append(message) + future = asyncio.get_running_loop().create_future() + self.pending.append(future) + return future + + def resolve_all(self): + for i, future in enumerate(self.pending): + if not future.done(): + future.set_result(PublicWriteResult.Written(offset=i)) + + +class _KeyMapChooser(PublicPartitionChooser): + """Deterministic chooser: routes by message key via a caller-controlled map.""" + + def __init__(self, mapping): + self._mapping = mapping + self.partitions = set() + + def add_partitions(self, partitions): + for p in partitions: + self.partitions.add(p.partition_id) + + def remove_partition(self, partition_id): + self.partitions.discard(partition_id) + + def choose_partition(self, message): + return self._mapping[message.key] + + +class _FlushControlledSubWriter(_FakeSubWriter): + """Sub-writer that acks buffered messages only when flush() is called.""" + + def __init__(self, driver, settings): + super().__init__(driver, settings) + self.pending: List = [] + + async def write_with_ack_future(self, message): + self.messages.append(message) + future = asyncio.get_running_loop().create_future() + self.pending.append(future) + return future + + async def flush(self): + for i, future in enumerate(self.pending): + if not future.done(): + future.set_result(PublicWriteResult.Written(offset=i)) + + +class _RaisingSubWriter(_FakeSubWriter): + """Sub-writer whose admission always fails.""" + + async def write_with_ack_future(self, message): + raise RuntimeError("admission failed") + + +class _CloseRaisesSubWriter(_ControllableSubWriter): + """Sub-writer whose close() re-raises the split stop reason (like a real hook-stopped writer).""" + + async def close(self, flush=True): + self.closed = True + raise TopicWriterPartitionSplitError() + + +class _SeqnoGuardSubWriter(_ControllableSubWriter): + """Sub-writer that enforces the real writer's explicit-seqno guard. + + With auto_seqno=False (what the multi-writer always uses) WriterAsyncIO seeds + _last_known_seq_no from the server's last_seqno at init and then rejects any message with + seq_no <= it -- see _prepare_internal_messages in topic_writer_asyncio.py. The plain fakes + accept every seqno, which hides resend bugs on that boundary. + """ + + def __init__(self, driver, settings): + super().__init__(driver, settings) + self._last_known_seqno: Optional[int] = None + + async def wait_init(self): + info = await super().wait_init() + if self._last_known_seqno is None: + self._last_known_seqno = info.last_seqno + return info + + async def write_with_ack_future(self, message): + if self._last_known_seqno is None: + self._last_known_seqno = _FAKE_LAST_SEQNO.get(self.partition_id, 0) + if message.seqno <= self._last_known_seqno: + raise TopicWriterError("Message seqno is duplicated: %s" % message.seqno) + self._last_known_seqno = message.seqno + return await super().write_with_ack_future(message) + + +# Partitions whose sub-writer never finishes init; models a writer opened against a partition +# that is already inactive (observed live: such a writer hangs in init forever). +_FAKE_HANGING_INIT_PARTITIONS: set = set() + + +class _HangingInitSubWriter(_ControllableSubWriter): + async def wait_init(self): + if self.partition_id in _FAKE_HANGING_INIT_PARTITIONS: + await asyncio.Event().wait() + return await super().wait_init() + + +class _AckOnCloseSubWriter(_ControllableSubWriter): + """Sub-writer whose outstanding acks land while the stream is being torn down. + + Models the server persisting (and acking) a message in the same moment the split is + detected -- the ack races the quiesce that the repartition performs before reading its + dedup cut. + """ + + async def close(self, flush=True): + self.resolve_all() + self.closed = True + + +def _retrieve_exceptions(futures) -> None: + """Consume results so asyncio does not warn about never-retrieved exceptions.""" + for future in futures: + if future.done() and not future.cancelled(): + future.exception() + + +def _multi_partition(partition_id, parents=None, children=None, from_bound=None, to_bound=None, active=True): + key_range = None + if from_bound is not None or to_bound is not None: + key_range = PublicDescribeTopicResult.PartitionKeyRange(from_bound=from_bound or b"", to_bound=to_bound or b"") + return PublicDescribeTopicResult.PartitionInfo( + partition_id=partition_id, + active=active, + child_partition_ids=children or [], + parent_partition_ids=parents or [], + partition_stats=None, + key_range=key_range, + ) + + +# Real YDB split topology (observed live against a cloud cluster): the split parent stays in the +# DescribeTopic result as an INACTIVE partition whose child_partition_ids point at the new leaves, +# and each child is active with parent_partition_ids == [parent]. Mocks below mirror that so the +# tests exercise the orchestrator's active/child filtering on realistic input. +def _split_parent(partition_id, children, parents=None): + return _multi_partition(partition_id, parents=parents, children=children, active=False) + + +@pytest.mark.asyncio +class TestTopicWriterMultiAsyncIO: + async def test_routes_messages_by_key(self): + driver = _MultiFakeDescribeDriver([[_multi_partition(0), _multi_partition(1), _multi_partition(2)]]) + settings = MultiWriterSettings( + topic="/local/topic", + producer_id_prefix="pfx", + partition_chooser=PublicPartitionByKeyKafka(), + ) + with mock.patch("ydb._topic_writer.topic_writer_multi_asyncio.WriterAsyncIO", _FakeSubWriter): + writer = TopicWriterMultiAsyncIO(driver, settings) + await writer.wait_init() + + keys = ["a", "user-42", "hello", "мурмур2-хэш", "0", "zzz"] + for key in keys: + await writer.write(PublicMessage(b"payload", key=key)) + + for key in keys: + partition_id = (murmur2_32(key.encode("utf-8"), 0) & 0x7FFFFFFF) % 3 + sub = writer._writers[partition_id] + assert sub.producer_id == "pfx-%d" % partition_id + assert any(m.key == key for m in sub.messages) + + assert sum(len(w.messages) for w in writer._writers.values()) == len(keys) + await writer.close(flush=False) + + async def test_split_reroutes_to_child_partitions(self): + before = [_multi_partition(0), _multi_partition(1), _multi_partition(2)] + after = [ + _split_parent(0, children=[3, 4]), # split parent stays, inactive, with children + _multi_partition(1), + _multi_partition(2), + _multi_partition(3, parents=[0]), + _multi_partition(4, parents=[0]), + ] + driver = _MultiFakeDescribeDriver([before, after]) + settings = MultiWriterSettings( + topic="/local/topic", + producer_id_prefix="pfx", + partition_chooser=PublicPartitionByKeyKafka(), + ) + with mock.patch("ydb._topic_writer.topic_writer_multi_asyncio.WriterAsyncIO", _FakeSubWriter): + writer = TopicWriterMultiAsyncIO(driver, settings) + await writer.wait_init() + + sub0 = await writer._get_or_create_writer(0) + await writer._on_partition_overloaded(0) + + assert sub0.closed + assert 0 not in writer._writers + assert 0 not in writer._partitions + assert set(writer._partitions) == {1, 2, 3, 4} + assert set(writer._chooser._partitions) == {1, 2, 3, 4} + await writer.close(flush=False) + + async def test_split_resends_unacked_messages_with_dedup_cut(self): + _FAKE_LAST_SEQNO.clear() + # Route all three keys to partition 0 initially; after the split, spread them + # across the two children. + mapping = {"a": 0, "b": 0, "c": 0} + chooser = _KeyMapChooser(mapping) + before = [_multi_partition(0), _multi_partition(1)] + after = [ + _split_parent(0, children=[2, 3]), # split parent stays, inactive, with children + _multi_partition(1), + _multi_partition(2, parents=[0]), + _multi_partition(3, parents=[0]), + ] + driver = _MultiFakeDescribeDriver([before, after]) + settings = MultiWriterSettings(topic="/local/topic", producer_id_prefix="pfx", partition_chooser=chooser) + + with mock.patch("ydb._topic_writer.topic_writer_multi_asyncio.WriterAsyncIO", _ControllableSubWriter): + writer = TopicWriterMultiAsyncIO(driver, settings) + await writer.wait_init() + + f_a = await writer.write_with_ack_future(PublicMessage(b"a", key="a")) # partition 0, seqno 1 + f_b = await writer.write_with_ack_future(PublicMessage(b"b", key="b")) # partition 0, seqno 2 + f_c = await writer.write_with_ack_future(PublicMessage(b"c", key="c")) # partition 0, seqno 3 + assert set(writer._inflight[0]) == {1, 2, 3} + + # Ack the first message before the split: seqno 1 is now the dedup cut (max_acked), + # and an acked message is already out of the in-flight set (never resent). + writer._writers[0].pending[0].set_result(PublicWriteResult.Written(offset=0)) + await asyncio.sleep(0) + assert f_a.done() and writer._max_acked.get(0) == 1 + assert set(writer._inflight[0]) == {2, 3} + + # Split: the un-acked b, c (seqno > cut) are re-routed to the children. + mapping.update({"a": 2, "b": 3, "c": 2}) + await writer._on_partition_overloaded(0) + + assert 0 not in writer._inflight + assert [m.key for m in writer._writers[3].messages] == ["b"] + assert [m.key for m in writer._writers[2].messages] == ["c"] + assert not f_b.done() and not f_c.done() + + writer._writers[2].resolve_all() + writer._writers[3].resolve_all() + await asyncio.sleep(0) + assert f_b.done() and f_c.done() + + await writer.close(flush=False) + + async def test_adaptive_default_chooser(self): + # A topic without key ranges -> Kafka hash chooser. + driver_plain = _MultiFakeDescribeDriver([[_multi_partition(0), _multi_partition(1)]]) + # A topic that reports key ranges (auto-partitioning) -> bound chooser. + driver_auto = _MultiFakeDescribeDriver( + [[_multi_partition(0, from_bound=b"", to_bound=b"\x80"), _multi_partition(1, from_bound=b"\x80")]] + ) + with mock.patch("ydb._topic_writer.topic_writer_multi_asyncio.WriterAsyncIO", _FakeSubWriter): + plain = TopicWriterMultiAsyncIO(driver_plain, MultiWriterSettings(topic="/local/topic")) + await plain.wait_init() + assert isinstance(plain._chooser, PublicPartitionByKeyKafka) + await plain.close(flush=False) + + auto = TopicWriterMultiAsyncIO(driver_auto, MultiWriterSettings(topic="/local/topic")) + await auto.wait_init() + assert isinstance(auto._chooser, PublicPartitionByKeyBound) + await auto.close(flush=False) + + async def test_idle_writer_eviction(self): + driver = _MultiFakeDescribeDriver([[_multi_partition(0), _multi_partition(1)]]) + chooser = _KeyMapChooser({"a": 0, "b": 1}) + settings = MultiWriterSettings( + topic="/local/topic", producer_id_prefix="pfx", partition_chooser=chooser, writer_idle_timeout_sec=1000 + ) + with mock.patch("ydb._topic_writer.topic_writer_multi_asyncio.WriterAsyncIO", _ControllableSubWriter): + writer = TopicWriterMultiAsyncIO(driver, settings) + await writer.wait_init() + + f_a = await writer.write_with_ack_future(PublicMessage(b"a", key="a")) # partition 0 + f_b = await writer.write_with_ack_future(PublicMessage(b"b", key="b")) # partition 1 + assert set(writer._writers) == {0, 1} + sub0 = writer._writers[0] + + # ack partition 0 (it goes idle); leave partition 1 un-acked + sub0.resolve_all() + await asyncio.sleep(0) + assert not writer._inflight.get(0) and writer._inflight.get(1) + + # make both look old: only the idle partition 0 is evictable + old = writer._loop.time() - 5000 + writer._last_write_at[0] = old + writer._last_write_at[1] = old + await writer._evict_idle_writers() + + assert 0 not in writer._writers and sub0.closed # idle -> evicted + assert 1 in writer._writers # pending in-flight -> kept + assert writer._seqno == 2 # writer-wide cursor survives eviction (a -> 1, b -> 2) + + # writing to partition 0 again recreates a fresh sub-writer; numbering continues from + # the shared cursor rather than restarting for the re-opened partition + f_a2 = await writer.write_with_ack_future(PublicMessage(b"a2", key="a")) + assert 0 in writer._writers and writer._writers[0] is not sub0 + assert next(iter(writer._inflight[0])) == 3 + + writer._writers[0].resolve_all() + writer._writers[1].resolve_all() + await asyncio.sleep(0) + assert f_a.done() and f_b.done() and f_a2.done() + await writer.close(flush=False) + + async def test_split_hook_detects_overloaded_only(self): + # How the server signals a split (observed live): a split partition goes inactive, so the + # next write to it fails on the write stream with OVERLOADED (status_code 400060), + # message "Write to inactive partition N", surfaced by the SDK as issues.Overloaded. + # The hook triggers on that exception TYPE (the message text is not inspected); any other + # error is left to the writer's normal retry path. + driver = _MultiFakeDescribeDriver([[_multi_partition(0)]]) + settings = MultiWriterSettings(topic="/local/topic", partition_chooser=PublicPartitionByKeyKafka()) + with mock.patch("ydb._topic_writer.topic_writer_multi_asyncio.WriterAsyncIO", _FakeSubWriter): + writer = TopicWriterMultiAsyncIO(driver, settings) + await writer.wait_init() + + observed: List[int] = [] + + async def fake_split(partition_id): + observed.append(partition_id) + + writer._on_partition_overloaded = fake_split + hook = writer._make_overloaded_hook(0) + + split_signal = issues.Overloaded("status is not ok: Write to inactive partition 0") + assert hook(split_signal) is True + await asyncio.sleep(0) + assert observed == [0] + assert hook(RuntimeError("some other error")) is False + + await writer.close(flush=False) + + async def test_repartition_tolerates_subwriter_close_raising(self): + # Regression: a hook-stopped sub-writer's close() re-raises TopicWriterPartitionSplitError. + # Repartition must swallow it and still migrate to the child (not fall back to recovery). + _FAKE_LAST_SEQNO.clear() + mapping = {"a": 0} + chooser = _KeyMapChooser(mapping) + before = [_multi_partition(0), _multi_partition(1)] + after = [ + _split_parent(0, children=[2, 3]), # split parent stays, inactive, with children + _multi_partition(1), + _multi_partition(2, parents=[0]), + _multi_partition(3, parents=[0]), + ] + driver = _MultiFakeDescribeDriver([before, after]) + settings = MultiWriterSettings(topic="/local/topic", producer_id_prefix="pfx", partition_chooser=chooser) + + with mock.patch("ydb._topic_writer.topic_writer_multi_asyncio.WriterAsyncIO", _CloseRaisesSubWriter): + writer = TopicWriterMultiAsyncIO(driver, settings) + await writer.wait_init() + + f_a = await writer.write_with_ack_future(PublicMessage(b"a", key="a")) # partition 0 + mapping["a"] = 2 # after split, route to child 2 + + await writer._on_partition_overloaded(0) # must not raise despite close() raising + + assert 0 not in writer._partitions # retired cleanly (no recovery fallback) + assert [m.key for m in writer._writers[2].messages] == ["a"] # migrated to the child + assert 0 not in writer._inflight + + writer._writers[2].resolve_all() + await asyncio.sleep(0) + assert f_a.done() + + await writer.close(flush=False) + + async def test_merge_migrates_both_parents_to_shared_child(self): + _FAKE_LAST_SEQNO.clear() + # Two parents (0, 1) merge into one child (2), which lists both as parents. The bounds are + # real ones: a merge child owns the ranges of BOTH parents, so it covers strictly more of + # the key space than the parent whose OVERLOADED triggered the handler -- the child-range + # coverage check must accept that, not just an exact tiling. + mapping = {"x": 0, "y": 1} + chooser = _KeyMapChooser(mapping) + before = [ + _multi_partition(0, from_bound=b"", to_bound=b"\x80"), + _multi_partition(1, from_bound=b"\x80", to_bound=b""), + ] + after = [_multi_partition(2, parents=[0, 1], from_bound=b"", to_bound=b"")] + driver = _MultiFakeDescribeDriver([before, after]) + settings = MultiWriterSettings(topic="/local/topic", producer_id_prefix="pfx", partition_chooser=chooser) + + with mock.patch("ydb._topic_writer.topic_writer_multi_asyncio.WriterAsyncIO", _ControllableSubWriter): + writer = TopicWriterMultiAsyncIO(driver, settings) + await writer.wait_init() + + f_x = await writer.write_with_ack_future(PublicMessage(b"x", key="x")) # partition 0 + f_y = await writer.write_with_ack_future(PublicMessage(b"y", key="y")) # partition 1 + # One sequence for the whole writer, so the two partitions do not reuse a number and + # the merge can carry both into the shared child unchanged. + assert set(writer._inflight[0]) == {1} and set(writer._inflight[1]) == {2} + + # After the merge both keys route to the shared child 2. + mapping.update({"x": 2, "y": 2}) + # Overloaded fired for partition 0 only; the handler must retire partition 1 too. + await writer._on_partition_overloaded(0) + + assert set(writer._partitions) == {2} + assert writer._chooser.partitions == {2} + assert 0 not in writer._writers and 1 not in writer._writers + assert sorted(m.key for m in writer._writers[2].messages) == ["x", "y"] + assert 0 not in writer._inflight and 1 not in writer._inflight + + writer._writers[2].resolve_all() + await asyncio.sleep(0) + assert f_x.done() and f_y.done() + + await writer.close(flush=False) + + async def test_close_flushes_buffered_messages(self): + driver = _MultiFakeDescribeDriver([[_multi_partition(0), _multi_partition(1), _multi_partition(2)]]) + settings = MultiWriterSettings( + topic="/local/topic", producer_id_prefix="pfx", partition_chooser=PublicPartitionByKeyKafka() + ) + with mock.patch("ydb._topic_writer.topic_writer_multi_asyncio.WriterAsyncIO", _FlushControlledSubWriter): + writer = TopicWriterMultiAsyncIO(driver, settings) + futures = [ + await writer.write_with_ack_future(PublicMessage(("m%d" % i).encode(), key="k%d" % i)) for i in range(5) + ] + assert not any(f.done() for f in futures) # nothing acked yet + + await writer.close() # flush=True must deliver the buffered messages + + assert all(f.done() and not f.cancelled() and f.exception() is None for f in futures) + + async def test_adaptive_chooser_single_open_range_partition(self): + # A single auto-partitioned partition owns the fully-open range b""..b"". + driver = _MultiFakeDescribeDriver([[_multi_partition(0, from_bound=b"", to_bound=b"")]]) + with mock.patch("ydb._topic_writer.topic_writer_multi_asyncio.WriterAsyncIO", _FakeSubWriter): + writer = TopicWriterMultiAsyncIO(driver, MultiWriterSettings(topic="/local/topic")) + await writer.wait_init() + assert isinstance(writer._chooser, PublicPartitionByKeyBound) + await writer.close(flush=False) + + async def test_transient_overload_recovers_partition_in_place(self): + # DescribeTopic never shows children -> ordinary overload, not a repartition. + driver = _MultiFakeDescribeDriver([[_multi_partition(0), _multi_partition(1)]]) + chooser = _KeyMapChooser({"a": 0}) + settings = MultiWriterSettings(topic="/local/topic", producer_id_prefix="pfx", partition_chooser=chooser) + with mock.patch( + "ydb._topic_writer.topic_writer_multi_asyncio.WriterAsyncIO", _ControllableSubWriter + ), mock.patch("ydb._topic_writer.topic_writer_multi_asyncio._REPARTITION_DISCOVER_DELAY", 0), mock.patch( + "ydb._topic_writer.topic_writer_multi_asyncio._REPARTITION_DISCOVER_ATTEMPTS", 2 + ): + writer = TopicWriterMultiAsyncIO(driver, settings) + await writer.wait_init() + + future = await writer.write_with_ack_future(PublicMessage(b"a", key="a")) + old_sub = writer._writers[0] + + await writer._on_partition_overloaded(0) + + # partition kept (not retired); a fresh sub-writer resends the message + assert 0 in writer._partitions + new_sub = writer._writers[0] + assert new_sub is not old_sub + assert old_sub.closed + assert [m.key for m in new_sub.messages] == ["a"] + assert not future.done() + + new_sub.resolve_all() + await asyncio.sleep(0) + assert future.done() + + await writer.close(flush=False) + + async def test_enqueue_failure_does_not_leak_inflight(self): + driver = _MultiFakeDescribeDriver([[_multi_partition(0), _multi_partition(1)]]) + chooser = _KeyMapChooser({"a": 0}) + settings = MultiWriterSettings(topic="/local/topic", producer_id_prefix="pfx", partition_chooser=chooser) + with mock.patch("ydb._topic_writer.topic_writer_multi_asyncio.WriterAsyncIO", _RaisingSubWriter): + writer = TopicWriterMultiAsyncIO(driver, settings) + await writer.wait_init() + + with pytest.raises(RuntimeError): + await writer.write_with_ack_future(PublicMessage(b"a", key="a")) + + assert not writer._inflight.get(0) # no leaked entry, no pending future + await writer.close(flush=False) + + async def test_duplicate_seqno_rejected_without_leak(self): + driver = _MultiFakeDescribeDriver([[_multi_partition(0), _multi_partition(1)]]) + chooser = _KeyMapChooser({"a": 0, "b": 0}) + settings = MultiWriterSettings( + topic="/local/topic", producer_id_prefix="pfx", partition_chooser=chooser, auto_seqno=False + ) + with mock.patch("ydb._topic_writer.topic_writer_multi_asyncio.WriterAsyncIO", _ControllableSubWriter): + writer = TopicWriterMultiAsyncIO(driver, settings) + await writer.wait_init() + + first = await writer.write_with_ack_future(PublicMessage(b"a", key="a", seqno=5)) + with pytest.raises(TopicWriterError): + await writer.write_with_ack_future(PublicMessage(b"b", key="b", seqno=5)) + + assert set(writer._inflight[0]) == {5} + await writer.close(flush=False) + assert isinstance(first.exception(), TopicWriterStopped) # retrieve to avoid warning + + async def test_unusable_partition_fails_its_inflight_instead_of_stranding_it(self): + """When neither repartition nor recovery can serve a partition, its messages must fail. + + The children never complete the parent's range, so the topology is not committed; the + parent is already inactive, so re-opening a writer for it times out too. Nothing is left + that could ever ack these messages, and flush()/close(flush=True) wait on user futures + without a deadline -- so leaving them pending hangs the caller forever. + """ + _FAKE_LAST_SEQNO.clear() + _FAKE_HANGING_INIT_PARTITIONS.clear() + chooser = _KeyMapChooser({"a": 0}) + before = [_multi_partition(0, from_bound=b"", to_bound=b"")] + # Only one child of the split ever shows up -> coverage check refuses to retire parent 0. + partial = [ + _split_parent(0, children=[1, 2]), + _multi_partition(1, parents=[0], from_bound=b"", to_bound=b"\x80"), + ] + driver = _MultiFakeDescribeDriver([before, partial]) + settings = MultiWriterSettings(topic="/local/topic", producer_id_prefix="pfx", partition_chooser=chooser) + + with mock.patch( + "ydb._topic_writer.topic_writer_multi_asyncio.WriterAsyncIO", _HangingInitSubWriter + ), mock.patch("ydb._topic_writer.topic_writer_multi_asyncio._WRITER_INIT_TIMEOUT", 0.1), mock.patch( + "ydb._topic_writer.topic_writer_multi_asyncio._REPARTITION_DISCOVER_DELAY", 0 + ), mock.patch( + "ydb._topic_writer.topic_writer_multi_asyncio._REPARTITION_DISCOVER_ATTEMPTS", 2 + ): + writer = TopicWriterMultiAsyncIO(driver, settings) + await writer.wait_init() + future = await writer.write_with_ack_future(PublicMessage(b"a", key="a")) + + # Parent 0 has gone inactive: recovering it in place cannot finish init either. + _FAKE_HANGING_INIT_PARTITIONS.add(0) + try: + await writer._on_partition_overloaded(0) + + assert future.done(), "in-flight message left without an owner" + assert future.exception() is not None + assert not writer._inflight.get(0) + + # The caller must be able to finish; both would hang on a pending future. + await asyncio.wait_for(writer.flush(), timeout=1) + await asyncio.wait_for(writer.close(flush=True), timeout=1) + finally: + _FAKE_HANGING_INIT_PARTITIONS.clear() + await writer.close(flush=False) + _retrieve_exceptions([future]) + + async def test_repartition_tasks_are_coalesced_and_closed_with_the_writer(self): + """Repartition must be owned by the writer, not fire-and-forget. + + A burst of OVERLOADED on one partition otherwise starts several concurrent recoveries of + it, and any of them can outlive close() -- still describing the topic and opening + sub-writers for a multi-writer the caller believes is shut down. + """ + driver = _MultiFakeDescribeDriver([[_multi_partition(0), _multi_partition(1)]]) + settings = MultiWriterSettings( + topic="/local/topic", producer_id_prefix="pfx", partition_chooser=PublicPartitionByKeyKafka() + ) + with mock.patch("ydb._topic_writer.topic_writer_multi_asyncio.WriterAsyncIO", _FakeSubWriter): + writer = TopicWriterMultiAsyncIO(driver, settings) + await writer.wait_init() + + started: List[int] = [] + + async def slow_repartition(partition_id): + started.append(partition_id) + await asyncio.sleep(30) + + writer._on_partition_overloaded = slow_repartition + hook = writer._make_overloaded_hook(0) + + for _ in range(5): # a burst for the same partition + assert hook(issues.Overloaded("Write to inactive partition 0")) is True + await asyncio.sleep(0) + + assert started == [0], "repeated signals for one partition must coalesce" + task = writer._repartition_tasks[0] + + await writer.close(flush=False) + assert task.done(), "close() must cancel and await the repartition task" + assert not writer._repartition_tasks + + # After close no further signal may start work. + assert hook(issues.Overloaded("again")) is True + await asyncio.sleep(0) + assert started == [0] + + async def test_merge_does_not_overwrite_a_colliding_manual_seqno(self): + """Manual seqnos are unique per partition, so a merge can collide them in the child. + + Writing the migrated entry over the existing one would silently detach the displaced + message: its ack callback becomes stale and its user future never resolves. + """ + _FAKE_LAST_SEQNO.clear() + mapping = {"x": 0, "y": 1} + chooser = _KeyMapChooser(mapping) + before = [_multi_partition(0), _multi_partition(1)] + after = [_multi_partition(2, parents=[0, 1])] + driver = _MultiFakeDescribeDriver([before, after]) + settings = MultiWriterSettings( + topic="/local/topic", producer_id_prefix="pfx", partition_chooser=chooser, auto_seqno=False + ) + + with mock.patch("ydb._topic_writer.topic_writer_multi_asyncio.WriterAsyncIO", _ControllableSubWriter): + writer = TopicWriterMultiAsyncIO(driver, settings) + await writer.wait_init() + + # Same seqno on two different partitions: allowed today, both are in flight. + f_x = await writer.write_with_ack_future(PublicMessage(b"x", key="x", seqno=7)) + f_y = await writer.write_with_ack_future(PublicMessage(b"y", key="y", seqno=7)) + assert set(writer._inflight[0]) == {7} and set(writer._inflight[1]) == {7} + + mapping.update({"x": 2, "y": 2}) + await writer._on_partition_overloaded(0) + + # Whatever the resolution, neither message may be silently dropped. + assert len(writer._inflight.get(2, {})) + sum(f.done() for f in (f_x, f_y)) == 2 + + writer._writers[2].resolve_all() + await asyncio.sleep(0) + assert f_x.done() and f_y.done(), "a colliding migration stranded a user future" + + await writer.close(flush=False) + _retrieve_exceptions([f_x, f_y]) + + async def test_message_persisted_with_a_lost_ack_is_not_resent_to_the_child(self): + """The dedup cut has to come from the server, not from the acks we happened to receive. + + A split kills the session, and a message the server already persisted can lose its ack on + the way back. Judging only by acks we saw, such a message looks unwritten, so it gets + resent to the child -- where the parent's producer id no longer covers it, because each + partition writes under its own. Nothing on the server can collapse the two, so the reader + sees the message twice. Reading the cut from the retiring producer instead closes that + window: the message is below it, and is reported written rather than resent. + """ + _FAKE_LAST_SEQNO.clear() + mapping = {"a": 0} + chooser = _KeyMapChooser(mapping) + before = [_multi_partition(0), _multi_partition(1)] + after = [ + _split_parent(0, children=[2, 3]), + _multi_partition(1), + _multi_partition(2, parents=[0]), + _multi_partition(3, parents=[0]), + ] + driver = _MultiFakeDescribeDriver([before, after]) + settings = MultiWriterSettings(topic="/local/topic", producer_id_prefix="pfx", partition_chooser=chooser) + + with mock.patch("ydb._topic_writer.topic_writer_multi_asyncio.WriterAsyncIO", _ControllableSubWriter): + writer = TopicWriterMultiAsyncIO(driver, settings) + await writer.wait_init() + future = await writer.write_with_ack_future(PublicMessage(b"a", key="a")) # seqno 1 + assert set(writer._inflight[0]) == {1} + assert writer._max_acked.get(0, 0) == 0 # no ack ever reached us + + # The server did persist it: partition 0's producer is at seqno 1. + _FAKE_LAST_SEQNO[0] = 1 + + mapping["a"] = 2 + await writer._on_partition_overloaded(0) + + resent = [m.seqno for m in writer._writers[2].messages] if 2 in writer._writers else [] + assert resent == [], "a message the server had already persisted was resent to the child" + assert future.done() and future.exception() is None, "the persisted message must resolve as written" + assert not writer._inflight.get(0) + + await writer.close(flush=False) + + async def test_cut_ignores_ancestor_history(self): + """The cut is the retiring partition's own high-water mark, nothing else. + + Two different things live in an ancestor's producer, and neither belongs in this number. + A merge child has two parents whose branches numbered independently, so a sibling's + history says nothing about messages that came down this branch. And producer ids are + `prefix-`, so with a stable prefix a retired ancestor still holds whatever a + *previous run* wrote there. + + Either one, folded into the maximum, marks unsent messages as already written and drops + them silently. Nothing needs to be read from them: a message only reaches this partition + by having a number above the cut that retired the previous one. + """ + _FAKE_LAST_SEQNO.clear() + _FAKE_LAST_SEQNO[1] = 5 # this branch numbered modestly + _FAKE_LAST_SEQNO[2] = 100 # the sibling branch ran far ahead + _FAKE_LAST_SEQNO[3] = 7 # the surviving partition itself + + # The topology comes from describe, as it does in production: p3 is the merge of p1 and + # p2, so both are reachable as its parents. + described = [ + _multi_partition(1, children=[3]), + _multi_partition(2, children=[3]), + _multi_partition(3, parents=[1, 2]), + ] + driver = _MultiFakeDescribeDriver([described]) + settings = MultiWriterSettings( + topic="/local/topic", producer_id_prefix="pfx", partition_chooser=_KeyMapChooser({}) + ) + with mock.patch("ydb._topic_writer.topic_writer_multi_asyncio.WriterAsyncIO", _ControllableSubWriter): + writer = TopicWriterMultiAsyncIO(driver, settings) + await writer.wait_init() + # This writer wrote to both branches before they merged. + for partition_id in (1, 2, 3): + await writer._get_or_create_writer(partition_id) + + assert await writer._max_seqno_cut(3) == 7, "an ancestor's history leaked into the cut" + + await writer.close(flush=False) + + async def test_a_message_survives_repeated_retirements(self): + """A message must not be written off by a partition it already escaped. + + This is the induction the cut relies on, end to end: seqno 6 outruns the parent's cut of + 5 and moves to the child, and when the child is retired in turn it has to move again. If + an earlier partition's number could still reach it, it would be dropped here. + """ + _FAKE_LAST_SEQNO.clear() + _FAKE_LAST_SEQNO[0] = 5 # parent persisted up to 5; our message will be 6 + mapping = {"a": 0} + chooser = _KeyMapChooser(mapping) + before = [_multi_partition(0), _multi_partition(1)] + after_first = [ + _split_parent(0, children=[2, 3]), + _multi_partition(1), + _multi_partition(2, parents=[0]), + _multi_partition(3, parents=[0]), + ] + after_second = [ + _split_parent(0, children=[2, 3]), + _split_parent(2, children=[4, 5], parents=[0]), + _multi_partition(1), + _multi_partition(3, parents=[0]), + _multi_partition(4, parents=[2]), + _multi_partition(5, parents=[2]), + ] + driver = _MultiFakeDescribeDriver([before, after_first, after_second]) + settings = MultiWriterSettings(topic="/local/topic", producer_id_prefix="pfx", partition_chooser=chooser) + + with mock.patch("ydb._topic_writer.topic_writer_multi_asyncio.WriterAsyncIO", _ControllableSubWriter): + writer = TopicWriterMultiAsyncIO(driver, settings) + await writer.wait_init() + future = await writer.write_with_ack_future(PublicMessage(b"a", key="a")) + seqno = next(iter(writer._inflight[0])) + assert seqno == 6, "the cursor starts above the parent's persisted history" + + mapping["a"] = 2 + await writer._on_partition_overloaded(0) + assert [m.seqno for m in writer._writers[2].messages] == [6] + + # The child retires too, and the message has to keep going. + mapping["a"] = 4 + await writer._on_partition_overloaded(2) + assert [m.seqno for m in writer._writers[4].messages] == [6] + assert not future.done() + + writer._writers[4].resolve_all() + await asyncio.sleep(0) + assert future.done() and future.exception() is None + await writer.close(flush=False) + + async def test_missing_manual_seqno_is_a_validation_error(self): + driver = _MultiFakeDescribeDriver([[_multi_partition(0), _multi_partition(1)]]) + settings = MultiWriterSettings( + topic="/local/topic", + producer_id_prefix="pfx", + partition_chooser=_KeyMapChooser({"a": 0}), + auto_seqno=False, + ) + with mock.patch("ydb._topic_writer.topic_writer_multi_asyncio.WriterAsyncIO", _ControllableSubWriter): + writer = TopicWriterMultiAsyncIO(driver, settings) + await writer.wait_init() + + # The writer is healthy; the message is what is wrong. Reporting this as + # TopicWriterStopped tells the caller to give up on a writer that is still usable. + with pytest.raises(TopicWriterError) as err: + await writer.write_with_ack_future(PublicMessage(b"a", key="a")) + assert not isinstance(err.value, TopicWriterStopped) + + await writer.close(flush=False) + + async def test_recovery_after_lost_ack_resends_the_remaining_messages(self): + """A transient overload plus one ack lost with the stream must not strand the partition. + + The server persisted seqno 1 but its ack never reached us, so the message is still + in-flight. _recover_partition resends the whole in-flight set to a fresh sub-writer for + the same partition -- whose init reports last_seqno=1, so the real writer rejects seqno 1 + with "Message seqno is duplicated" before the server ever sees it. The resend loop is a + plain `for`, so that exception also skips seqnos 2 and 3, which are never retried and + never resolve. + + The comment on _recover_partition assumes the server dedups this retry; the client-side + guard fires first, so it never gets the chance. + """ + _FAKE_LAST_SEQNO.clear() + driver = _MultiFakeDescribeDriver([[_multi_partition(0), _multi_partition(1)]]) + chooser = _KeyMapChooser({"a": 0}) + settings = MultiWriterSettings(topic="/local/topic", producer_id_prefix="pfx", partition_chooser=chooser) + with mock.patch("ydb._topic_writer.topic_writer_multi_asyncio.WriterAsyncIO", _SeqnoGuardSubWriter), mock.patch( + "ydb._topic_writer.topic_writer_multi_asyncio._REPARTITION_DISCOVER_DELAY", 0 + ), mock.patch("ydb._topic_writer.topic_writer_multi_asyncio._REPARTITION_DISCOVER_ATTEMPTS", 2): + writer = TopicWriterMultiAsyncIO(driver, settings) + await writer.wait_init() + futures = [ + await writer.write_with_ack_future(PublicMessage(("m%d" % i).encode(), key="a")) for i in range(3) + ] + assert set(writer._inflight[0]) == {1, 2, 3} + + # The server persisted seqno 1 and acked it, but the ack died with the stream. + _FAKE_LAST_SEQNO[0] = 1 + + try: + await writer._on_partition_overloaded(0) + + new_sub = writer._writers[0] + assert [m.seqno for m in new_sub.messages] == [2, 3], "messages after the persisted one were dropped" + assert futures[0].done() and futures[0].exception() is None, "persisted message must resolve as written" + + new_sub.resolve_all() + await asyncio.sleep(0) + assert all(f.done() for f in futures) + finally: + await writer.close(flush=False) + _retrieve_exceptions(futures) + + async def test_split_migration_does_not_block_on_an_uninitializable_child(self): + """Migration must not wedge the whole multi-writer when a child cannot be opened. + + Splits cascade (1->3->7 was observed live), so by the time we migrate, the child the + chooser picked may itself have split and gone inactive. _get_or_create_writer awaits + wait_init() on it while holding the orchestrator lock, and a writer against an inactive + partition never finishes init -- so every write, flush and further repartition blocks + behind it. This is the same failure the maxSeqNo probe caused before it was removed; + the unbounded wait_init under the lock survived it. + """ + _FAKE_LAST_SEQNO.clear() + _FAKE_HANGING_INIT_PARTITIONS.clear() + mapping = {"a": 0, "b": 1} + chooser = _KeyMapChooser(mapping) + before = [_multi_partition(0), _multi_partition(1)] + after = [ + _split_parent(0, children=[2, 3]), + _multi_partition(1), + _multi_partition(2, parents=[0]), + _multi_partition(3, parents=[0]), + ] + driver = _MultiFakeDescribeDriver([before, after]) + settings = MultiWriterSettings(topic="/local/topic", producer_id_prefix="pfx", partition_chooser=chooser) + + with mock.patch( + "ydb._topic_writer.topic_writer_multi_asyncio.WriterAsyncIO", _HangingInitSubWriter + ), mock.patch("ydb._topic_writer.topic_writer_multi_asyncio._WRITER_INIT_TIMEOUT", 0.1): + writer = TopicWriterMultiAsyncIO(driver, settings) + await writer.wait_init() + future = await writer.write_with_ack_future(PublicMessage(b"a", key="a")) + + mapping["a"] = 2 + _FAKE_HANGING_INIT_PARTITIONS.add(2) # child 2 split again before we got to it + + try: + try: + await asyncio.wait_for(writer._on_partition_overloaded(0), timeout=1) + except asyncio.TimeoutError: + pytest.fail("repartition blocked on an uninitializable child while holding the lock") + + # And the lock must be free afterwards: an unrelated partition still accepts writes. + await asyncio.wait_for(writer.write_with_ack_future(PublicMessage(b"b", key="b")), timeout=1) + finally: + _FAKE_HANGING_INIT_PARTITIONS.clear() + await writer.close(flush=False) + _retrieve_exceptions([future]) + + async def test_split_waits_for_children_to_cover_the_parent_range(self): + """A partially visible split must not retire the parent. + + DescribeTopic can show one child before its sibling becomes active. Retiring the parent + on that view drops its key range down to the single child, leaving the rest of the key + space uncovered -- and since routing only compares from_bound, keys from the missing + range then land in the left sibling. That puts one key on two branches of the partition + graph, which is the invariant the whole design exists to protect. + + Per the C++ producer spec an incomplete graph is a retry-with-backoff state, not a + successful split: "producer не должен выбирать случайную partition". + """ + before = [_multi_partition(0, from_bound=b"", to_bound=b"")] + partial = [ + _split_parent(0, children=[1, 2]), + _multi_partition(1, parents=[0], from_bound=b"", to_bound=b"\x80"), + ] + complete = partial + [_multi_partition(2, parents=[0], from_bound=b"\x80", to_bound=b"")] + driver = _MultiFakeDescribeDriver([before, partial, complete]) + settings = MultiWriterSettings(topic="/local/topic", producer_id_prefix="pfx") + + with mock.patch("ydb._topic_writer.topic_writer_multi_asyncio.WriterAsyncIO", _FakeSubWriter), mock.patch( + "ydb._topic_writer.topic_writer_multi_asyncio._REPARTITION_DISCOVER_DELAY", 0 + ): + writer = TopicWriterMultiAsyncIO(driver, settings) + await writer.wait_init() + assert isinstance(writer._chooser, PublicPartitionByKeyBound) + + await writer._on_partition_overloaded(0) + + assert set(writer._partitions) == {1, 2}, "parent retired before both children were visible" + assert sorted(p[-1] for p in writer._chooser._partitions) == [1, 2] + await writer.close(flush=False) + + async def test_ack_racing_the_split_is_not_resent_to_the_child(self): + """An ack that lands during the quiesce must count towards the dedup cut. + + Closing a retired sub-writer fails the acks it still holds, and those failures must not + reach the user because the messages are about to be migrated. Suppressing them must not + also swallow a SUCCESS landing in the same window: that ack is real, and dropping it + leaves _max_acked -- the dedup cut -- too low, so a message the parent already persisted + is resent to the child. That is the duplicate the cut exists to prevent. + """ + _FAKE_LAST_SEQNO.clear() + mapping = {"a": 0} + chooser = _KeyMapChooser(mapping) + before = [_multi_partition(0), _multi_partition(1)] + after = [ + _split_parent(0, children=[2, 3]), + _multi_partition(1), + _multi_partition(2, parents=[0]), + _multi_partition(3, parents=[0]), + ] + driver = _MultiFakeDescribeDriver([before, after]) + settings = MultiWriterSettings(topic="/local/topic", producer_id_prefix="pfx", partition_chooser=chooser) + + with mock.patch("ydb._topic_writer.topic_writer_multi_asyncio.WriterAsyncIO", _AckOnCloseSubWriter): + writer = TopicWriterMultiAsyncIO(driver, settings) + await writer.wait_init() + future = await writer.write_with_ack_future(PublicMessage(b"a", key="a")) + + mapping["a"] = 2 + try: + await writer._on_partition_overloaded(0) + await asyncio.sleep(0) + + resent = [m.key for m in writer._writers[2].messages] if 2 in writer._writers else [] + assert resent == [], "a message the parent persisted was resent to the child" + assert future.done() and future.exception() is None + finally: + await writer.close(flush=False) + _retrieve_exceptions([future]) + + async def test_split_resend_preserves_the_original_seqno(self): + """Both reference implementations keep a message's seqno when resending it to a child. + + C++ `TProducer::TMessagesWorker::ScheduleResendMessages` reassigns only the target + partition and leaves `SeqNo` alone; Go's multiwriter does the same. That works because + their counter is global: `CurrentSeqNo` is a single cursor per producer (C++ + `producer.h`), as is Go's `o.currentSeqNo`. A number drawn from one global sequence stays + meaningful in whatever partition the message ends up in. + + Ours is per partition, so a migrated message would carry a number from the parent's + sequence into a child that has its own -- hence the renumbering this test pins down. + Adopting the reference model means replacing the per-partition cursors with one global + counter first; preserving the seqno without that would break monotonicity in the child. + + Note this is not what makes dedup work: producer_id is per partition in C++ and Go too + (`"{prefix}_{partitionId}"`), so the server cannot deduplicate across a split either way. + Both implementations rely on a client-side maxSeqNo cut, exactly as we do. + """ + _FAKE_LAST_SEQNO.clear() + mapping = {"a": 0, "b": 0} + chooser = _KeyMapChooser(mapping) + before = [_multi_partition(0), _multi_partition(1)] + after = [ + _split_parent(0, children=[2, 3]), + _multi_partition(1), + _multi_partition(2, parents=[0]), + _multi_partition(3, parents=[0]), + ] + driver = _MultiFakeDescribeDriver([before, after]) + settings = MultiWriterSettings(topic="/local/topic", producer_id_prefix="pfx", partition_chooser=chooser) + + with mock.patch("ydb._topic_writer.topic_writer_multi_asyncio.WriterAsyncIO", _ControllableSubWriter): + writer = TopicWriterMultiAsyncIO(driver, settings) + await writer.wait_init() + f_a = await writer.write_with_ack_future(PublicMessage(b"a", key="a")) # partition 0, seqno 1 + f_b = await writer.write_with_ack_future(PublicMessage(b"b", key="b")) # partition 0, seqno 2 + assert set(writer._inflight[0]) == {1, 2} + + # The split sends the two keys to different children; each keeps its own number. + mapping.update({"a": 2, "b": 3}) + try: + await writer._on_partition_overloaded(0) + + assert [m.seqno for m in writer._writers[2].messages] == [1] + assert [m.seqno for m in writer._writers[3].messages] == [2] + assert set(writer._inflight[2]) == {1} and set(writer._inflight[3]) == {2} + finally: + await writer.close(flush=False) + _retrieve_exceptions([f_a, f_b]) + + +@pytest.mark.asyncio +class TestTopicWriterMultiAsyncIOLifecycle: + """Lifecycle and error-path behaviour of the orchestrator. + + Separate from the routing/split tests: these are about what happens around the happy path -- + shutdown, destructors, background tasks and the branches that only run when something fails. + """ + + def _driver(self, partitions=None): + return _MultiFakeDescribeDriver([partitions or [_multi_partition(0), _multi_partition(1)]]) + + def _settings(self, **kwargs): + kwargs.setdefault("topic", "/local/topic") + kwargs.setdefault("producer_id_prefix", "pfx") + kwargs.setdefault("partition_chooser", _KeyMapChooser({"a": 0, "b": 1})) + return MultiWriterSettings(**kwargs) + + async def test_context_manager_closes_and_keeps_body_errors(self): + with mock.patch("ydb._topic_writer.topic_writer_multi_asyncio.WriterAsyncIO", _FakeSubWriter): + async with TopicWriterMultiAsyncIO(self._driver(), self._settings()) as writer: + await writer.write(PublicMessage(b"a", key="a")) + assert writer._closed + + class TestException(Exception): + pass + + # A failure inside the block must survive close(): losing it would report a real + # error as an unrelated teardown problem. + with pytest.raises(TestException): + async with TopicWriterMultiAsyncIO(self._driver(), self._settings()) as writer: + raise TestException() + assert writer._closed + + async def test_unclosed_writer_schedules_a_close_on_delete(self): + with mock.patch("ydb._topic_writer.topic_writer_multi_asyncio.WriterAsyncIO", _FakeSubWriter): + writer = TopicWriterMultiAsyncIO(self._driver(), self._settings()) + await writer.wait_init() + + writer.__del__() # forgotten writer: the streams still have to be released + await asyncio.sleep(0) + await asyncio.sleep(0) + assert writer._closed + + writer.__del__() # already closed -> nothing scheduled, still no raise + + async def test_writes_are_refused_after_close(self): + with mock.patch("ydb._topic_writer.topic_writer_multi_asyncio.WriterAsyncIO", _FakeSubWriter): + writer = TopicWriterMultiAsyncIO(self._driver(), self._settings()) + await writer.wait_init() + await writer.close(flush=False) + + with pytest.raises(TopicWriterClosedError): + await writer.write(PublicMessage(b"a", key="a")) + with pytest.raises(TopicWriterClosedError): + await writer.flush() + + async def test_close_cancels_an_init_that_never_finished(self): + class NeverDescribes(_MultiFakeDescribeDriver): + async def __call__(self, *args, **kwargs): + await asyncio.Event().wait() + + with mock.patch("ydb._topic_writer.topic_writer_multi_asyncio.WriterAsyncIO", _FakeSubWriter): + writer = TopicWriterMultiAsyncIO(NeverDescribes([[]]), self._settings()) + await asyncio.wait_for(writer.close(), timeout=1) + + # An init left running would keep describing a topic nobody writes to any more. + with pytest.raises(asyncio.CancelledError): + await writer._init_task + + async def test_close_tolerates_a_failing_flush(self): + class FlushRaises(_FakeSubWriter): + async def flush(self): + raise RuntimeError("flush failed") + + with mock.patch("ydb._topic_writer.topic_writer_multi_asyncio.WriterAsyncIO", FlushRaises): + writer = TopicWriterMultiAsyncIO(self._driver(), self._settings()) + await writer.write(PublicMessage(b"a", key="a")) + # close(flush=True) must still shut the writer down, not propagate the flush error. + await writer.close() + assert writer._closed + + async def test_write_with_ack_returns_results_for_one_and_many(self): + with mock.patch("ydb._topic_writer.topic_writer_multi_asyncio.WriterAsyncIO", _FakeSubWriter): + writer = TopicWriterMultiAsyncIO(self._driver(), self._settings()) + + single = await writer.write_with_ack(PublicMessage(b"a", key="a")) + assert isinstance(single, PublicWriteResult.Written) + + many = await writer.write_with_ack([PublicMessage(b"a", key="a"), PublicMessage(b"b", key="b")]) + assert isinstance(many, list) and len(many) == 2 + + await writer.close(flush=False) + + async def test_a_failed_ack_reaches_the_caller(self): + with mock.patch("ydb._topic_writer.topic_writer_multi_asyncio.WriterAsyncIO", _ControllableSubWriter): + writer = TopicWriterMultiAsyncIO(self._driver(), self._settings()) + future = await writer.write_with_ack_future(PublicMessage(b"a", key="a")) + + writer._writers[0].pending[0].set_exception(RuntimeError("write rejected")) + await asyncio.sleep(0) + + assert isinstance(future.exception(), RuntimeError) + assert not writer._inflight.get(0), "a settled message must not stay in flight" + await writer.close(flush=False) + + async def test_a_cancelled_ack_leaves_the_message_in_flight(self): + """Cancellation is not an outcome: the writer is being torn down, and the message is + still owned by the orchestrator until a repartition or recovery decides its fate.""" + with mock.patch("ydb._topic_writer.topic_writer_multi_asyncio.WriterAsyncIO", _ControllableSubWriter): + writer = TopicWriterMultiAsyncIO(self._driver(), self._settings()) + future = await writer.write_with_ack_future(PublicMessage(b"a", key="a")) + + writer._writers[0].pending[0].cancel() + await asyncio.sleep(0) + + assert not future.done() + assert set(writer._inflight[0]) == {1} + await writer.close(flush=False) + _retrieve_exceptions([future]) + + async def test_ack_failures_are_suppressed_while_the_partition_is_being_retired(self): + with mock.patch("ydb._topic_writer.topic_writer_multi_asyncio.WriterAsyncIO", _ControllableSubWriter): + writer = TopicWriterMultiAsyncIO(self._driver(), self._settings()) + future = await writer.write_with_ack_future(PublicMessage(b"a", key="a")) + + writer._retiring.add(0) + writer._writers[0].pending[0].set_exception(RuntimeError("session closed by us")) + await asyncio.sleep(0) + + # Expected noise from our own teardown: the message waits to be resent instead. + assert not future.done() + assert set(writer._inflight[0]) == {1} + await writer.close(flush=False) + _retrieve_exceptions([future]) + + async def test_probe_result_is_cached_per_producer(self): + _FAKE_LAST_SEQNO.clear() + _FAKE_LAST_SEQNO[0] = 42 + with mock.patch("ydb._topic_writer.topic_writer_multi_asyncio.WriterAsyncIO", _FakeSubWriter): + writer = TopicWriterMultiAsyncIO(self._driver(), self._settings()) + await writer.wait_init() + + assert await writer._probe_server_seqno(0) == 42 + # A retired producer receives no further writes, so the answer is final: changing + # what the server would say must not change the cached cut. + _FAKE_LAST_SEQNO[0] = 99 + assert await writer._probe_server_seqno(0) == 42 + + await writer.close(flush=False) + + async def test_repartition_task_deregisters_itself_when_done(self): + with mock.patch("ydb._topic_writer.topic_writer_multi_asyncio.WriterAsyncIO", _FakeSubWriter): + writer = TopicWriterMultiAsyncIO(self._driver(), self._settings()) + await writer.wait_init() + + async def noop(partition_id): + return + + writer._on_partition_overloaded = noop + writer._schedule_repartition(0) + task = writer._repartition_tasks[0] + await task + await asyncio.sleep(0) + + # Left registered, a finished task would block every later signal for this partition. + assert 0 not in writer._repartition_tasks + await writer.close(flush=False) + + async def test_repartition_of_an_unknown_partition_is_a_no_op(self): + with mock.patch("ydb._topic_writer.topic_writer_multi_asyncio.WriterAsyncIO", _FakeSubWriter): + driver = self._driver() + writer = TopicWriterMultiAsyncIO(driver, self._settings()) + await writer.wait_init() + describes = driver.describe_calls + + # Already retired by a sibling's event: nothing left to do, and re-describing would + # only race the handler that did retire it. + await writer._handle_repartition(99) + + assert driver.describe_calls == describes + await writer.close(flush=False) + + async def test_failing_inflight_of_an_empty_partition_is_a_no_op(self): + with mock.patch("ydb._topic_writer.topic_writer_multi_asyncio.WriterAsyncIO", _FakeSubWriter): + writer = TopicWriterMultiAsyncIO(self._driver(), self._settings()) + await writer.wait_init() + writer._fail_partition_inflight(0, RuntimeError("boom")) # must not raise + await writer.close(flush=False) + + async def test_idle_reaper_evicts_and_survives_a_failing_pass(self): + with mock.patch("ydb._topic_writer.topic_writer_multi_asyncio.WriterAsyncIO", _FakeSubWriter): + writer = TopicWriterMultiAsyncIO(self._driver(), self._settings(writer_idle_timeout_sec=3)) + await writer.write(PublicMessage(b"a", key="a")) + assert 0 in writer._writers + + writer._last_write_at[0] = writer._loop.time() - 1000 + + # One failing pass must not kill the reaper: it has to keep collecting later. + failed = {"once": False} + real_evict = writer._evict_idle_writers + + async def flaky(): + if not failed["once"]: + failed["once"] = True + raise RuntimeError("eviction failed") + await real_evict() + + async def flaky_then_stop(): + await flaky() + if failed["once"] and 0 not in writer._writers: + writer._closed = True # end the reaper loop once it has done its job + + writer._evict_idle_writers = flaky_then_stop + with mock.patch("ydb._topic_writer.topic_writer_multi_asyncio.asyncio.sleep", _no_sleep): + await asyncio.wait_for(TopicWriterMultiAsyncIO._idle_reaper(weakref.ref(writer), 3), timeout=2) + + assert failed["once"] + assert 0 not in writer._writers, "an idle sub-writer should have been closed" + await writer.close(flush=False) + + async def test_idle_reaper_stops_when_the_writer_is_gone(self): + with mock.patch("ydb._topic_writer.topic_writer_multi_asyncio.asyncio.sleep", _no_sleep): + # A dead weakref is how the reaper learns the writer was garbage collected; it must + # end rather than keep a task alive forever. + await asyncio.wait_for(TopicWriterMultiAsyncIO._idle_reaper(lambda: None, 3), timeout=2) + + async def test_a_busy_partition_is_never_evicted(self): + with mock.patch("ydb._topic_writer.topic_writer_multi_asyncio.WriterAsyncIO", _ControllableSubWriter): + writer = TopicWriterMultiAsyncIO(self._driver(), self._settings(writer_idle_timeout_sec=3)) + future = await writer.write_with_ack_future(PublicMessage(b"a", key="a")) + + writer._last_write_at[0] = writer._loop.time() - 1000 + await writer._evict_idle_writers() + + # Closing it would strand the un-acked message on a dead session. + assert 0 in writer._writers + await writer.close(flush=False) + _retrieve_exceptions([future]) + + async def test_close_failure_on_exit_is_raised_only_on_a_clean_body(self): + class CloseRaises(_FakeSubWriter): + async def close(self, flush=True): + raise RuntimeError("close failed") + + with mock.patch("ydb._topic_writer.topic_writer_multi_asyncio.WriterAsyncIO", _FakeSubWriter): + writer = TopicWriterMultiAsyncIO(self._driver(), self._settings()) + await writer.wait_init() + writer._flush_impl = mock.AsyncMock(side_effect=RuntimeError("close failed")) + writer.close = mock.AsyncMock(side_effect=RuntimeError("close failed")) + + with pytest.raises(RuntimeError, match="close failed"): + await writer.__aexit__(None, None, None) + + # With an exception already travelling, the close failure must not replace it. + await writer.__aexit__(TypeError, TypeError("original"), None) + + async def test_delete_survives_a_loop_that_cannot_schedule(self): + """__del__ can run at interpreter shutdown, when scheduling no longer works.""" + with mock.patch("ydb._topic_writer.topic_writer_multi_asyncio.WriterAsyncIO", _FakeSubWriter): + writer = TopicWriterMultiAsyncIO(self._driver(), self._settings()) + await writer.wait_init() + writer._loop = mock.Mock(is_closed=lambda: False, create_task=mock.Mock(side_effect=RuntimeError)) + + writer.__del__() # must not raise + + writer._closed = True + + async def test_close_survives_a_failing_flush(self): + with mock.patch("ydb._topic_writer.topic_writer_multi_asyncio.WriterAsyncIO", _FakeSubWriter): + writer = TopicWriterMultiAsyncIO(self._driver(), self._settings()) + await writer.wait_init() + writer._flush_impl = mock.AsyncMock(side_effect=RuntimeError("flush failed")) + + # Refusing to close on a flush error would leak every stream the writer holds. + await writer.close() + assert writer._closed + + async def test_repartition_cancellation_is_not_swallowed(self): + with mock.patch("ydb._topic_writer.topic_writer_multi_asyncio.WriterAsyncIO", _FakeSubWriter): + writer = TopicWriterMultiAsyncIO(self._driver(), self._settings()) + await writer.wait_init() + + async def cancelled(partition_id): + raise asyncio.CancelledError() + + # Cancellation means "we are shutting down", not "this partition failed": turning it + # into a recovery attempt would restart work close() is trying to stop. + writer._handle_repartition = cancelled + with pytest.raises(asyncio.CancelledError): + await writer._on_partition_overloaded(0) + + writer._handle_repartition = mock.AsyncMock(side_effect=RuntimeError("boom")) + writer._recover_partition = cancelled + with pytest.raises(asyncio.CancelledError): + await writer._on_partition_overloaded(0) + + await writer.close(flush=False) + + async def test_recovery_after_a_failed_repartition_keeps_the_partition(self): + with mock.patch("ydb._topic_writer.topic_writer_multi_asyncio.WriterAsyncIO", _ControllableSubWriter): + writer = TopicWriterMultiAsyncIO(self._driver(), self._settings()) + future = await writer.write_with_ack_future(PublicMessage(b"a", key="a")) + + writer._handle_repartition = mock.AsyncMock(side_effect=RuntimeError("describe blew up")) + await writer._on_partition_overloaded(0) + + # Repartition failed, but the partition is still ours, so the message is resent + # rather than failed. + assert 0 in writer._partitions + assert not future.done() + writer._writers[0].resolve_all() + await asyncio.sleep(0) + assert future.done() + + await writer.close(flush=False) + + async def test_migration_refuses_a_seqno_the_child_already_persisted(self): + _FAKE_LAST_SEQNO.clear() + with mock.patch("ydb._topic_writer.topic_writer_multi_asyncio.WriterAsyncIO", _FakeSubWriter): + writer = TopicWriterMultiAsyncIO(self._driver(), self._settings()) + await writer.wait_init() + writer._server_init_seqno[2] = 100 + + # Resending it anyway would be rejected by the child's writer as a duplicate seqno, + # taking down the migration of everything behind it. + conflict = writer._migration_conflict(2, 50, {}) + assert isinstance(conflict, TopicWriterError) + assert "already persisted" in str(conflict) + + await writer.close(flush=False) + + async def test_idle_eviction_leaves_recently_used_writers_alone(self): + with mock.patch("ydb._topic_writer.topic_writer_multi_asyncio.WriterAsyncIO", _FakeSubWriter): + writer = TopicWriterMultiAsyncIO(self._driver(), self._settings(writer_idle_timeout_sec=1000)) + await writer.write(PublicMessage(b"a", key="a")) + await asyncio.sleep(0) # let the ack settle so the writer counts as idle + assert not writer._inflight.get(0) + + await writer._evict_idle_writers() + + assert 0 in writer._writers, "a writer used a moment ago is not idle" + await writer.close(flush=False) + + +def test_children_must_cover_the_parent_range(): + """Coverage decides whether a split may be committed to, so each way it can fail matters.""" + cover = TopicWriterMultiAsyncIO._children_cover_parent + parent = _multi_partition(0, from_bound=b"", to_bound=b"") + + # A child with no range at all cannot be reconciled with a bounded parent. + assert cover(parent, [_multi_partition(1, parents=[0])]) is False + + # A hole between the children: keys in it would fall back to the left sibling. + assert ( + cover( + parent, + [ + _multi_partition(1, parents=[0], from_bound=b"", to_bound=b"\x40"), + _multi_partition(2, parents=[0], from_bound=b"\x80", to_bound=b""), + ], + ) + is False + ) + + # A bounded parent fully tiled by its children. + bounded = _multi_partition(0, from_bound=b"a", to_bound=b"m") + assert ( + cover( + bounded, + [ + _multi_partition(1, parents=[0], from_bound=b"a", to_bound=b"f"), + _multi_partition(2, parents=[0], from_bound=b"f", to_bound=b"m"), + ], + ) + is True + ) + + # The same parent left short: the tail of its range has no owner. + assert cover(bounded, [_multi_partition(1, parents=[0], from_bound=b"a", to_bound=b"f")]) is False + + +class _AckOnChildInitSubWriter(_ControllableSubWriter): + """Child sub-writer whose init settles the parent's outstanding ack. + + Models the real race: opening the child's session yields, and the parent's ack can land in + that window -- after the migration decided to resend the message, but before it does. + """ + + parent_writer = None + child_partition_id = None + + async def wait_init(self): + # Only when the CHILD's own session is opened. The dedup-cut probe also opens a session + # (unpinned, so no partition id); acking there would settle the message through the cut + # instead of the race this is meant to reproduce. + parent = type(self).parent_writer + if parent is not None and self.partition_id == type(self).child_partition_id: + type(self).parent_writer = None + parent.resolve_all() + await asyncio.sleep(0) + return await super().wait_init() + + +@pytest.mark.asyncio +async def test_message_acked_while_opening_the_child_is_not_resent(): + """An ack that arrives mid-migration means the message is already written. + + Resending it would duplicate it, and the ack has already completed the caller's future, so + the migration has to notice it lost ownership rather than push it again. + """ + _FAKE_LAST_SEQNO.clear() + mapping = {"a": 0} + chooser = _KeyMapChooser(mapping) + before = [_multi_partition(0), _multi_partition(1)] + after = [ + _split_parent(0, children=[2, 3]), + _multi_partition(1), + _multi_partition(2, parents=[0]), + _multi_partition(3, parents=[0]), + ] + driver = _MultiFakeDescribeDriver([before, after]) + settings = MultiWriterSettings(topic="/local/topic", producer_id_prefix="pfx", partition_chooser=chooser) + + with mock.patch("ydb._topic_writer.topic_writer_multi_asyncio.WriterAsyncIO", _AckOnChildInitSubWriter): + writer = TopicWriterMultiAsyncIO(driver, settings) + await writer.wait_init() + future = await writer.write_with_ack_future(PublicMessage(b"a", key="a")) + + _AckOnChildInitSubWriter.parent_writer = writer._writers[0] + _AckOnChildInitSubWriter.child_partition_id = 2 + mapping["a"] = 2 + await writer._on_partition_overloaded(0) + + assert future.done() and future.exception() is None + resent = [m.seqno for m in writer._writers[2].messages] if 2 in writer._writers else [] + assert resent == [], "a message acked mid-migration was resent anyway" + + await writer.close(flush=False) + + +@pytest.mark.asyncio +class TestTopicWriterMultiAsyncIOBranches: + """The remaining decision points: the side of each branch that only a specific state reaches. + + Most of these are guards against resolving a caller's future twice. A future is settled once + and raises if settled again, and that second attempt happens inside an ack callback where the + exception would be swallowed and the message silently stuck, so the guards are load-bearing. + """ + + def _driver(self): + return _MultiFakeDescribeDriver([[_multi_partition(0), _multi_partition(1)]]) + + def _settings(self, **kwargs): + kwargs.setdefault("topic", "/local/topic") + kwargs.setdefault("producer_id_prefix", "pfx") + kwargs.setdefault("partition_chooser", _KeyMapChooser({"a": 0, "b": 1})) + return MultiWriterSettings(**kwargs) + + def _settled_entry(self, writer, partition_id, seqno): + """Put an already-resolved message into the in-flight map.""" + future = writer._loop.create_future() + future.set_result(PublicWriteResult.Written(offset=1)) + entry = topic_writer_multi_asyncio._InflightMessage( + message=PublicMessage(b"x", key="a"), + user_future=future, + seqno=seqno, + partition_id=partition_id, + ) + writer._inflight.setdefault(partition_id, {})[seqno] = entry + return entry + + async def test_idle_eviction_can_be_switched_off(self): + with mock.patch("ydb._topic_writer.topic_writer_multi_asyncio.WriterAsyncIO", _FakeSubWriter): + writer = TopicWriterMultiAsyncIO(self._driver(), self._settings(writer_idle_timeout_sec=0)) + await writer.wait_init() + + # No reaper task at all, so nothing to cancel on close either. + assert writer._reaper_task is None + await writer.close(flush=False) + + async def test_describe_accepts_a_synchronous_driver(self): + """The sync facade hands the sync driver to the same code, which returns a result + directly instead of a coroutine.""" + + class SyncDriver(_MultiFakeDescribeDriver): + def __call__(self, request, stub, method, wrapper=None, *args, **kwargs): + description = _PublicDescription([_multi_partition(0), _multi_partition(1)]) + + class _Result: + def to_public(self): + return description + + return _Result() + + with mock.patch("ydb._topic_writer.topic_writer_multi_asyncio.WriterAsyncIO", _FakeSubWriter): + writer = TopicWriterMultiAsyncIO(SyncDriver([[]]), self._settings()) + await writer.wait_init() + + assert set(writer._partitions) == {0, 1} + await writer.close(flush=False) + + async def test_already_settled_futures_are_left_alone(self): + with mock.patch("ydb._topic_writer.topic_writer_multi_asyncio.WriterAsyncIO", _ControllableSubWriter): + writer = TopicWriterMultiAsyncIO(self._driver(), self._settings()) + await writer.wait_init() + + # close(): a message whose ack arrived while close was collecting the pending set. + self._settled_entry(writer, 0, 1) + await writer.close(flush=False) + + writer2 = TopicWriterMultiAsyncIO(self._driver(), self._settings()) + await writer2.wait_init() + entry = self._settled_entry(writer2, 0, 1) + # An unusable partition must not try to fail a message that already succeeded. + writer2._fail_partition_inflight(0, RuntimeError("boom")) + assert entry.user_future.exception() is None + await writer2.close(flush=False) + + async def test_settled_futures_survive_recovery_and_migration(self): + _FAKE_LAST_SEQNO.clear() + with mock.patch("ydb._topic_writer.topic_writer_multi_asyncio.WriterAsyncIO", _FakeSubWriter): + writer = TopicWriterMultiAsyncIO(self._driver(), self._settings()) + await writer.wait_init() + + # Below the cut on the recovery path: reported as written, but it already is. + _FAKE_LAST_SEQNO[0] = 5 + entry = self._settled_entry(writer, 0, 3) + await writer._recover_partition(0) + assert entry.user_future.result().offset == 1, "the real ack must not be overwritten" + + # Same on the migration path. + entry2 = self._settled_entry(writer, 1, 3) + await writer._migrate_messages(1, max_seqno=5) + assert entry2.user_future.result().offset == 1 + + await writer.close(flush=False) + + async def test_settled_futures_survive_a_failed_migration(self): + _FAKE_LAST_SEQNO.clear() + chooser = _KeyMapChooser({"a": 0}) + with mock.patch("ydb._topic_writer.topic_writer_multi_asyncio.WriterAsyncIO", _FakeSubWriter): + writer = TopicWriterMultiAsyncIO(self._driver(), self._settings(partition_chooser=chooser)) + await writer.wait_init() + + entry = self._settled_entry(writer, 0, 7) + # Nothing owns the key any more: the tail is failed, but a settled message keeps its + # successful result rather than being turned into an error after the fact. + writer._chooser = _KeyMapChooser({}) + await writer._migrate_messages(0, max_seqno=0) + + assert entry.user_future.exception() is None + await writer.close(flush=False) + + async def test_conflicting_migration_leaves_a_settled_future_alone(self): + _FAKE_LAST_SEQNO.clear() + chooser = _KeyMapChooser({"a": 2}) + with mock.patch("ydb._topic_writer.topic_writer_multi_asyncio.WriterAsyncIO", _FakeSubWriter): + writer = TopicWriterMultiAsyncIO(self._driver(), self._settings(partition_chooser=chooser)) + await writer.wait_init() + writer._partitions[2] = _multi_partition(2) + + entry = self._settled_entry(writer, 0, 3) + # Set it where the fake reports it from: opening the child's writer would otherwise + # overwrite _server_init_seqno with what its init returns. + _FAKE_LAST_SEQNO[2] = 100 # the child already persisted past this seqno + + await writer._migrate_messages(0, max_seqno=0) + + assert entry.user_future.exception() is None + await writer.close(flush=False) + + async def test_recovery_is_skipped_for_an_already_retired_partition(self): + with mock.patch("ydb._topic_writer.topic_writer_multi_asyncio.WriterAsyncIO", _FakeSubWriter): + writer = TopicWriterMultiAsyncIO(self._driver(), self._settings()) + await writer.wait_init() + + writer._handle_repartition = mock.AsyncMock(side_effect=RuntimeError("boom")) + writer._partitions.pop(0) # retired by a sibling's event while we were failing + + # Recreating a sub-writer for it would resurrect a partition we already gave up. + await writer._on_partition_overloaded(0) + + assert 0 not in writer._writers + await writer.close(flush=False) + + async def test_repartition_ignores_unknown_parents_and_known_children(self): + before = [_multi_partition(0), _multi_partition(1)] + # The child lists a parent we never held, and both children are already in our view. + after = [ + _split_parent(0, children=[2, 3]), + _multi_partition(1), + _multi_partition(2, parents=[0, 99]), + _multi_partition(3, parents=[0]), + ] + driver = _MultiFakeDescribeDriver([before, after]) + with mock.patch("ydb._topic_writer.topic_writer_multi_asyncio.WriterAsyncIO", _FakeSubWriter): + writer = TopicWriterMultiAsyncIO(driver, self._settings()) + await writer.wait_init() + writer._partitions[2] = _multi_partition(2) + writer._partitions[3] = _multi_partition(3) + + await writer._on_partition_overloaded(0) + + assert set(writer._partitions) == {1, 2, 3} + await writer.close(flush=False) + + async def test_a_late_ack_does_not_resolve_a_settled_message_twice(self): + """Both ack paths must tolerate a future that is already finished. + + A message can be completed by a dedup cut or by close() while its sub-writer ack is still + on its way. Setting a result twice raises, and this runs inside a done-callback where the + exception would be swallowed -- leaving a message that looks in flight forever. + """ + with mock.patch("ydb._topic_writer.topic_writer_multi_asyncio.WriterAsyncIO", _ControllableSubWriter): + writer = TopicWriterMultiAsyncIO(self._driver(), self._settings()) + await writer.wait_init() + + for seqno, outcome in ((1, "success"), (2, "failure")): + entry = self._settled_entry(writer, 0, seqno) + sub_future = writer._loop.create_future() + entry.sub_future = sub_future + writer._attach_ack(entry) + + if outcome == "success": + sub_future.set_result(PublicWriteResult.Written(offset=9)) + else: + sub_future.set_exception(RuntimeError("too late")) + await asyncio.sleep(0) + + # The original outcome stands, and the message is no longer in flight. + assert entry.user_future.result().offset == 1 + assert seqno not in writer._inflight.get(0, {}) + + await writer.close(flush=False) + + +def test_overlapping_children_do_not_extend_the_covered_range(): + """A child contained in what is already covered adds nothing. + + Letting it move the cursor would make an incomplete set look complete and retire a parent + whose range still has an unowned tail. + """ + cover = TopicWriterMultiAsyncIO._children_cover_parent + parent = _multi_partition(0, from_bound=b"", to_bound=b"") + + assert ( + cover( + parent, + [ + _multi_partition(1, parents=[0], from_bound=b"", to_bound=b"\x80"), + _multi_partition(2, parents=[0], from_bound=b"\x40", to_bound=b"\x60"), + ], + ) + is False + ) + + +def test_bound_chooser_keeps_existing_message_metadata(): + chooser = PublicPartitionByKeyBound() + chooser.add_partitions([PublicDescribeTopicResult.PartitionInfo(0, True, [], [], None, None)]) + message = PublicMessage(b"x", key="user-42", metadata_items={"trace": b"abc"}) + + chooser.choose_partition(message) + + assert message.metadata_items["trace"] == b"abc", "routing must not drop the caller's metadata" + assert PARTITION_KEY_METADATA_KEY in message.metadata_items diff --git a/ydb/_topic_writer/topic_writer_multi_asyncio.py b/ydb/_topic_writer/topic_writer_multi_asyncio.py new file mode 100644 index 000000000..aec5fa5f2 --- /dev/null +++ b/ydb/_topic_writer/topic_writer_multi_asyncio.py @@ -0,0 +1,872 @@ +from __future__ import annotations + +import asyncio +import concurrent.futures +import inspect +import uuid +import logging +import weakref +from dataclasses import dataclass, field +from typing import Callable, Dict, List, Mapping, Optional, Union + +from .topic_writer import ( + Message, + PublicMessage, + PublicWriterSettings, + PublicWriteResult, + PublicWriteResultTypes, + TopicWriterClosedError, + TopicWriterError, + TopicWriterPartitionSplitError, + TopicWriterStopped, +) +from .topic_writer_asyncio import WriterAsyncIO +from .topic_writer_partition_chooser import ( + PublicPartitionByKeyBound, + PublicPartitionByKeyKafka, + PublicPartitionChooser, +) +from .. import _apis, issues +from .._topic_common.common import create_result_wrapper +from .._grpc.grpcwrapper import ydb_topic as _ydb_topic +from .._grpc.grpcwrapper import ydb_topic_public_types as _ydb_topic_public_types +from .._grpc.grpcwrapper.ydb_topic_public_types import PublicAutoPartitioningStrategy, PublicCodec + +_PartitionInfo = _ydb_topic_public_types.PublicDescribeTopicResult.PartitionInfo + +logger = logging.getLogger(__name__) + +# An OVERLOADED response may be an in-progress split/merge (children not yet visible in +# DescribeTopic) or ordinary transient overload. Re-describe a few times before deciding. +_REPARTITION_DISCOVER_ATTEMPTS = 4 +_REPARTITION_DISCOVER_DELAY = 0.25 + +# A per-partition sub-writer with no in-flight messages and no writes for this long is closed; +# it is recreated on demand. Set <= 0 to disable idle eviction. +_DEFAULT_WRITER_IDLE_TIMEOUT = 60.0 + +# A writer opened against a partition that is no longer active never finishes its init handshake. +# Since sub-writers are created while the orchestrator lock is held, an unbounded wait there stalls +# every write, flush and repartition, so the wait is capped. +_WRITER_INIT_TIMEOUT = 30.0 + + +@dataclass +class MultiWriterSettings: + """Settings for the multi-partition (write-by-key) topic writer. + + order of fields IS NOT stable, use keywords only + """ + + topic: str + producer_id_prefix: Optional[str] = None + partition_chooser: Optional[PublicPartitionChooser] = None + auto_seqno: bool = True + auto_created_at: bool = True + codec: Optional[PublicCodec] = None + encoders: Optional[Mapping[PublicCodec, Callable[[bytes], bytes]]] = None + encoder_executor: Optional[concurrent.futures.Executor] = None + max_buffer_size_bytes: Optional[int] = None + max_buffer_messages: Optional[int] = None + buffer_wait_timeout_sec: Optional[float] = None + # Idle per-partition sub-writers are closed after this many seconds and recreated on demand. + # None -> default; <= 0 disables eviction. + writer_idle_timeout_sec: Optional[float] = None + + def __post_init__(self): + if self.producer_id_prefix is None: + self.producer_id_prefix = uuid.uuid4().hex + # partition_chooser is left as-is; when None the writer picks one adaptively + # after describing the topic (Bound if the topic reports key ranges, else Kafka). + + +@dataclass +class _InflightMessage: + message: PublicMessage + user_future: asyncio.Future + seqno: int + partition_id: int + sub_future: Optional[asyncio.Future] = field(default=None) + + +def _is_overloaded(err: BaseException) -> bool: + return isinstance(err, issues.Overloaded) + + +class TopicWriterMultiAsyncIO: + """One logical writer that routes messages to per-partition sub-writers by key. + + Each partition is served by an ordinary :class:`WriterAsyncIO` (buffering, + encoding, reconnection and token refresh are reused as-is). On top of that this + class: + + * routes each message to a partition via the partition chooser; + * owns the in-flight messages and assigns their sequence numbers, so that on an + auto-partition split it can transparently resend the un-acked messages of the + split partition to its children — without duplicating messages that were + already persisted (the ``maxSeqNo`` cut). + """ + + def __init__(self, driver, settings: MultiWriterSettings, _parent=None): + self._loop = asyncio.get_running_loop() + self._driver = driver + self._parent = _parent # keep parent client alive against GC + self._settings = settings + # producer_id_prefix is guaranteed set in __post_init__ + prefix = settings.producer_id_prefix + assert prefix is not None + self._prefix: str = prefix + # Resolved in _init(): the configured chooser, or an adaptive default. + self._chooser: Optional[PublicPartitionChooser] = settings.partition_chooser + self._closed = False + self._lock = asyncio.Lock() + self._writers: Dict[int, WriterAsyncIO] = {} + self._partitions: Dict[int, object] = {} + # partition_id -> {seqno -> in-flight message}, un-acked messages we may resend. + self._inflight: Dict[int, Dict[int, _InflightMessage]] = {} + # One sequence for the whole writer, seeded above the last_seqno of every producer we + # open, so a message keeps its number when a split moves it to another partition. + self._seqno: int = 0 + # partition_id -> last persisted seqno reported by the current sub-writer's init. + self._server_init_seqno: Dict[int, int] = {} + # partition_id -> server's last persisted seqno for a retired producer. Final once read: + # nothing writes under that producer id again. + self._retired_max_seqno: Dict[int, int] = {} + # partition_id -> highest acked seqno (fallback maxSeqNo if a probe fails). + self._max_acked: Dict[int, int] = {} + # Partitions whose sub-writer is being torn down for a repartition or a recovery. Ack + # failures from such a writer are expected and must not reach the user. + self._retiring: set = set() + # partition_id -> in-progress repartition task, so repeated OVERLOADED for one partition + # coalesce and close() can cancel and await them. + self._repartition_tasks: Dict[int, asyncio.Future] = {} + # partition_id -> monotonic time of the last write/creation, for idle eviction. + self._last_write_at: Dict[int, float] = {} + self._idle_timeout = ( + settings.writer_idle_timeout_sec + if settings.writer_idle_timeout_sec is not None + else _DEFAULT_WRITER_IDLE_TIMEOUT + ) + self._init_task = asyncio.ensure_future(self._init()) + self._reaper_task: Optional[asyncio.Future] = None + if self._idle_timeout > 0: + # Hold only a weakref so the reaper does not keep the writer alive against GC. + self._reaper_task = asyncio.ensure_future(self._idle_reaper(weakref.ref(self), self._idle_timeout)) + + async def __aenter__(self) -> "TopicWriterMultiAsyncIO": + return self + + async def __aexit__(self, exc_type, exc_val, exc_tb): + try: + await self.close() + except BaseException: + if exc_val is None: + raise + + def __del__(self): + if self._closed or self._loop.is_closed(): + return + try: + logger.debug("Topic multi-writer was not closed properly. Consider using method close().") + task = self._loop.create_task(self.close(flush=False)) + task.set_name("close multiwriter") + except BaseException: + logger.warning("Something went wrong during multi-writer close in __del__") + + async def _describe(self): + req = _ydb_topic_public_types.DescribeTopicRequestParams(path=self._settings.topic, include_stats=False) + # The async driver returns a coroutine; the sync driver (used behind the + # sync facade) returns the result directly. Support both. + res = self._driver( + req.to_proto(), + _apis.TopicService.Stub, + _apis.TopicService.DescribeTopic, + create_result_wrapper(_ydb_topic.DescribeTopicResult), + ) + if inspect.isawaitable(res): + res = await res + return res.to_public() + + async def _init(self): + description = await self._describe() + leaves = [p for p in description.partitions if p.active and not p.child_partition_ids] + self._partitions = {p.partition_id: p for p in leaves} + if self._chooser is None: + self._chooser = self._default_chooser(leaves, description) + self._chooser.add_partitions(leaves) + + @staticmethod + def _default_chooser(partitions, description=None) -> PublicPartitionChooser: + # Route by key range on auto-partitioned topics, else by Kafka hash. Prefer the + # topic's auto-partitioning strategy as the signal: a single auto partition can + # report no key_range at all (open b""..b""), yet its split-children are bounded, + # which only the Bound chooser can accept. + aps = getattr(description, "auto_partitioning_settings", None) + auto_enabled = aps is not None and aps.strategy not in ( + None, + PublicAutoPartitioningStrategy.UNSPECIFIED, + PublicAutoPartitioningStrategy.DISABLED, + ) + has_key_ranges = any(p.key_range is not None for p in partitions) + if auto_enabled or has_key_ranges: + return PublicPartitionByKeyBound() + return PublicPartitionByKeyKafka() + + async def wait_init(self): + await self._init_task + + def _check_closed(self): + if self._closed: + raise TopicWriterClosedError() + + def _build_writer_settings( + self, + partition_id: int, + with_split_hook: bool, + pin_partition: bool = True, + ) -> PublicWriterSettings: + return PublicWriterSettings( + topic=self._settings.topic, + producer_id="%s-%d" % (self._prefix, partition_id), + # Unpinned (pin_partition=False) leaves the session routed by message group instead, + # which is the only way to reach a partition that is no longer active. + partition_id=partition_id if pin_partition else None, + # The multi-writer assigns sequence numbers itself so it can resend + # messages to child partitions after a split, keeping them monotonic. + auto_seqno=False, + auto_created_at=self._settings.auto_created_at, + codec=self._settings.codec, + encoders=self._settings.encoders, + encoder_executor=self._settings.encoder_executor, + max_buffer_size_bytes=self._settings.max_buffer_size_bytes, + max_buffer_messages=self._settings.max_buffer_messages, + buffer_wait_timeout_sec=self._settings.buffer_wait_timeout_sec, + _on_check_retriable_error=self._make_overloaded_hook(partition_id) if with_split_hook else None, + ) + + async def _get_or_create_writer(self, partition_id: int) -> WriterAsyncIO: + writer = self._writers.get(partition_id) + if writer is None: + writer = WriterAsyncIO(self._driver, self._build_writer_settings(partition_id, with_split_hook=True)) + # Seed the seqno cursor from the producer's last persisted seqno so a + # stable producer_id_prefix resumes numbering instead of colliding. The wait is + # bounded: a partition that went inactive between the routing decision and this call + # never completes init, and we hold the orchestrator lock here. + try: + init_info = await asyncio.wait_for(writer.wait_init(), timeout=_WRITER_INIT_TIMEOUT) + except BaseException: + # Do not register a writer that never became usable, and do not leak its stream. + await self._safe_close(writer) + raise + self._writers[partition_id] = writer + self._last_write_at[partition_id] = self._loop.time() + last_seqno = init_info.last_seqno or 0 + self._server_init_seqno[partition_id] = last_seqno + # Lift the shared cursor above every producer we have opened: each partition has its + # own producer_id and therefore its own persisted history, and a seqno we hand out + # must be new in whichever partition the message ends up in. + self._seqno = max(self._seqno, last_seqno) + return writer + + def _assign_seqno(self, message: PublicMessage) -> int: + """Draw the message's sequence number from the writer-wide cursor. + + One counter for the whole multi-writer, not one per partition: that is what lets a + message keep its seqno when a split moves it to a child. A per-partition number would + mean nothing in another partition's sequence, so it would have to be reassigned -- and a + message that changes identity mid-flight cannot be reconciled with the attempt that may + already have been persisted. Both reference implementations do the same (C++ + `TProducer` `CurrentSeqNo`, Go `orchestrator.currentSeqNo`). + """ + if self._settings.auto_seqno: + self._seqno += 1 + message.seqno = self._seqno + return self._seqno + + if message.seqno is None: + # Bad input, not a stopped writer: the caller disabled auto_seqno and owes us + # a seqno. Reporting this as TopicWriterStopped misleads retry handling. + raise TopicWriterError("message seqno is required when auto_seqno is disabled") + self._seqno = max(self._seqno, message.seqno) + return message.seqno + + def _schedule_repartition(self, partition_id: int) -> None: + """Start (or join) the repartition of one partition. + + A burst of OVERLOADED for the same partition must not start several concurrent + recoveries of it, and the writer must own the task so close() can cancel and await it + instead of leaving it to describe topics and open sub-writers after shutdown. + """ + if self._closed: + return + running = self._repartition_tasks.get(partition_id) + if running is not None and not running.done(): + return + + task = self._loop.create_task(self._on_partition_overloaded(partition_id)) + task.set_name("multiwriter repartition %d" % partition_id) + self._repartition_tasks[partition_id] = task + + def _forget(finished: asyncio.Future) -> None: + if self._repartition_tasks.get(partition_id) is finished: + del self._repartition_tasks[partition_id] + + task.add_done_callback(_forget) + + def _make_overloaded_hook(self, partition_id: int): + def hook(err: BaseException) -> bool: + if _is_overloaded(err): + logger.debug("multi-writer: partition %d overloaded, re-describing (split/merge)", partition_id) + self._schedule_repartition(partition_id) + return True + return False + + return hook + + def _attach_ack(self, entry: _InflightMessage) -> None: + sub_future = entry.sub_future + assert sub_future is not None + sub_future.add_done_callback(lambda f: self._on_sub_result(entry.partition_id, entry.seqno, f)) + + def _on_sub_result(self, partition_id: int, seqno: int, sub_future: asyncio.Future) -> None: + entry = self._inflight.get(partition_id, {}).get(seqno) + if entry is None or entry.sub_future is not sub_future: + return # stale: the message was already resolved or moved to a child + + if sub_future.cancelled(): + return + + exc = sub_future.exception() + if exc is not None: + if isinstance(exc, TopicWriterPartitionSplitError) or partition_id in self._retiring: + # Expected while the partition is torn down: leave the message in flight so the + # repartition (or the in-place recovery) resends it. + return + self._inflight.get(partition_id, {}).pop(seqno, None) + if not entry.user_future.done(): + entry.user_future.set_exception(exc) + return + + # A success is always honoured, including one that lands while the sub-writer is being + # closed: the server persisted the message, and ignoring the ack would leave the dedup + # cut too low and resend an already-written message to the child. + self._inflight.get(partition_id, {}).pop(seqno, None) + self._max_acked[partition_id] = max(self._max_acked.get(partition_id, 0), seqno) + if not entry.user_future.done(): + entry.user_future.set_result(sub_future.result()) + + async def _quiesce_writer(self, partition_id: int) -> None: + """Close the partition's sub-writer and settle the acks it was still holding. + + Closing fails every pending ack, which must not reach the user because those messages + are about to be resent. ``_retiring`` suppresses exactly those failures while letting + successes through, and the yield below gives already-scheduled ack callbacks a chance to + run before the caller reads the dedup cut. + """ + self._retiring.add(partition_id) + writer = self._writers.pop(partition_id, None) + if writer is not None: + await self._safe_close(writer) + await asyncio.sleep(0) + + @staticmethod + async def _safe_close(writer: WriterAsyncIO) -> None: + # A sub-writer stopped by the split hook re-raises TopicWriterPartitionSplitError from + # close(); closing a writer we are discarding must never abort a repartition. + try: + await writer.close(flush=False) + except Exception: # noqa: BLE001 + logger.debug("multi-writer: ignoring error while closing a discarded sub-writer", exc_info=True) + + async def _probe_server_seqno(self, partition_id: int) -> int: + """Ask the server how far this partition's producer actually got. + + The answer must come from the server, not from the acks we happened to receive: a message + can be persisted and its ack lost when the session dies, and treating it as unwritten is + exactly what produces a duplicate on resend. + + The session is opened WITHOUT a partition id. By the time we ask, the partition is + usually inactive, and a session pinned to an inactive partition never completes its init + handshake -- it would hang here holding the orchestrator lock. Unpinned, the session is + routed by message group and still reports this producer's persisted seqno. Both reference + implementations read the cut the same way (C++ `CreateWriteSession(..., false)` sets only + ProducerId/MessageGroupId, Go `createNonDirectWriter` only WithProducerID). + + Results are cached: a retired producer receives no further writes, so its value is final. + """ + cached = self._retired_max_seqno.get(partition_id) + if cached is not None: + return cached + + probe = WriterAsyncIO( + self._driver, + self._build_writer_settings(partition_id, with_split_hook=False, pin_partition=False), + ) + try: + init_info = await asyncio.wait_for(probe.wait_init(), timeout=_WRITER_INIT_TIMEOUT) + finally: + await self._safe_close(probe) + + last_seqno = init_info.last_seqno or 0 + self._retired_max_seqno[partition_id] = last_seqno + logger.debug("multi-writer: partition %d persisted up to seqno %d (server)", partition_id, last_seqno) + return last_seqno + + async def _max_seqno_cut(self, partition_id: int) -> int: + """Dedup cut for a repartition: messages at or below it were persisted to the retiring + partition and must not be resent to the child. + + Only this partition's own producer is asked. A message sitting here cannot have been + persisted under any producer it used earlier in its life: every move is gated by a cut of + at least that producer's server seqno, so anything already stored there was resolved as + written on the spot and never travelled. Its number is therefore strictly above that + value, and a retired producer never grows -- asking again would only repeat an answer that + is already known to be too low to matter. + + Walking further would not merely be redundant. A merge child has two parents whose + branches numbered independently, so the maximum over them pulls in the sibling branch's + history, which says nothing about messages that came down this branch and can be + arbitrarily higher than theirs. + + Raises if the server cannot be asked. Falling back to the highest ack we saw would look + like it worked while quietly reopening the duplicate window this exists to close; the + caller turns the failure into terminal errors on the affected messages instead. + """ + acked = self._max_acked.get(partition_id, 0) + cut = max(acked, await self._probe_server_seqno(partition_id)) + # Worth seeing: this number alone decides resend vs. drop, so a cut below `acked` would + # mean duplicates and one above it would mean loss. + logger.debug( + "multi-writer: dedup cut for partition %d is %d (highest ack seen %d)", + partition_id, + cut, + acked, + ) + return cut + + @staticmethod + def _children_cover_parent(parent, children: List[_PartitionInfo]) -> bool: + """True if the children's key ranges tile the parent's range with no gap. + + A split becomes visible in DescribeTopic one child at a time, so a mid-split describe can + report a single child of two. Retiring the parent on that view would leave the rest of its + key space unowned, and routing -- which locates a partition by the greatest from_bound at + or below the key -- would then send those keys to the left sibling instead: one key on two + branches of the partition graph. + """ + parent_range = getattr(parent, "key_range", None) + if parent_range is None: + return True # topic without key ranges: nothing to verify (and nothing that splits) + + ranges = [] + for child in children: + child_range = getattr(child, "key_range", None) + if child_range is None: + return False # inconsistent with a bounded parent -> treat the view as incomplete + ranges.append((child_range.from_bound, child_range.to_bound)) + ranges.sort(key=lambda r: r[0]) + + # Coverage is "at least the parent's range", not "exactly": a merge child owns the ranges + # of both its parents, so it legitimately covers more than the parent we started from. + parent_end = parent_range.to_bound # empty == end of the key space + cursor = parent_range.from_bound # covered up to here, exclusive + for from_bound, to_bound in ranges: + if from_bound > cursor: + return False # gap between the covered prefix and this child + if not to_bound: + return True # this child runs to the end of the key space + if to_bound > cursor: + cursor = to_bound + if parent_end and cursor >= parent_end: + return True + return False + + async def _discover_children(self, partition_id: int) -> List[_PartitionInfo]: + """Re-describe until the split/merge children of ``partition_id`` appear. + + Returns an empty list if none appear (ordinary transient overload), or if the children + that did appear never covered the parent's key range -- an incomplete graph is a + retry-later state, not a topology we may commit to. + """ + parent = self._partitions.get(partition_id) + children: List[_PartitionInfo] = [] + for attempt in range(_REPARTITION_DISCOVER_ATTEMPTS): + description = await self._describe() + children = [ + p + for p in description.partitions + if p.active and not p.child_partition_ids and partition_id in p.parent_partition_ids + ] + if children and self._children_cover_parent(parent, children): + return children + if attempt + 1 < _REPARTITION_DISCOVER_ATTEMPTS: + await asyncio.sleep(_REPARTITION_DISCOVER_DELAY) + if children: + logger.warning( + "multi-writer: children of partition %d never covered its key range; keeping the" + " partition instead of routing keys into the uncovered range", + partition_id, + ) + return [] + + def _fail_partition_inflight(self, partition_id: int, err: BaseException) -> None: + """Give every in-flight message of an unusable partition a terminal outcome. + + Once repartition and recovery have both failed there is no sub-writer left to ack these + messages, so their futures would never resolve and flush()/close(flush=True) would wait + on them forever. An accepted message must always end in success or failure. + """ + entries = self._inflight.pop(partition_id, {}) + if not entries: + return + logger.error( + "multi-writer: partition %d is unusable, failing %d in-flight messages: %s", + partition_id, + len(entries), + err, + ) + for entry in entries.values(): + entry.sub_future = None + if not entry.user_future.done(): + entry.user_future.set_exception(err) + + async def _on_partition_overloaded(self, partition_id: int): + """Entry point for the OVERLOADED hook: handle a repartition, or recover on failure. + + The hook force-stops the sub-writer, so if handling fails we must not leave the + partition's messages stranded — recreate the writer and resend them. If that fails too, + the messages are failed explicitly rather than left without an owner. + """ + try: + await self._handle_repartition(partition_id) + return + except asyncio.CancelledError: + raise + except Exception as err: + logger.exception("multi-writer: repartition of partition %d failed; recovering", partition_id) + failure: BaseException = err + + try: + async with self._lock: + if partition_id in self._partitions: + await self._recover_partition(partition_id) + return + except asyncio.CancelledError: + raise + except Exception as err: + logger.exception("multi-writer: recovery of partition %d failed", partition_id) + failure = err + + async with self._lock: + self._fail_partition_inflight(partition_id, failure) + + async def _handle_repartition(self, partition_id: int): + """Resolve an OVERLOADED partition: split, merge, or ordinary transient overload. + + A split turns one partition into two children (each with a single parent); a merge + turns two into one child (with both as parents). Both are discovered by finding the + active leaf partitions that list ``partition_id`` as a parent. All parents of those + children that we still hold are retired together, so a merge does not leave the + sibling parent lingering with an overlapping key range. If no children ever appear + the overload was transient and the partition is recovered in place. + """ + async with self._lock: + if partition_id not in self._partitions: + return # already handled by a sibling parent's event + assert self._chooser is not None # resolved by _init() before any repartition + + children = await self._discover_children(partition_id) + if not children: + # Transient overload, not a topology change: keep the partition. + await self._recover_partition(partition_id) + return + + retired = {partition_id} + for child in children: + for parent in child.parent_partition_ids: + if parent in self._partitions: + retired.add(parent) + + # Update the routing view first: add children, drop every retired parent, so + # migration re-routes only to the surviving partitions (no overlapping ranges). + new_children = [c for c in children if c.partition_id not in self._partitions] + if new_children: + self._chooser.add_partitions(new_children) + for child in new_children: + self._partitions[child.partition_id] = child + for old in retired: + self._chooser.remove_partition(old) + self._partitions.pop(old, None) + + # Quiesce every retired parent BEFORE reading its cutoff, so a sibling cannot + # persist a message after its maxSeqNo was probed (which would duplicate on resend). + try: + for old in retired: + await self._quiesce_writer(old) + + for old in retired: + if self._inflight.get(old): + await self._migrate_messages(old, await self._max_seqno_cut(old)) + else: + self._inflight.pop(old, None) + finally: + self._retiring.difference_update(retired) + + async def _recover_partition(self, partition_id: int): + # The hook stopped the sub-writer; drop it and resend the partition's in-flight + # messages to a fresh writer for the SAME partition, keeping their seqnos. + await self._quiesce_writer(partition_id) + try: + writer = await self._get_or_create_writer(partition_id) + + # The partition is still active, so this writer's init carries the server's real last + # persisted seqno -- an exact dedup cut, unlike the split case. Messages at or below it + # were written and their ack was lost with the stream; resending one is not idempotent + # here, because the writer rejects a seqno it has already seen ("Message seqno is + # duplicated") before the server ever gets a chance to deduplicate it, and that error + # would abort the resend of every message after it. + cut = max(self._server_init_seqno.get(partition_id, 0), self._max_acked.get(partition_id, 0)) + for seqno, entry in sorted(self._inflight.get(partition_id, {}).items()): + if seqno <= cut: + self._inflight.get(partition_id, {}).pop(seqno, None) + self._max_acked[partition_id] = max(self._max_acked.get(partition_id, 0), seqno) + if not entry.user_future.done(): + entry.user_future.set_result(PublicWriteResult.Written(offset=-1)) + continue + entry.message.seqno = seqno + sub_future = await writer.write_with_ack_future(entry.message) + assert not isinstance(sub_future, list) # single message -> single future + entry.sub_future = sub_future + self._attach_ack(entry) + finally: + self._retiring.discard(partition_id) + + def _migration_conflict( + self, + child_id: int, + seqno: int, + target: Dict[int, _InflightMessage], + ) -> Optional[BaseException]: + """Why this seqno cannot be carried into ``child_id``, or None if it can. + + Carrying the number over is only safe while it is free in the child's sequence, and two + things can take it: another in-flight message (manual seqnos are checked for uniqueness + only within one partition, so a merge can bring two equal ones into the same child), or + the child's own producer, if a stable ``producer_id_prefix`` already wrote that far in an + earlier run. In the second case the sub-writer would reject the resend as a duplicate + seqno and take down the migration of every message behind it. + """ + if seqno in target: + return TopicWriterError( + "seqno %d is already in flight on partition %d: manual seqnos must be unique" + " across partitions to survive a repartition" % (seqno, child_id) + ) + server_seqno = self._server_init_seqno.get(child_id, 0) + if seqno <= server_seqno: + return TopicWriterError( + "seqno %d cannot be resent to partition %d: its producer has already persisted" + " up to %d" % (seqno, child_id, server_seqno) + ) + return None + + async def _migrate_messages(self, partition_id: int, max_seqno: int): + entries = self._inflight.get(partition_id, {}) + # Snapshot in seqno order so re-routed messages keep their relative order. + for seqno, entry in sorted(entries.items()): + if seqno <= max_seqno: + # Already persisted to the retired partition: resolve as written + # (offset is unknown because the ack was lost) and do not resend. + self._inflight.get(partition_id, {}).pop(seqno, None) + if not entry.user_future.done(): + entry.user_future.set_result(PublicWriteResult.Written(offset=-1)) + continue + + assert self._chooser is not None + try: + child_id = self._chooser.choose_partition(entry.message) + child_writer = await self._get_or_create_writer(child_id) + except Exception as err: # noqa: BLE001 + # The message cannot be placed: no ready leaf owns its key, or the child itself + # went inactive (splits cascade). Fail it and everything after it -- dropping + # would lose the message silently, and skipping ahead would reorder the key. + logger.warning( + "multi-writer: cannot migrate messages of partition %d, failing seqno >= %d: %s", + partition_id, + seqno, + err, + ) + # Everything before this message has already been popped by an earlier + # iteration, so what is left in `entries` is exactly this one and its tail. + for pending_seqno, pending in sorted(entries.items()): + self._inflight.get(partition_id, {}).pop(pending_seqno, None) + if not pending.user_future.done(): + pending.user_future.set_exception(err) + break + + if self._inflight.get(partition_id, {}).get(seqno) is not entry: + continue # acked while the child writer was being opened -> nothing to resend + + # The message keeps its seqno: only its target partition changes. The number comes + # from the writer-wide sequence, so it stays valid -- and unchanged identity is what + # lets the child's cut still describe this exact message. + target = self._inflight.setdefault(child_id, {}) + conflict = self._migration_conflict(child_id, seqno, target) + if conflict is not None: + logger.error("multi-writer: %s", conflict) + self._inflight.get(partition_id, {}).pop(seqno, None) + if not entry.user_future.done(): + entry.user_future.set_exception(conflict) + continue + + self._inflight.get(partition_id, {}).pop(seqno, None) + entry.partition_id = child_id + entry.message.seqno = seqno + target[seqno] = entry + sub_future = await child_writer.write_with_ack_future(entry.message) + assert not isinstance(sub_future, list) # single message -> single future + entry.sub_future = sub_future + self._attach_ack(entry) + + self._inflight.pop(partition_id, None) + + async def write_with_ack_future( + self, + messages: Union[Message, List[Message]], + ) -> Union[asyncio.Future, List[asyncio.Future]]: + self._check_closed() + await self.wait_init() + + input_single_message = not isinstance(messages, list) + raw = messages if isinstance(messages, list) else [messages] + converted = [PublicMessage._create_message(m) for m in raw] + + futures: List[asyncio.Future] = [] + async with self._lock: + assert self._chooser is not None # resolved by _init(), awaited above + for message in converted: + partition_id = self._chooser.choose_partition(message) + writer = await self._get_or_create_writer(partition_id) + self._last_write_at[partition_id] = self._loop.time() + seqno = self._assign_seqno(message) + if seqno in self._inflight.get(partition_id, {}): + raise TopicWriterError("duplicate in-flight seqno %d for partition %d" % (seqno, partition_id)) + + user_future: asyncio.Future = self._loop.create_future() + entry = _InflightMessage( + message=message, + user_future=user_future, + seqno=seqno, + partition_id=partition_id, + ) + # Record the message only after the sub-writer accepts it, so a failed + # admission (buffer timeout, stopped writer) does not leak an in-flight entry + # whose future would never resolve. + sub_future = await writer.write_with_ack_future(message) + assert not isinstance(sub_future, list) # single message -> single future + entry.sub_future = sub_future + self._inflight.setdefault(partition_id, {})[seqno] = entry + self._attach_ack(entry) + futures.append(user_future) + + return futures[0] if input_single_message else futures + + async def write(self, messages: Union[Message, List[Message]]): + await self.write_with_ack_future(messages) + + async def write_with_ack( + self, + messages: Union[Message, List[Message]], + ) -> Union[PublicWriteResultTypes, List[PublicWriteResultTypes]]: + futures = await self.write_with_ack_future(messages) + future_list = futures if isinstance(futures, list) else [futures] + await asyncio.wait(future_list) + results = [f.result() for f in future_list] + return results if isinstance(futures, list) else results[0] + + @staticmethod + async def _idle_reaper(mw_ref: "weakref.ref", idle_timeout: float): + interval = max(1.0, idle_timeout / 3) + try: + while True: + await asyncio.sleep(interval) + mw = mw_ref() + if mw is None or mw._closed: + return + try: + await mw._evict_idle_writers() + except Exception: # noqa: BLE001 + logger.debug("multi-writer: idle eviction pass failed", exc_info=True) + del mw # drop the strong ref so the writer can be GC'd while we sleep + except asyncio.CancelledError: + pass + + async def _evict_idle_writers(self): + # Close sub-writers with no in-flight messages that have not been written to for the idle + # timeout; they are recreated on demand. Done under the lock so a concurrent write cannot + # open a second session for the same producer while we close the old one. + now = self._loop.time() + async with self._lock: + for partition_id in list(self._writers.keys()): + if self._inflight.get(partition_id): + continue # has un-acked messages -> not idle + if now - self._last_write_at.get(partition_id, now) < self._idle_timeout: + continue + writer = self._writers.pop(partition_id) + self._last_write_at.pop(partition_id, None) + logger.debug("multi-writer: evicting idle sub-writer for partition %d", partition_id) + await self._safe_close(writer) + + def _pending_user_futures(self) -> List[asyncio.Future]: + return [entry.user_future for part in self._inflight.values() for entry in part.values()] + + async def _flush_impl(self): + await self.wait_init() + async with self._lock: + writers = list(self._writers.values()) + pending = self._pending_user_futures() + await asyncio.gather(*(w.flush() for w in writers), return_exceptions=True) + if pending: + await asyncio.gather(*pending, return_exceptions=True) + + async def flush(self): + self._check_closed() + await self._flush_impl() + + async def close(self, *, flush: bool = True): + if self._closed: + return + + # Flush BEFORE marking closed (flush() itself refuses to run on a closed writer), + # but only if init completed — otherwise nothing was written. + init_done = self._init_task.done() and not self._init_task.cancelled() and self._init_task.exception() is None + if flush and init_done: + try: + await self._flush_impl() + except BaseException: + logger.debug("multi-writer: flush during close failed", exc_info=True) + + self._closed = True + if not self._init_task.done(): + self._init_task.cancel() + if self._reaper_task is not None and not self._reaper_task.done(): + self._reaper_task.cancel() + + # Stop repartitions before touching the writers: an in-flight one holds the lock, opens + # sub-writers and mutates the topology, none of which may outlive the multi-writer. + repartitions = list(self._repartition_tasks.values()) + self._repartition_tasks.clear() + for task in repartitions: + if not task.done(): + task.cancel() + if repartitions: + await asyncio.gather(*repartitions, return_exceptions=True) + + async with self._lock: + writers = list(self._writers.values()) + self._writers.clear() + pending = self._pending_user_futures() + self._inflight.clear() + await asyncio.gather(*(w.close(flush=False) for w in writers), return_exceptions=True) + for future in pending: + if not future.done(): + future.set_exception(TopicWriterStopped()) diff --git a/ydb/_topic_writer/topic_writer_multi_sync.py b/ydb/_topic_writer/topic_writer_multi_sync.py new file mode 100644 index 000000000..5e599a745 --- /dev/null +++ b/ydb/_topic_writer/topic_writer_multi_sync.py @@ -0,0 +1,132 @@ +from __future__ import annotations + +import asyncio +import logging +import typing +from concurrent.futures import Future +from typing import List, Optional, Union + +from .._grpc.grpcwrapper.common_utils import SupportedDriverType +from .topic_writer import ( + Message, + PublicWriteResult, + TopicWriterClosedError, +) +from .topic_writer_multi_asyncio import MultiWriterSettings, TopicWriterMultiAsyncIO +from .._topic_common.common import ( + _get_shared_event_loop, + TimeoutType, + CallFromSyncToAsync, +) + +logger = logging.getLogger(__name__) + + +class TopicWriterMultiSync: + _caller: CallFromSyncToAsync + _async_writer: TopicWriterMultiAsyncIO + _closed: bool + _parent: typing.Any # need for prevent close parent client by GC + + def __init__( + self, + driver: SupportedDriverType, + settings: MultiWriterSettings, + *, + eventloop: Optional[asyncio.AbstractEventLoop] = None, + _parent=None, + ): + self._closed = False + + if eventloop: + loop = eventloop + else: + loop = _get_shared_event_loop() + + self._caller = CallFromSyncToAsync(loop) + + async def create_async_writer(): + return TopicWriterMultiAsyncIO(driver, settings) + + self._async_writer = self._caller.safe_call_with_result(create_async_writer(), None) + self._parent = _parent + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + try: + self.close() + except BaseException: + if exc_val is None: + raise + + def __del__(self): + if not self._closed: + try: + logger.debug("Topic multi-writer was not closed properly. Consider using method close().") + self.close(flush=False) + except BaseException: + logger.warning("Something went wrong during multi-writer close in __del__") + + def close(self, *, flush: bool = True, timeout: TimeoutType = None): + if self._closed: + return + + logger.debug("Close topic multi-writer") + self._closed = True + + self._caller.safe_call_with_result(self._async_writer.close(flush=flush), timeout) + + def _check_closed(self): + if self._closed: + raise TopicWriterClosedError() + + def async_flush(self) -> Future: + self._check_closed() + return self._caller.unsafe_call_with_future(self._async_writer.flush()) + + def flush(self, *, timeout: TimeoutType = None): + self._check_closed() + logger.debug("flush multi-writer") + return self._caller.unsafe_call_with_result(self._async_writer.flush(), timeout) + + def async_wait_init(self) -> Future: + self._check_closed() + return self._caller.unsafe_call_with_future(self._async_writer.wait_init()) + + def wait_init(self, *, timeout: TimeoutType = None): + self._check_closed() + logger.debug("wait multi-writer init") + return self._caller.unsafe_call_with_result(self._async_writer.wait_init(), timeout) + + def write( + self, + messages: Union[Message, List[Message]], + timeout: TimeoutType = None, + ): + self._check_closed() + logger.debug( + "write %s messages", + len(messages) if isinstance(messages, list) else 1, + ) + self._caller.safe_call_with_result(self._async_writer.write(messages), timeout) + + def async_write_with_ack( + self, + messages: Union[Message, List[Message]], + ) -> Future[Union[PublicWriteResult, List[PublicWriteResult]]]: + self._check_closed() + return self._caller.unsafe_call_with_future(self._async_writer.write_with_ack(messages)) + + def write_with_ack( + self, + messages: Union[Message, List[Message]], + timeout: Union[float, None] = None, + ) -> Union[PublicWriteResult, List[PublicWriteResult]]: + self._check_closed() + logger.debug( + "write_with_ack %s messages", + len(messages) if isinstance(messages, list) else 1, + ) + return self._caller.unsafe_call_with_result(self._async_writer.write_with_ack(messages), timeout=timeout) diff --git a/ydb/_topic_writer/topic_writer_partition_chooser.py b/ydb/_topic_writer/topic_writer_partition_chooser.py new file mode 100644 index 000000000..38e0dfa1e --- /dev/null +++ b/ydb/_topic_writer/topic_writer_partition_chooser.py @@ -0,0 +1,195 @@ +import abc +import bisect +from typing import Callable, List, Tuple + +from .topic_writer import PublicMessage +from .._grpc.grpcwrapper.ydb_topic_public_types import PublicDescribeTopicResult + +_MASK32 = 0xFFFFFFFF +_MASK64 = 0xFFFFFFFFFFFFFFFF + +# Metadata key the server uses to validate bound-based partition selection. +PARTITION_KEY_METADATA_KEY = "__partition_key" + +PartitionInfo = PublicDescribeTopicResult.PartitionInfo + + +def murmur2_32(data: bytes, seed: int = 0) -> int: + """MurmurHash2, 32-bit, little-endian. Matches the Kafka BuiltInPartitioner + and the YDB Go SDK ``xhash.Murmur2Hash32``.""" + m = 0x5BD1E995 + r = 24 + n = len(data) + h = (seed ^ n) & _MASK32 + + body = n - (n % 4) + for i in range(0, body, 4): + k = int.from_bytes(data[i : i + 4], "little") + k = (k * m) & _MASK32 + k ^= k >> r + k = (k * m) & _MASK32 + h = (h * m) & _MASK32 + h ^= k + + rem = n % 4 + if rem: + tail = data[body:] + if rem >= 3: + h ^= tail[2] << 16 + if rem >= 2: + h ^= tail[1] << 8 + h ^= tail[0] + h = (h * m) & _MASK32 + + h ^= h >> 13 + h = (h * m) & _MASK32 + h ^= h >> 15 + return h & _MASK32 + + +def murmur64a(data: bytes, seed: int = 0) -> int: + """MurmurHash64A, 64-bit, little-endian. Matches the YDB Go SDK + ``xhash.Murmur2Hash64A`` and the C++ default partitioning key hasher.""" + m = 0xC6A4A7935BD1E995 + r = 47 + n = len(data) + h = (seed ^ ((n * m) & _MASK64)) & _MASK64 + + body = n - (n % 8) + for i in range(0, body, 8): + k = int.from_bytes(data[i : i + 8], "little") + k = (k * m) & _MASK64 + k ^= k >> r + k = (k * m) & _MASK64 + h ^= k + h = (h * m) & _MASK64 + + rem = n % 8 + if rem: + tail = data[body:] + for j in range(rem - 1, -1, -1): + h ^= tail[j] << (8 * j) + h = (h * m) & _MASK64 + + h ^= h >> r + h = (h * m) & _MASK64 + h ^= h >> r + return h & _MASK64 + + +def default_bound_key_hasher(key: str) -> bytes: + """Hash a routing key the way the YDB server expects for bound-based + partition selection: MurmurHash64A(seed=0) as 8 big-endian bytes.""" + return murmur64a(key.encode("utf-8"), 0).to_bytes(8, "big") + + +class PublicPartitionChooser(abc.ABC): + """Maps a message to a target partition id and tracks the live partition set.""" + + @abc.abstractmethod + def choose_partition(self, message: PublicMessage) -> int: ... + + @abc.abstractmethod + def add_partitions(self, partitions: List[PartitionInfo]) -> None: ... + + @abc.abstractmethod + def remove_partition(self, partition_id: int) -> None: ... + + +class PublicPartitionByKeyKafka(PublicPartitionChooser): + """Kafka-compatible routing: ``murmur2_32(key) % partitions_count``. + + Ignores server key ranges, so it fits topics with a fixed partition count. + """ + + def __init__(self): + self._partitions: List[int] = [] + + def add_partitions(self, partitions: List[PartitionInfo]) -> None: + for p in partitions: + # Any key range at all means the topic is auto-partitioned, including the fully open + # one a single partition reports before its first split. Accepting that would work + # until the split, then break: the children come back bounded, and modulo routing + # cannot place a key by bounds. + if p.key_range is not None: + raise ValueError("PublicPartitionByKeyKafka does not support partition key ranges") + self._partitions.append(p.partition_id) + self._partitions.sort() + + def remove_partition(self, partition_id: int) -> None: + self._partitions = [p for p in self._partitions if p != partition_id] + + def choose_partition(self, message: PublicMessage) -> int: + if not self._partitions: + raise ValueError("no partitions configured for partition chooser") + # Apache Kafka's DefaultPartitioner applies toPositive() (mask the sign bit) to the + # murmur2 hash before the modulo; match it so the same key lands on the same partition. + h = murmur2_32((message.key or "").encode("utf-8"), 0) & 0x7FFFFFFF + return self._partitions[h % len(self._partitions)] + + +class PublicPartitionByKeyBound(PublicPartitionChooser): + """Server-accurate routing: hashes the key and selects the partition whose + ``[from_bound, to_bound)`` key range owns it. Mirrors YDB auto-partitioning, + so the same key lands where the server expects it. + """ + + def __init__(self, key_hasher: Callable[[str], bytes] = default_bound_key_hasher): + self._key_hasher = key_hasher + # (from_bound, to_bound, partition_id) sorted by from_bound. An empty from_bound is the + # start of the key space, an empty to_bound is its end. + self._partitions: List[Tuple[bytes, bytes, int]] = [] + + def add_partitions(self, partitions: List[PartitionInfo]) -> None: + # Normalise and sort before validating: the caller may pass partitions in any order + # (DescribeTopic does not promise one), and only the partition that sorts leftmost may + # carry an open lower bound. Validating in argument order rejects a valid set that + # happens to arrive reversed. + added: List[Tuple[bytes, bytes, int]] = [] + for p in partitions: + from_bound = p.key_range.from_bound if p.key_range is not None else b"" + to_bound = p.key_range.to_bound if p.key_range is not None else b"" + added.append((from_bound, to_bound, p.partition_id)) + added.sort(key=lambda x: x[0]) + + for i, (from_bound, _to_bound, partition_id) in enumerate(added): + if i > 0 and not from_bound: + raise ValueError( + "partition %d has no from_bound key range, but only the leftmost partition may" + " have an open lower bound" % partition_id + ) + + self._partitions.extend(added) + self._partitions.sort(key=lambda x: x[0]) + + def remove_partition(self, partition_id: int) -> None: + self._partitions = [p for p in self._partitions if p[-1] != partition_id] + + def choose_partition(self, message: PublicMessage) -> int: + if not self._partitions: + raise ValueError("no partitions configured for partition chooser") + hashed = self._key_hasher(message.key or "") + + if message.metadata_items is None: + message.metadata_items = {} + message.metadata_items[PARTITION_KEY_METADATA_KEY] = hashed + + # First partition whose from_bound is strictly greater than the hashed key; the owning + # partition is the previous one. Searched in place with a key rather than by building a + # list of bounds, which would allocate per message on the hot path. + idx = bisect.bisect_right(self._partitions, hashed, key=lambda entry: entry[0]) + if idx == 0: + raise RuntimeError("inconsistent partition bounds: lower-bound search returned 0") + + from_bound, to_bound, partition_id = self._partitions[idx - 1] + # A key range is [from_bound, to_bound); an empty to_bound means the range runs to the + # end of the key space. Landing past to_bound means the known partitions have a hole -- + # e.g. only one child of a split is visible yet -- and the greatest-lower-bound search + # silently points at the *sibling* that owns the range to the left. Refuse instead: + # routing a key into a sibling branch is exactly what breaks one-key-one-lineage. + if to_bound and hashed >= to_bound: + raise RuntimeError( + "key is not covered by any known partition: it is above partition %d to_bound;" + " the partition set is incomplete" % partition_id + ) + return partition_id diff --git a/ydb/_topic_writer/topic_writer_test.py b/ydb/_topic_writer/topic_writer_test.py index 5b857ab75..141a0ca18 100644 --- a/ydb/_topic_writer/topic_writer_test.py +++ b/ydb/_topic_writer/topic_writer_test.py @@ -11,12 +11,38 @@ PublicMessage, PublicWriterSettings, TopicWriterBufferFullError, + TopicWriterClosedError, _split_messages_by_size, _split_messages_for_send, messages_to_proto_requests, ) from .topic_writer_asyncio import WriterAsyncIOReconnector from .topic_writer_sync import WriterSync +from .topic_writer_multi_asyncio import MultiWriterSettings +from .topic_writer_multi_sync import TopicWriterMultiSync +from .topic_writer_partition_chooser import ( + murmur2_32, + murmur64a, + default_bound_key_hasher, + PublicPartitionByKeyKafka, + PublicPartitionByKeyBound, + PARTITION_KEY_METADATA_KEY, +) +from .._grpc.grpcwrapper.ydb_topic_public_types import PublicDescribeTopicResult + + +def _partition_info(partition_id: int, from_bound: bytes = None, to_bound: bytes = None): + key_range = None + if from_bound is not None or to_bound is not None: + key_range = PublicDescribeTopicResult.PartitionKeyRange(from_bound=from_bound or b"", to_bound=to_bound or b"") + return PublicDescribeTopicResult.PartitionInfo( + partition_id=partition_id, + active=True, + child_partition_ids=[], + parent_partition_ids=[], + partition_stats=None, + key_range=key_range, + ) @pytest.mark.parametrize( @@ -270,3 +296,327 @@ def do_write(): assert not write_errors, f"unexpected error: {write_errors}" writer.close(flush=False) + + +# Golden vectors generated from the YDB Go SDK pkg/xhash reference implementation +# (Murmur2Hash32 / Murmur2Hash64A, seed=0). A mismatch means keys would route to +# the wrong partition, so these are byte-exact assertions. +_MURMUR_GOLDEN = [ + ("", 0, 0), + ("a", 2456313694, 510903276987443985), + ("hello", 3848350155, 2191231550387646743), + ("hello world, murmur2 hash", 1305234166, 15193844207144850389), + ("мурмур2-хэш", 1364064206, 2094682108092698226), + ("user-42", 3766944517, 17854748655353381905), + ("0", 1111412596, 5533571732986600803), + ("key-with-длинный-unicode-🚀", 901185517, 8765931500921732560), +] + + +@pytest.mark.parametrize("text,h32,h64", _MURMUR_GOLDEN) +def test_murmur_hashes_match_go_golden_vectors(text, h32, h64): + data = text.encode("utf-8") + assert murmur2_32(data, 0) == h32 + assert murmur64a(data, 0) == h64 + + +def test_default_bound_key_hasher_is_big_endian_murmur64a(): + for text, _h32, h64 in _MURMUR_GOLDEN: + assert default_bound_key_hasher(text) == h64.to_bytes(8, "big") + + +def test_kafka_chooser_routes_by_murmur2_modulo(): + chooser = PublicPartitionByKeyKafka() + chooser.add_partitions([_partition_info(0), _partition_info(1), _partition_info(2)]) + for key in ["", "a", "user-42", "hello", "мурмур2-хэш"]: + expected = (murmur2_32(key.encode("utf-8"), 0) & 0x7FFFFFFF) % 3 + assert chooser.choose_partition(PublicMessage(b"x", key=key)) == expected + + +def test_kafka_chooser_matches_apache_kafka_partition(): + # Apache Kafka DefaultPartitioner: toPositive(murmur2(key)) % numPartitions. + # Golden index for key "a" with 3 partitions is 2 (not 1, which the raw hash gives). + chooser = PublicPartitionByKeyKafka() + chooser.add_partitions([_partition_info(i) for i in range(3)]) + assert chooser.choose_partition(PublicMessage(b"x", key="a")) == 2 + + +def test_kafka_chooser_rejects_key_ranges(): + chooser = PublicPartitionByKeyKafka() + with pytest.raises(ValueError): + chooser.add_partitions([_partition_info(0, from_bound=b"\x10")]) + + +def test_kafka_chooser_raises_without_partitions(): + with pytest.raises(ValueError): + PublicPartitionByKeyKafka().choose_partition(PublicMessage(b"x", key="k")) + + +def test_bound_chooser_routes_into_owning_key_range(): + lo = b"\x55" * 8 + hi = b"\xaa" * 8 + chooser = PublicPartitionByKeyBound() + chooser.add_partitions([_partition_info(0), _partition_info(1, from_bound=lo), _partition_info(2, from_bound=hi)]) + + for key in ["a", "b", "user-42", "hello", "zzz", "мурмур2-хэш", "0"]: + hashed = default_bound_key_hasher(key) + if hashed < lo: + expected = 0 + elif hashed < hi: + expected = 1 + else: + expected = 2 + assert chooser.choose_partition(PublicMessage(b"x", key=key)) == expected + + +def test_bound_chooser_stamps_partition_key_metadata(): + chooser = PublicPartitionByKeyBound() + chooser.add_partitions([_partition_info(0)]) + message = PublicMessage(b"x", key="user-42") + chooser.choose_partition(message) + assert message.metadata_items[PARTITION_KEY_METADATA_KEY] == default_bound_key_hasher("user-42") + + +def test_bound_chooser_requires_from_bound_on_non_first_partition(): + chooser = PublicPartitionByKeyBound() + with pytest.raises(ValueError): + chooser.add_partitions([_partition_info(0), _partition_info(1)]) + + +def test_bound_chooser_accepts_partitions_in_any_order(): + """A valid range set must not be rejected just because it arrived out of order. + + DescribeTopic does not promise an ordering, so validation has to run on the sorted set: + the partition with the open lower bound is the leftmost one, not necessarily the first + argument. + """ + lo = b"\x55" * 8 + reversed_order = [_partition_info(1, from_bound=lo), _partition_info(0)] + + chooser = PublicPartitionByKeyBound() + chooser.add_partitions(reversed_order) + + assert sorted(p[-1] for p in chooser._partitions) == [0, 1] + + +def test_bound_chooser_rejects_key_above_the_last_to_bound(): + """A key outside every known range must fail routing instead of landing on a neighbour. + + The chooser only indexes `from_bound` and takes the greatest one <= the hashed key, so a + key above the last partition's `to_bound` silently lands in that partition. That happens + whenever the known partition set has a hole -- e.g. a split whose second child is not + visible in DescribeTopic yet -- and sends the key into a *sibling* branch of the partition + graph, breaking "one key -> one lineage". + + Per the C++ producer spec the lookup is two steps, not one: + partition = greatest from_bound <= key + check key < partition.to_bound + and a key that maps to no ready leaf must be retried against a refreshed graph. + + The exact exception type is a fix-time decision; the contract asserted here is only that + routing refuses to guess. + """ + chooser = PublicPartitionByKeyBound(key_hasher=lambda key: key.encode("utf-8")) + chooser.add_partitions([_partition_info(0, from_bound=b"", to_bound=b"m")]) + + assert chooser.choose_partition(PublicMessage(b"x", key="apple")) == 0 + with pytest.raises((ValueError, RuntimeError)): + chooser.choose_partition(PublicMessage(b"x", key="zebra")) + + +def test_bound_chooser_raises_without_partitions(): + with pytest.raises(ValueError): + PublicPartitionByKeyBound().choose_partition(PublicMessage(b"x", key="k")) + + +def test_bound_chooser_rejects_key_below_every_from_bound(): + """A key under the leftmost bound means the partition set lost its first partition. + + Nothing owns that key, and picking the nearest partition anyway would put it in a range the + server does not associate with it, so routing has to refuse instead of guessing. + """ + chooser = PublicPartitionByKeyBound(key_hasher=lambda key: key.encode("utf-8")) + chooser.add_partitions([_partition_info(1, from_bound=b"m", to_bound=b"")]) + + with pytest.raises(RuntimeError): + chooser.choose_partition(PublicMessage(b"x", key="apple")) + + +def test_kafka_chooser_rejects_a_fully_open_key_range(): + """An open range still means the topic is auto-partitioned. + + A single auto-partitioned partition reports ["", "") before its first split. Accepting it + here would work right up until that split, and then break: the children come back bounded + and modulo routing cannot place a key by bounds. + """ + chooser = PublicPartitionByKeyKafka() + with pytest.raises(ValueError): + chooser.add_partitions([_partition_info(0, from_bound=b"", to_bound=b"")]) + + +class _FakeMultiAsyncWriter: + """Stand-in for the async multi-writer that the sync facade drives.""" + + def __init__(self, driver=None, settings=None, _parent=None): + self.calls = [] + self.closed_with = None + + async def wait_init(self): + self.calls.append("wait_init") + + async def write(self, messages): + self.calls.append(("write", messages)) + + async def write_with_ack(self, messages): + self.calls.append(("write_with_ack", messages)) + return "ack" + + async def flush(self): + self.calls.append("flush") + + async def close(self, *, flush=True): + self.closed_with = flush + self.calls.append(("close", flush)) + + +class TestTopicWriterMultiSync: + """The sync facade owns no logic of its own: it must hand every call to the async writer + on the shared loop and report closure honestly. These check that nothing is dropped or + silently reordered on the way across the thread boundary.""" + + @pytest.fixture + def writer(self, background_loop, monkeypatch): + monkeypatch.setattr( + "ydb._topic_writer.topic_writer_multi_sync.TopicWriterMultiAsyncIO", + _FakeMultiAsyncWriter, + ) + settings = MultiWriterSettings(topic="/local/topic", producer_id_prefix="pfx") + writer = TopicWriterMultiSync(mock.Mock(), settings, eventloop=background_loop) + yield writer + if not writer._closed: + writer.close(flush=False) + + def test_calls_are_forwarded_to_the_async_writer(self, writer): + message = PublicMessage(data=b"hello", key="k") + + writer.wait_init() + writer.write(message) + assert writer.write_with_ack(message) == "ack" + writer.flush() + + kinds = [c if isinstance(c, str) else c[0] for c in writer._async_writer.calls] + assert kinds == ["wait_init", "write", "write_with_ack", "flush"] + + def test_async_variants_return_futures(self, writer): + assert writer.async_wait_init().result(timeout=5) is None + assert writer.async_write_with_ack(PublicMessage(data=b"x", key="k")).result(timeout=5) == "ack" + assert writer.async_flush().result(timeout=5) is None + + def test_write_accepts_a_batch(self, writer): + messages = [PublicMessage(data=b"a", key="k"), PublicMessage(data=b"b", key="k")] + writer.write(messages) + assert writer._async_writer.calls[-1][0] == "write" + + def test_close_is_idempotent_and_flushes_by_default(self, writer): + writer.close() + assert writer._async_writer.closed_with is True + + # A second close must not reach the async writer again. + calls_before = len(writer._async_writer.calls) + writer.close() + assert len(writer._async_writer.calls) == calls_before + + def test_every_entry_point_refuses_to_work_after_close(self, writer): + writer.close(flush=False) + + message = PublicMessage(data=b"x", key="k") + for call in ( + lambda: writer.write(message), + lambda: writer.write_with_ack(message), + lambda: writer.flush(), + lambda: writer.wait_init(), + lambda: writer.async_flush(), + lambda: writer.async_wait_init(), + lambda: writer.async_write_with_ack(message), + ): + with pytest.raises(TopicWriterClosedError): + call() + + def test_context_manager_closes_on_exit(self, background_loop, monkeypatch): + monkeypatch.setattr( + "ydb._topic_writer.topic_writer_multi_sync.TopicWriterMultiAsyncIO", + _FakeMultiAsyncWriter, + ) + settings = MultiWriterSettings(topic="/local/topic", producer_id_prefix="pfx") + with TopicWriterMultiSync(mock.Mock(), settings, eventloop=background_loop) as writer: + writer.write(PublicMessage(data=b"x", key="k")) + assert writer._closed + assert writer._async_writer.closed_with is True + + def test_unclosed_writer_is_closed_on_delete(self, background_loop, monkeypatch): + """__del__ is the last chance to release the streams a forgotten writer still holds.""" + monkeypatch.setattr( + "ydb._topic_writer.topic_writer_multi_sync.TopicWriterMultiAsyncIO", + _FakeMultiAsyncWriter, + ) + settings = MultiWriterSettings(topic="/local/topic", producer_id_prefix="pfx") + writer = TopicWriterMultiSync(mock.Mock(), settings, eventloop=background_loop) + inner = writer._async_writer + + writer.__del__() + + assert writer._closed + assert inner.closed_with is False # never flush from a destructor + + def test_close_failure_surfaces_unless_the_body_already_failed(self, background_loop, monkeypatch): + """A failing close must not mask the caller's own exception. + + Losing the body's exception to a teardown error is how a real failure gets reported as + something unrelated, so close errors are only raised when the block exited cleanly. + """ + + class FailingClose(_FakeMultiAsyncWriter): + async def close(self, *, flush=True): + raise RuntimeError("close failed") + + monkeypatch.setattr("ydb._topic_writer.topic_writer_multi_sync.TopicWriterMultiAsyncIO", FailingClose) + settings = MultiWriterSettings(topic="/local/topic", producer_id_prefix="pfx") + + with pytest.raises(RuntimeError, match="close failed"): + with TopicWriterMultiSync(mock.Mock(), settings, eventloop=background_loop): + pass + + class TestException(Exception): + pass + + with pytest.raises(TestException): + with TopicWriterMultiSync(mock.Mock(), settings, eventloop=background_loop): + raise TestException() + + def test_delete_swallows_close_failure(self, background_loop, monkeypatch): + """__del__ runs during garbage collection: raising there only produces noise.""" + + class FailingClose(_FakeMultiAsyncWriter): + async def close(self, *, flush=True): + raise RuntimeError("close failed") + + monkeypatch.setattr("ydb._topic_writer.topic_writer_multi_sync.TopicWriterMultiAsyncIO", FailingClose) + settings = MultiWriterSettings(topic="/local/topic", producer_id_prefix="pfx") + writer = TopicWriterMultiSync(mock.Mock(), settings, eventloop=background_loop) + + writer.__del__() # must not raise + + def test_falls_back_to_the_shared_event_loop(self, monkeypatch): + """Without an explicit loop the writer joins the client's shared one.""" + monkeypatch.setattr("ydb._topic_writer.topic_writer_multi_sync.TopicWriterMultiAsyncIO", _FakeMultiAsyncWriter) + loop = mock.Mock() + shared = mock.Mock(return_value=loop) + monkeypatch.setattr("ydb._topic_writer.topic_writer_multi_sync._get_shared_event_loop", shared) + caller = mock.Mock() + monkeypatch.setattr("ydb._topic_writer.topic_writer_multi_sync.CallFromSyncToAsync", caller) + + settings = MultiWriterSettings(topic="/local/topic", producer_id_prefix="pfx") + TopicWriterMultiSync(mock.Mock(), settings) + + shared.assert_called_once() + caller.assert_called_once_with(loop) diff --git a/ydb/topic.py b/ydb/topic.py index 988592937..8969cb839 100644 --- a/ydb/topic.py +++ b/ydb/topic.py @@ -27,6 +27,12 @@ "TopicWriteResult", "TopicWriter", "TopicWriterAsyncIO", + "TopicWriterMulti", + "TopicWriterMultiAsyncIO", + "TopicWriterMultiSettings", + "TopicWriterPartitionChooser", + "TopicWriterPartitionByKeyKafka", + "TopicWriterPartitionByKeyBound", "TopicTxWriter", "TopicTxWriterAsyncIO", "TopicWriterInitInfo", @@ -81,6 +87,17 @@ from ._topic_writer.topic_writer_sync import WriterSync as TopicWriter from ._topic_writer.topic_writer_sync import TxWriterSync as TopicTxWriter +from ._topic_writer.topic_writer_multi_asyncio import ( # noqa: F401 + MultiWriterSettings as TopicWriterMultiSettings, + TopicWriterMultiAsyncIO, +) +from ._topic_writer.topic_writer_multi_sync import TopicWriterMultiSync as TopicWriterMulti +from ._topic_writer.topic_writer_partition_chooser import ( # noqa: F401 + PublicPartitionChooser as TopicWriterPartitionChooser, + PublicPartitionByKeyKafka as TopicWriterPartitionByKeyKafka, + PublicPartitionByKeyBound as TopicWriterPartitionByKeyBound, +) + from ._topic_common.common import ( wrap_operation as _wrap_operation, create_result_wrapper as _create_result_wrapper, @@ -335,7 +352,7 @@ def writer( partition_id: Union[int, None] = None, auto_seqno: bool = True, auto_created_at: bool = True, - codec: Optional[TopicCodec] = None, # default mean auto-select + codec: Optional[TopicCodec] = None, # default means auto-select # encoders: map[codec_code] func(encoded_bytes)->decoded_bytes # the func will be called from multiply threads in parallel. encoders: Optional[Mapping[_ydb_topic_public_types.PublicCodec, Callable[[bytes], bytes]]] = None, @@ -367,7 +384,7 @@ def tx_writer( partition_id: Union[int, None] = None, auto_seqno: bool = True, auto_created_at: bool = True, - codec: Optional[TopicCodec] = None, # default mean auto-select + codec: Optional[TopicCodec] = None, # default means auto-select # encoders: map[codec_code] func(encoded_bytes)->decoded_bytes # the func will be called from multiply threads in parallel. encoders: Optional[Mapping[_ydb_topic_public_types.PublicCodec, Callable[[bytes], bytes]]] = None, @@ -390,6 +407,36 @@ def tx_writer( return TopicTxWriterAsyncIO(tx=tx, driver=self._driver, settings=settings, _client=self) + def multiwriter( + self, + topic, + *, + producer_id_prefix: Optional[str] = None, # default - random + partition_chooser: Optional[TopicWriterPartitionChooser] = None, # default - route by key (Kafka hash) + auto_seqno: bool = True, + auto_created_at: bool = True, + codec: Optional[TopicCodec] = None, # default means auto-select + encoders: Optional[Mapping[_ydb_topic_public_types.PublicCodec, Callable[[bytes], bytes]]] = None, + encoder_executor: Optional[concurrent.futures.Executor] = None, # default shared client executor pool + max_buffer_size_bytes: Optional[int] = None, + max_buffer_messages: Optional[int] = None, + buffer_wait_timeout_sec: Optional[float] = None, + # Close idle per-partition sub-writers after this many seconds (recreated on demand). + # None = default; <= 0 disables eviction. + writer_idle_timeout_sec: Optional[float] = None, + ) -> TopicWriterMultiAsyncIO: + logger.debug("Create multi-writer for topic=%s", topic) + args = locals().copy() + del args["self"] + self._check_closed() + + settings = TopicWriterMultiSettings(**args) + + if not settings.encoder_executor: + settings.encoder_executor = self._executor + + return TopicWriterMultiAsyncIO(self._driver, settings, _parent=self) + @ydb_retry(retry_cancelled=True, idempotent=True) async def commit_offset( self, path: str, consumer: str, partition_id: int, offset: int, read_session_id: Optional[str] = None @@ -453,7 +500,7 @@ def __init__(self, driver: driver.Driver, settings: Optional[TopicClientSettings def __del__(self): if not self._closed: try: - logger.warning("Topic client was not closed properly. Consider using method close().") + logger.debug("Topic client was not closed properly. Consider using method close().") self.close() except BaseException: logger.warning("Something went wrong during topic client close in __del__") @@ -669,7 +716,7 @@ def writer( partition_id: Union[int, None] = None, auto_seqno: bool = True, auto_created_at: bool = True, - codec: Optional[TopicCodec] = None, # default mean auto-select + codec: Optional[TopicCodec] = None, # default means auto-select # encoders: map[codec_code] func(encoded_bytes)->decoded_bytes # the func will be called from multiply threads in parallel. encoders: Optional[Mapping[_ydb_topic_public_types.PublicCodec, Callable[[bytes], bytes]]] = None, @@ -702,7 +749,7 @@ def tx_writer( partition_id: Union[int, None] = None, auto_seqno: bool = True, auto_created_at: bool = True, - codec: Optional[TopicCodec] = None, # default mean auto-select + codec: Optional[TopicCodec] = None, # default means auto-select # encoders: map[codec_code] func(encoded_bytes)->decoded_bytes # the func will be called from multiply threads in parallel. encoders: Optional[Mapping[_ydb_topic_public_types.PublicCodec, Callable[[bytes], bytes]]] = None, @@ -726,6 +773,36 @@ def tx_writer( return TopicTxWriter(tx, self._driver, settings, _parent=self) + def multiwriter( + self, + topic, + *, + producer_id_prefix: Optional[str] = None, # default - random + partition_chooser: Optional[TopicWriterPartitionChooser] = None, # default - route by key (Kafka hash) + auto_seqno: bool = True, + auto_created_at: bool = True, + codec: Optional[TopicCodec] = None, # default means auto-select + encoders: Optional[Mapping[_ydb_topic_public_types.PublicCodec, Callable[[bytes], bytes]]] = None, + encoder_executor: Optional[concurrent.futures.Executor] = None, # default shared client executor pool + max_buffer_size_bytes: Optional[int] = None, + max_buffer_messages: Optional[int] = None, + buffer_wait_timeout_sec: Optional[float] = None, + # Close idle per-partition sub-writers after this many seconds (recreated on demand). + # None = default; <= 0 disables eviction. + writer_idle_timeout_sec: Optional[float] = None, + ) -> TopicWriterMulti: + logger.debug("Create multi-writer for topic=%s", topic) + args = locals().copy() + del args["self"] + self._check_closed() + + settings = TopicWriterMultiSettings(**args) + + if not settings.encoder_executor: + settings.encoder_executor = self._executor + + return TopicWriterMulti(self._driver, settings, _parent=self) + @ydb_retry(retry_cancelled=True, idempotent=True) def commit_offset( self, path: str, consumer: str, partition_id: int, offset: int, read_session_id: Optional[str] = None diff --git a/ydb/topic_test.py b/ydb/topic_test.py new file mode 100644 index 000000000..a1533fa51 --- /dev/null +++ b/ydb/topic_test.py @@ -0,0 +1,115 @@ +from unittest import mock + +import pytest + +from . import issues, topic + + +class _CapturedWriter: + """Records what the client factory built, so the test can check the wiring only.""" + + last = None + + def __init__(self, driver, settings, _parent=None): + self.driver = driver + self.settings = settings + self.parent = _parent + type(self).last = self + + +@pytest.fixture +def async_client(): + client = topic.TopicClientAsyncIO(mock.Mock()) + yield client + client._closed = True # skip the real close: nothing was actually opened + + +@pytest.fixture +def sync_client(): + # The sync client takes settings positionally, unlike the async one. + client = topic.TopicClient(mock.Mock(), None) + yield client + client._closed = True + + +@pytest.mark.parametrize( + "client_fixture, patched", [("async_client", "TopicWriterMultiAsyncIO"), ("sync_client", "TopicWriterMulti")] +) +def test_multiwriter_factory_passes_settings_through(client_fixture, patched, request, monkeypatch): + """The factory is the only place the caller's arguments turn into writer settings. + + A silently dropped argument here (a codec, a chooser, a buffer limit) would leave the writer + running on defaults while the caller believes otherwise, so every one of them is checked. + """ + client = request.getfixturevalue(client_fixture) + monkeypatch.setattr(topic, patched, _CapturedWriter) + + chooser = topic.TopicWriterPartitionByKeyKafka() + writer = client.multiwriter( + "/local/topic", + producer_id_prefix="pfx", + partition_chooser=chooser, + auto_seqno=False, + auto_created_at=False, + codec=topic.TopicCodec.RAW, + max_buffer_size_bytes=1024, + max_buffer_messages=10, + buffer_wait_timeout_sec=1.5, + writer_idle_timeout_sec=30, + ) + + assert writer is _CapturedWriter.last + assert writer.parent is client, "the writer must keep the client alive" + + settings = writer.settings + assert settings.topic == "/local/topic" + assert settings.producer_id_prefix == "pfx" + assert settings.partition_chooser is chooser + assert settings.auto_seqno is False + assert settings.auto_created_at is False + assert settings.codec == topic.TopicCodec.RAW + assert settings.max_buffer_size_bytes == 1024 + assert settings.max_buffer_messages == 10 + assert settings.buffer_wait_timeout_sec == 1.5 + assert settings.writer_idle_timeout_sec == 30 + # Encoding runs on the client's shared pool unless the caller brought its own. + assert settings.encoder_executor is client._executor + + +@pytest.mark.parametrize( + "client_fixture, patched", [("async_client", "TopicWriterMultiAsyncIO"), ("sync_client", "TopicWriterMulti")] +) +def test_multiwriter_keeps_a_caller_supplied_executor(client_fixture, patched, request, monkeypatch): + client = request.getfixturevalue(client_fixture) + monkeypatch.setattr(topic, patched, _CapturedWriter) + executor = mock.Mock() + + writer = client.multiwriter("/local/topic", encoder_executor=executor) + + assert writer.settings.encoder_executor is executor + + +@pytest.mark.parametrize("client_fixture", ["async_client", "sync_client"]) +def test_multiwriter_refuses_a_closed_client(client_fixture, request): + client = request.getfixturevalue(client_fixture) + client._closed = True + + with pytest.raises(issues.Error): + client.multiwriter("/local/topic") + + +def test_unclosed_sync_client_is_closed_on_delete(): + """__del__ is the last chance to release the executor a forgotten client still holds.""" + client = topic.TopicClient(mock.Mock(), None) + with mock.patch.object(topic.TopicClient, "close") as close: + client.__del__() + close.assert_called_once() + client._closed = True + + +def test_delete_of_a_closed_sync_client_does_nothing(): + client = topic.TopicClient(mock.Mock(), None) + client._closed = True + with mock.patch.object(topic.TopicClient, "close") as close: + client.__del__() + close.assert_not_called()