|
| 1 | +"""Shared HIMPORT wire-execution helpers for the synchronous clients. |
| 2 | +
|
| 3 | +The PREPARE / SET / DISCARD packed-write drain loops, per-connection version |
| 4 | +bookkeeping, and ``NoSuchFieldsetError`` re-prepare-and-retry are identical for |
| 5 | +the standalone (:class:`redis.Redis`) and cluster (:class:`redis.RedisCluster`) |
| 6 | +sync clients, differing only in (a) which object provides ``parse_response`` and |
| 7 | +(b) the cluster-only ASK-redirect handling. These functions take that object as |
| 8 | +``node`` and an ``asking`` flag (``False`` -- and a no-op -- for standalone), so |
| 9 | +both clients share one implementation instead of copy-pasting the logic. The |
| 10 | +per-class ``_himport_*`` methods are thin delegators to these functions. |
| 11 | +
|
| 12 | +The async mirror lives in :mod:`redis.asyncio._himport_exec`; the two are kept |
| 13 | +separate on purpose (the project maintains parallel sync/async stacks by hand). |
| 14 | +""" |
| 15 | + |
| 16 | +from redis.exceptions import NoSuchFieldsetError, ResponseError |
| 17 | +from redis.himport import ( |
| 18 | + HIMPORT_DISCARD, |
| 19 | + HIMPORT_PREPARE, |
| 20 | + HIMPORT_SET, |
| 21 | + HImportRegistry, |
| 22 | + himport_discard_command, |
| 23 | + himport_prepare_command, |
| 24 | + himport_set_command, |
| 25 | + parse_himport_set_args, |
| 26 | +) |
| 27 | + |
| 28 | + |
| 29 | +def reconcile_discards(node, conn): |
| 30 | + """DISCARD, on ``conn``, any prepared fieldset removed from the registry. |
| 31 | +
|
| 32 | + Runs at most once per registry mutation: the connection records the registry |
| 33 | + ``revision`` it last reconciled against, so unchanged registries are a no-op. |
| 34 | + ``node`` supplies ``parse_response`` (the standalone client itself, or the |
| 35 | + owning node's client in cluster mode). |
| 36 | + """ |
| 37 | + registry = conn.himport_registry |
| 38 | + if registry is None or conn._himport_reconciled_revision == registry.revision: |
| 39 | + return |
| 40 | + # Snapshot the revision *before* computing ``stale`` so the value stamped at |
| 41 | + # the end is never newer than the registry state ``stale`` reflects. A |
| 42 | + # concurrent ``himport_discard`` (thread-shared sync client) that lands after |
| 43 | + # this point only leaves the connection marked behind the live revision, so |
| 44 | + # the next reconcile re-runs and catches it. Re-reading ``registry.revision`` |
| 45 | + # at the end instead would stamp a discard this connection never sent. |
| 46 | + reconciled_to = registry.revision |
| 47 | + stale = registry.names_to_discard(list(conn._himport_prepared)) |
| 48 | + if stale: |
| 49 | + conn.send_packed_command( |
| 50 | + conn.pack_commands([himport_discard_command(n) for n in stale]) |
| 51 | + ) |
| 52 | + # One reply per packed DISCARD must be read regardless of a per-command |
| 53 | + # ResponseError, otherwise the unread replies desync the pooled socket. |
| 54 | + # Drain every reply, then surface the first error (ConnectionError is not |
| 55 | + # caught: it tears the socket down, so no desync is possible). |
| 56 | + first_error = None |
| 57 | + for n in stale: |
| 58 | + try: |
| 59 | + node.parse_response(conn, HIMPORT_DISCARD) |
| 60 | + except ResponseError as e: |
| 61 | + first_error = first_error or e |
| 62 | + conn._himport_prepared.pop(n, None) |
| 63 | + if first_error is not None: |
| 64 | + raise first_error |
| 65 | + conn._himport_reconciled_revision = reconciled_to |
| 66 | + |
| 67 | + |
| 68 | +def prepare_and_set(node, conn, key, fieldset_name, values, fieldset, asking=False): |
| 69 | + """PREPARE ``fieldset`` bundled with the SET on ``conn`` (one packed write). |
| 70 | +
|
| 71 | + When ``asking`` is set (an ASK-redirected cluster SET) the batch becomes |
| 72 | + ``[PREPARE, ASKING, SET]`` so the per-command ASKING allowance falls |
| 73 | + immediately before the SET -- the only slot-scoped command. PREPARE is a |
| 74 | + connection-session command the ASKING flag does not gate, so placing it |
| 75 | + before ASKING is safe. Every reply is drained even on a per-command error so |
| 76 | + the packed replies never desync the pooled socket. |
| 77 | + """ |
| 78 | + commands = [himport_prepare_command(fieldset_name, fieldset.fields)] |
| 79 | + if asking: |
| 80 | + commands.append(("ASKING",)) |
| 81 | + commands.append(himport_set_command(key, fieldset_name, values)) |
| 82 | + conn.send_packed_command(conn.pack_commands(commands)) |
| 83 | + prep_error = ask_error = set_error = None |
| 84 | + set_resp = None |
| 85 | + try: |
| 86 | + node.parse_response(conn, HIMPORT_PREPARE) |
| 87 | + except ResponseError as e: |
| 88 | + prep_error = e |
| 89 | + if asking: |
| 90 | + try: |
| 91 | + node.parse_response(conn, "ASKING") |
| 92 | + except ResponseError as e: |
| 93 | + ask_error = e |
| 94 | + try: |
| 95 | + set_resp = node.parse_response(conn, HIMPORT_SET) |
| 96 | + except ResponseError as e: |
| 97 | + set_error = e |
| 98 | + |
| 99 | + if prep_error: |
| 100 | + raise prep_error # PREPARE failure is the root cause |
| 101 | + else: |
| 102 | + conn._himport_prepared[fieldset_name] = fieldset.version |
| 103 | + |
| 104 | + if ask_error: |
| 105 | + raise ask_error |
| 106 | + if set_error: |
| 107 | + raise set_error |
| 108 | + return set_resp |
| 109 | + |
| 110 | + |
| 111 | +def execute_set(node, conn, key, fieldset_name, values, asking=False): |
| 112 | + """Execute an ``HIMPORT SET`` on ``conn`` with the required session setup. |
| 113 | +
|
| 114 | + Reconciles deferred discards, lazily bundles PREPARE with the SET on first |
| 115 | + use of a fieldset, and recovers once from a mid-connection fieldset loss |
| 116 | + (``NoSuchFieldsetError``) by re-PREPARE-and-retry. When ``asking`` is set the |
| 117 | + ASKING allowance is folded into the SET's own packed write so it immediately |
| 118 | + precedes the (slot-scoped) SET; the session setup runs first, since those are |
| 119 | + connection-session commands the flag does not gate. |
| 120 | + """ |
| 121 | + reconcile_discards(node, conn) |
| 122 | + |
| 123 | + registry = conn.himport_registry |
| 124 | + fieldset = registry.get(fieldset_name) if registry is not None else None |
| 125 | + # Lazy PREPARE bundled with SET on first use of this fieldset. |
| 126 | + if ( |
| 127 | + fieldset is not None |
| 128 | + and conn._himport_prepared.get(fieldset_name) != fieldset.version |
| 129 | + ): |
| 130 | + return prepare_and_set( |
| 131 | + node, conn, key, fieldset_name, values, fieldset, asking=asking |
| 132 | + ) |
| 133 | + |
| 134 | + # Believed already prepared (or an unregistered fieldset): bare SET, with |
| 135 | + # ASKING packed immediately before it when this is an ASK redirect. |
| 136 | + if asking: |
| 137 | + conn.send_packed_command( |
| 138 | + conn.pack_commands( |
| 139 | + [("ASKING",), himport_set_command(key, fieldset_name, values)] |
| 140 | + ) |
| 141 | + ) |
| 142 | + try: |
| 143 | + node.parse_response(conn, "ASKING") |
| 144 | + except ResponseError as ask_error: |
| 145 | + # ASKING and SET were one packed write, so the SET reply is still |
| 146 | + # queued. Drain it before surfacing the ASKING error, otherwise the |
| 147 | + # connection returns to the pool with an unread reply and desyncs the |
| 148 | + # next borrower. |
| 149 | + try: |
| 150 | + node.parse_response(conn, HIMPORT_SET) |
| 151 | + except ResponseError: |
| 152 | + pass |
| 153 | + raise ask_error |
| 154 | + else: |
| 155 | + conn.send_command(*himport_set_command(key, fieldset_name, values)) |
| 156 | + try: |
| 157 | + return node.parse_response(conn, HIMPORT_SET) |
| 158 | + except NoSuchFieldsetError: |
| 159 | + # Server dropped the fieldset mid-connection without dropping the socket |
| 160 | + # (e.g. RESET / maxmemory-clients eviction): re-PREPARE on this healthy |
| 161 | + # connection and retry the SET once rather than reconnecting. Only for |
| 162 | + # registry-backed fieldsets; manual/unregistered usage propagates. |
| 163 | + if fieldset is None: |
| 164 | + raise |
| 165 | + conn._himport_prepared.pop(fieldset_name, None) |
| 166 | + return prepare_and_set( |
| 167 | + node, conn, key, fieldset_name, values, fieldset, asking=asking |
| 168 | + ) |
| 169 | + |
| 170 | + |
| 171 | +def prepare_pipeline(node, conn, command_arg_lists): |
| 172 | + """Pre-flight ``conn`` for a pipeline batch containing ``HIMPORT SET``s. |
| 173 | +
|
| 174 | + The packed pipeline write bypasses the per-command lazy-PREPARE path, so the |
| 175 | + fieldsets referenced by the buffered SETs must be PREPAREd on ``conn`` first. |
| 176 | + Reconciles deferred discards, then PREPAREs every distinct registered fieldset |
| 177 | + the batch references that this connection has not already prepared, in one |
| 178 | + packed write. ``command_arg_lists`` is the batch's per-command positional-arg |
| 179 | + sequences (the caller extracts them from its own command representation). |
| 180 | + No-op when the batch has no registry-backed ``HIMPORT SET``. |
| 181 | + """ |
| 182 | + # Selection (registry check, deferred-discard reconcile, scan/dedup/version) |
| 183 | + # is shared with pipeline_prepares. This path differs only in that it sends the |
| 184 | + # PREPAREs as their own packed exchange -- rather than folding them into a |
| 185 | + # queued write -- then drains their replies and raises the first error. |
| 186 | + to_prepare = pipeline_prepares(node, conn, command_arg_lists) |
| 187 | + if not to_prepare: |
| 188 | + return |
| 189 | + conn.send_packed_command(conn.pack_commands(prepare_wire_commands(to_prepare))) |
| 190 | + # Every reply must be drained even on a per-command error, or the unread |
| 191 | + # replies desync the socket before the buffered batch is sent; then raise. |
| 192 | + first_error = drain_pipeline_prepares(node, conn, to_prepare) |
| 193 | + if first_error is not None: |
| 194 | + raise first_error |
| 195 | + |
| 196 | + |
| 197 | +def pipeline_prepares(node, conn, command_arg_lists): |
| 198 | + """Return the fieldsets that must be PREPAREd on ``conn`` for this batch. |
| 199 | +
|
| 200 | + Like :func:`prepare_pipeline`, but does **not** send the PREPAREs: the caller |
| 201 | + folds them into the same packed write as the queued commands (see the pipeline |
| 202 | + executors), so the first pipeline use of a fieldset on a fresh or reconnected |
| 203 | + connection stays a single round trip instead of a separate PREPARE exchange |
| 204 | + followed by the batch. Deferred-discard reconciliation is still performed here, |
| 205 | + but it only touches the socket when discards are actually pending (rare); the |
| 206 | + common warm-up cost -- the first-use PREPARE -- is what gets folded. Returns an |
| 207 | + empty list when the batch references no not-yet-prepared registered fieldset, |
| 208 | + or when ``conn`` carries no real HIMPORT registry. |
| 209 | + """ |
| 210 | + registry = getattr(conn, "himport_registry", None) |
| 211 | + if not isinstance(registry, HImportRegistry): |
| 212 | + return [] |
| 213 | + reconcile_discards(node, conn) |
| 214 | + to_prepare = [] |
| 215 | + seen = set() |
| 216 | + for args in command_arg_lists: |
| 217 | + parsed = parse_himport_set_args(args) |
| 218 | + if parsed is None: |
| 219 | + continue |
| 220 | + fieldset_name = parsed[1] |
| 221 | + if fieldset_name in seen: |
| 222 | + continue |
| 223 | + seen.add(fieldset_name) |
| 224 | + fieldset = registry.get(fieldset_name) |
| 225 | + if ( |
| 226 | + fieldset is not None |
| 227 | + and conn._himport_prepared.get(fieldset_name) != fieldset.version |
| 228 | + ): |
| 229 | + to_prepare.append(fieldset) |
| 230 | + return to_prepare |
| 231 | + |
| 232 | + |
| 233 | +def prepare_wire_commands(fieldsets): |
| 234 | + """The leading ``HIMPORT PREPARE`` wire commands the caller folds into a batch.""" |
| 235 | + return [himport_prepare_command(fs.name, fs.fields) for fs in fieldsets] |
| 236 | + |
| 237 | + |
| 238 | +def drain_pipeline_prepares(node, conn, fieldsets): |
| 239 | + """Drain the ``len(fieldsets)`` leading PREPARE replies of a folded pipeline |
| 240 | + write, marking each fieldset prepared on success. |
| 241 | +
|
| 242 | + Returns the first ``ResponseError`` (or ``None``). The caller must still drain |
| 243 | + the queued command replies and only then surface this error: every reply on |
| 244 | + the wire has to be read before raising, or the pooled socket desyncs. |
| 245 | + """ |
| 246 | + first_error = None |
| 247 | + for fs in fieldsets: |
| 248 | + try: |
| 249 | + node.parse_response(conn, HIMPORT_PREPARE) |
| 250 | + except ResponseError as e: |
| 251 | + first_error = first_error or e |
| 252 | + continue |
| 253 | + conn._himport_prepared[fs.name] = fs.version |
| 254 | + return first_error |
0 commit comments