-
Notifications
You must be signed in to change notification settings - Fork 8
Expand file tree
/
Copy pathwrapper.py
More file actions
307 lines (273 loc) · 11.5 KB
/
Copy pathwrapper.py
File metadata and controls
307 lines (273 loc) · 11.5 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
from __future__ import annotations
import logging
import time
from collections.abc import Callable
from typing import Any
import workflows.recipe
logger = logging.getLogger("workflows.recipe.wrapper")
class RecipeWrapper:
"""A wrapper object which contains a recipe and a number of functions to make
life easier for recipe users.
"""
def __init__(
self, message=None, transport=None, recipe=None, environment=None, **kwargs
):
"""Create a RecipeWrapper object from a wrapped message.
References to the transport layer are required to send directly to
connected downstream processes.
"""
if message:
self.recipe = workflows.recipe.Recipe(message["recipe"])
self.recipe_pointer = int(message["recipe-pointer"])
self.recipe_step = self.recipe[self.recipe_pointer]
self.recipe_path = message.get("recipe-path", [])
if environment is None:
self.environment = message.get("environment", {})
else:
self.environment = environment
self.payload = message.get("payload")
elif recipe:
if isinstance(recipe, workflows.recipe.Recipe):
self.recipe = recipe
else:
self.recipe = workflows.recipe.Recipe(recipe)
self.recipe_pointer = None
self.recipe_step = None
self.recipe_path = []
self.environment = environment or {}
self.payload = None
else:
raise ValueError(
"A message or recipe is required to create a RecipeWrapper object."
)
self.default_channel = None
self.transport = transport
def send(self, *args, **kwargs):
"""Send messages to another service that is connected to the currently
running service via the recipe. The 'send' method will either use a
default channel name, set via the set_default_channel method, or an
unnamed output definition.
"""
if not self.transport:
raise ValueError(
"This RecipeWrapper object does not contain "
"a reference to a transport object."
)
if not self.recipe_step:
raise ValueError(
"This RecipeWrapper object does not contain "
"a recipe with a selected step."
)
if "output" not in self.recipe_step:
# The current recipe step does not have output channels.
return
if isinstance(self.recipe_step["output"], dict):
# The current recipe step does have named output channels.
if self.default_channel:
# Use named output channel
self.send_to(self.default_channel, *args, **kwargs)
else:
# The current recipe step does have unnamed output channels.
self._send_to_destinations(self.recipe_step["output"], *args, **kwargs)
def send_to(self, channel, *args, **kwargs):
"""Send messages to another service that is connected to the currently
running service via the recipe. Discard messages if the recipe does
not have anything connected to the specified output channel.
"""
if not self.transport:
raise ValueError(
"This RecipeWrapper object does not contain "
"a reference to a transport object."
)
if not self.recipe_step:
raise ValueError(
"This RecipeWrapper object does not contain "
"a recipe with a selected step."
)
if "output" not in self.recipe_step:
# The current recipe step does not have output channels.
return
if not isinstance(self.recipe_step["output"], dict):
# The current recipe step does not have named output channels.
if self.default_channel == channel:
# Use unnamed output channels
self.send(*args, **kwargs)
return
if channel not in self.recipe_step["output"]:
# The current recipe step does not have an output channel with this name.
return
self._send_to_destinations(self.recipe_step["output"][channel], *args, **kwargs)
def set_default_channel(self, channel):
"""Define one named output channel to be equivalent to unnamed output
channels. For this channel send() and send_to() will be identical."""
self.default_channel = channel
def start(self, header=None, **kwargs):
"""Trigger the start of a recipe, sending the defined payloads to the
recipients set in the recipe. Any parameters to this function are
passed to the transport send/broadcast methods.
If the wrapped recipe has already been started then a ValueError will
be raised.
"""
if not self.transport:
raise ValueError(
"This RecipeWrapper object does not contain "
"a reference to a transport object."
)
if self.recipe_step:
raise ValueError("This recipe has already been started.")
for destination, payload in self.recipe["start"]:
self._send_to_destination(destination, header, payload, kwargs)
def checkpoint(self, message, header=None, delay=0, **kwargs):
"""Send a message to the current recipe destination. This can be used to
keep a state for longer processing tasks.
:param delay: Delay transport of message by this many seconds
"""
if not self.transport:
raise ValueError(
"This RecipeWrapper object does not contain "
"a reference to a transport object."
)
if not self.recipe_step:
raise ValueError(
"This RecipeWrapper object does not contain "
"a recipe with a selected step."
)
kwargs["delay"] = delay
self._send_to_destination(
self.recipe_pointer, header, message, kwargs, add_path_step=False
)
def apply_parameters(self, parameters):
"""Recursively apply parameter replacement (see recipe.py) to the wrapped
recipe, updating internal references afterwards.
While this operation is useful for testing it should not be used in
production. Replacing parameters means that the recipe changes as it is
passed down the chain of services. This makes debugging very difficult.
"""
self.recipe.apply_parameters(parameters)
self.recipe_step = self.recipe[self.recipe_pointer]
def _generate_full_recipe_message(self, destination, message, add_path_step):
"""Factory function to generate independent message objects for
downstream recipients with different destinations."""
if add_path_step and self.recipe_pointer:
recipe_path = self.recipe_path + [self.recipe_pointer]
else:
recipe_path = self.recipe_path
return {
"environment": self.environment,
"payload": message,
"recipe": self.recipe.recipe,
"recipe-path": recipe_path,
"recipe-pointer": destination,
}
def _send_to_destinations(
self,
destinations,
message,
header=None,
mangle_for_sending: Callable[[Any], Any] | None = None,
**kwargs,
):
"""Send messages to a list of numbered destinations. This is an internal
helper method used by the public 'send' methods.
"""
if not isinstance(destinations, list):
destinations = (destinations,)
for destination in destinations:
self._send_to_destination(
destination,
header,
message,
kwargs,
mangle_for_sending=mangle_for_sending,
)
def _send_to_destination(
self,
destination,
header,
payload,
transport_kwargs,
add_path_step=True,
mangle_for_sending: Callable[[Any], Any] | None = None,
):
"""Helper function to send a message to a specific recipe destination."""
if header:
header = header.copy()
header["workflows-recipe"] = True
else:
header = {"workflows-recipe": True}
dest_kwargs = transport_kwargs.copy()
if (
"transport-delay" in self.recipe[destination]
and "delay" not in transport_kwargs
):
dest_kwargs["delay"] = self.recipe[destination]["transport-delay"]
if "exchange" in self.recipe[destination]:
dest_kwargs.setdefault("exchange", self.recipe[destination]["exchange"])
if self.recipe[destination].get("queue"):
send = (
self.transport.raw_send if mangle_for_sending else self.transport.send
)
message = self._generate_full_recipe_message(
destination, payload, add_path_step
)
if mangle_for_sending:
message = mangle_for_sending(message)
self._retry_transport(
send,
self.recipe[destination]["queue"],
message,
headers=header,
**dest_kwargs,
)
if self.recipe[destination].get("topic"):
broadcast = (
self.transport.raw_broadcast
if mangle_for_sending
else self.transport.broadcast
)
message = self._generate_full_recipe_message(
destination, payload, add_path_step
)
if mangle_for_sending:
message = mangle_for_sending(message)
self._retry_transport(
broadcast,
self.recipe[destination]["exchange"],
message,
headers=header,
topic=self.recipe[destination].get("topic"),
**dest_kwargs,
)
def _retry_transport(self, function, *args, **kwargs):
"""Attempt to send a message, and in case the connection has been lost,
attempt to reconnect. Reconnecting only works on the assumption that
the previous connection did not include any subscriptions, which should
be true for wrappers."""
attempt = 0
limit = 12
initial_failure = None
while True:
try:
if initial_failure:
logger.info("Reconnection attempt %d of %d", attempt, limit)
self.transport.connect()
return function(*args, **kwargs)
except workflows.Disconnected as e:
if not self.transport.is_reconnectable:
raise
attempt = attempt + 1
if not initial_failure:
initial_failure = e
if attempt > limit:
logger.error("Transport connection failure", exc_info=True)
raise initial_failure from None
delay = attempt * attempt * 3
# wait 3, 12, 27, 48, 75, 108, 147, 192, 243, 300, 363, 432 seconds
# between attempts for a total maximum of 32.5 minutes
logger.warning(
"Connection failure detected during send attempt."
" Retrying in %d seconds",
delay,
exc_info=True,
)
time.sleep(delay)