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