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
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
|
# Copyright 2022 The Pigweed Authors
#
# Licensed under the Apache License, Version 2.0 (the "License"); you may not
# use this file except in compliance with the License. You may obtain a copy of
# the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
# License for the specific language governing permissions and limitations under
# the License.
"""Client for the pw_transfer service, which transmits data over pw_rpc."""
import asyncio
import logging
import threading
from typing import Any, Dict, Optional, Union
from pw_rpc.callback_client import BidirectionalStreamingCall
from pw_status import Status
from pw_transfer.transfer import (
ProgressCallback,
ProtocolVersion,
ReadTransfer,
Transfer,
WriteTransfer,
)
from pw_transfer.chunk import Chunk
try:
from pw_transfer import transfer_pb2
except ImportError:
# For the bazel build, which puts generated protos in a different location.
from pigweed.pw_transfer import transfer_pb2 # type: ignore
_LOG = logging.getLogger(__package__)
_TransferDict = Dict[int, Transfer]
class Manager: # pylint: disable=too-many-instance-attributes
"""A manager for transmitting data through an RPC TransferService.
This should be initialized with an active Manager over an RPC channel. Only
one instance of this class should exist for a configured RPC TransferService
-- the Manager supports multiple simultaneous transfers.
When created, a Manager starts a separate thread in which transfer
communications and events are handled.
"""
def __init__(
self,
rpc_transfer_service,
*,
default_response_timeout_s: float = 2.0,
initial_response_timeout_s: float = 4.0,
max_retries: int = 3,
max_lifetime_retries: int = 1500,
default_protocol_version=ProtocolVersion.LATEST,
):
"""Initializes a Manager on top of a TransferService.
Args:
rpc_transfer_service: the pw_rpc transfer service client
default_response_timeout_s: max time to wait between receiving packets
initial_response_timeout_s: timeout for the first packet; may be
longer to account for transfer handler initialization
max_retires: number of times to retry a single package after a timeout
max_lifetime_retires: Cumulative maximum number of times to retry over
the course of the transfer before giving up.
"""
self._service: Any = rpc_transfer_service
self._default_response_timeout_s = default_response_timeout_s
self._initial_response_timeout_s = initial_response_timeout_s
self.max_retries = max_retries
self.max_lifetime_retries = max_lifetime_retries
self._default_protocol_version = default_protocol_version
# Ongoing transfers in the service by resource ID.
self._read_transfers: _TransferDict = {}
self._write_transfers: _TransferDict = {}
# RPC streams for read and write transfers. These are shareable by
# multiple transfers of the same type.
self._read_stream: Optional[BidirectionalStreamingCall] = None
self._write_stream: Optional[BidirectionalStreamingCall] = None
self._loop = asyncio.new_event_loop()
# Set the event loop for the current thread.
asyncio.set_event_loop(self._loop)
# Queues are used for communication between the Manager context and the
# dedicated asyncio transfer thread.
self._new_transfer_queue: asyncio.Queue = asyncio.Queue()
self._read_chunk_queue: asyncio.Queue = asyncio.Queue()
self._write_chunk_queue: asyncio.Queue = asyncio.Queue()
self._quit_event = asyncio.Event()
self._thread = threading.Thread(
target=self._start_event_loop_thread, daemon=True
)
self._thread.start()
def __del__(self):
# Notify the thread that the transfer manager is being destroyed and
# wait for it to exit.
if self._thread.is_alive():
self._loop.call_soon_threadsafe(self._quit_event.set)
self._thread.join()
def read(
self,
resource_id: int,
progress_callback: Optional[ProgressCallback] = None,
protocol_version: Optional[ProtocolVersion] = None,
) -> bytes:
"""Receives ("downloads") data from the server.
Args:
resource_id: ID of the resource from which to read.
progress_callback: Optional callback periodically invoked throughout
the transfer with the transfer state. Can be used to provide user-
facing status updates such as progress bars.
Raises:
Error: the transfer failed to complete
"""
if resource_id in self._read_transfers:
raise ValueError(
f'Read transfer for resource {resource_id} already exists'
)
if protocol_version is None:
protocol_version = self._default_protocol_version
transfer = ReadTransfer(
resource_id,
self._send_read_chunk,
self._end_read_transfer,
self._default_response_timeout_s,
self._initial_response_timeout_s,
self.max_retries,
self.max_lifetime_retries,
protocol_version,
progress_callback=progress_callback,
)
self._start_read_transfer(transfer)
transfer.done.wait()
if not transfer.status.ok():
raise Error(transfer.resource_id, transfer.status)
return transfer.data
def write(
self,
resource_id: int,
data: Union[bytes, str],
progress_callback: Optional[ProgressCallback] = None,
protocol_version: Optional[ProtocolVersion] = None,
) -> None:
"""Transmits ("uploads") data to the server.
Args:
resource_id: ID of the resource to which to write.
data: Data to send to the server.
progress_callback: Optional callback periodically invoked throughout
the transfer with the transfer state. Can be used to provide user-
facing status updates such as progress bars.
Raises:
Error: the transfer failed to complete
"""
if isinstance(data, str):
data = data.encode()
if resource_id in self._write_transfers:
raise ValueError(
f'Write transfer for resource {resource_id} already exists'
)
if protocol_version is None:
protocol_version = self._default_protocol_version
transfer = WriteTransfer(
resource_id,
data,
self._send_write_chunk,
self._end_write_transfer,
self._default_response_timeout_s,
self._initial_response_timeout_s,
self.max_retries,
self.max_lifetime_retries,
protocol_version,
progress_callback=progress_callback,
)
self._start_write_transfer(transfer)
transfer.done.wait()
if not transfer.status.ok():
raise Error(transfer.resource_id, transfer.status)
def _send_read_chunk(self, chunk: Chunk) -> None:
assert self._read_stream is not None
self._read_stream.send(chunk.to_message())
def _send_write_chunk(self, chunk: Chunk) -> None:
assert self._write_stream is not None
self._write_stream.send(chunk.to_message())
def _start_event_loop_thread(self):
"""Entry point for event loop thread that starts an asyncio context."""
asyncio.set_event_loop(self._loop)
# Recreate the async communication channels in the context of the
# running event loop.
self._new_transfer_queue = asyncio.Queue()
self._read_chunk_queue = asyncio.Queue()
self._write_chunk_queue = asyncio.Queue()
self._quit_event = asyncio.Event()
self._loop.create_task(self._transfer_event_loop())
self._loop.run_forever()
async def _transfer_event_loop(self):
"""Main async event loop."""
exit_thread = self._loop.create_task(self._quit_event.wait())
new_transfer = self._loop.create_task(self._new_transfer_queue.get())
read_chunk = self._loop.create_task(self._read_chunk_queue.get())
write_chunk = self._loop.create_task(self._write_chunk_queue.get())
while not self._quit_event.is_set():
# Perform a select(2)-like wait for one of several events to occur.
done, _ = await asyncio.wait(
(exit_thread, new_transfer, read_chunk, write_chunk),
return_when=asyncio.FIRST_COMPLETED,
)
if exit_thread in done:
break
if new_transfer in done:
await new_transfer.result().begin()
new_transfer = self._loop.create_task(
self._new_transfer_queue.get()
)
if read_chunk in done:
self._loop.create_task(
self._handle_chunk(
self._read_transfers, read_chunk.result()
)
)
read_chunk = self._loop.create_task(
self._read_chunk_queue.get()
)
if write_chunk in done:
self._loop.create_task(
self._handle_chunk(
self._write_transfers, write_chunk.result()
)
)
write_chunk = self._loop.create_task(
self._write_chunk_queue.get()
)
self._loop.stop()
@staticmethod
async def _handle_chunk(
transfers: _TransferDict, message: transfer_pb2.Chunk
) -> None:
"""Processes an incoming chunk from a stream.
The chunk is dispatched to an active transfer based on its ID. If the
transfer indicates that it is complete, the provided completion callback
is invoked.
"""
chunk = Chunk.from_message(message)
# Find a transfer for the chunk in the list of active transfers.
try:
if chunk.resource_id is not None:
# Prioritize a resource_id if one is set.
transfer = transfers[chunk.resource_id]
else:
# Otherwise, match against either resource or session ID.
transfer = next(
t for t in transfers.values() if t.id == chunk.id()
)
except (KeyError, StopIteration):
_LOG.error(
'TransferManager received chunk for unknown transfer %d',
chunk.id(),
)
# TODO(frolv): What should be done here, if anything?
return
await transfer.handle_chunk(chunk)
def _open_read_stream(self) -> None:
self._read_stream = self._service.Read.invoke(
lambda _, chunk: self._loop.call_soon_threadsafe(
self._read_chunk_queue.put_nowait, chunk
),
on_error=lambda _, status: self._on_read_error(status),
)
def _on_read_error(self, status: Status) -> None:
"""Callback for an RPC error in the read stream."""
if status is Status.FAILED_PRECONDITION:
# FAILED_PRECONDITION indicates that the stream packet was not
# recognized as the stream is not open. This could occur if the
# server resets during an active transfer. Re-open the stream to
# allow pending transfers to continue.
self._open_read_stream()
else:
# Other errors are unrecoverable. Clear the stream and cancel any
# pending transfers with an INTERNAL status as this is a system
# error.
self._read_stream = None
for transfer in self._read_transfers.values():
transfer.finish(Status.INTERNAL, skip_callback=True)
self._read_transfers.clear()
_LOG.error('Read stream shut down: %s', status)
def _open_write_stream(self) -> None:
self._write_stream = self._service.Write.invoke(
lambda _, chunk: self._loop.call_soon_threadsafe(
self._write_chunk_queue.put_nowait, chunk
),
on_error=lambda _, status: self._on_write_error(status),
)
def _on_write_error(self, status: Status) -> None:
"""Callback for an RPC error in the write stream."""
if status is Status.FAILED_PRECONDITION:
# FAILED_PRECONDITION indicates that the stream packet was not
# recognized as the stream is not open. This could occur if the
# server resets during an active transfer. Re-open the stream to
# allow pending transfers to continue.
self._open_write_stream()
else:
# Other errors are unrecoverable. Clear the stream and cancel any
# pending transfers with an INTERNAL status as this is a system
# error.
self._write_stream = None
for transfer in self._write_transfers.values():
transfer.finish(Status.INTERNAL, skip_callback=True)
self._write_transfers.clear()
_LOG.error('Write stream shut down: %s', status)
def _start_read_transfer(self, transfer: Transfer) -> None:
"""Begins a new read transfer, opening the stream if it isn't."""
self._read_transfers[transfer.resource_id] = transfer
if not self._read_stream:
self._open_read_stream()
_LOG.debug('Starting new read transfer %d', transfer.id)
self._loop.call_soon_threadsafe(
self._new_transfer_queue.put_nowait, transfer
)
def _end_read_transfer(self, transfer: Transfer) -> None:
"""Completes a read transfer."""
del self._read_transfers[transfer.resource_id]
if not transfer.status.ok():
_LOG.error(
'Read transfer %d terminated with status %s',
transfer.id,
transfer.status,
)
# TODO(frolv): This doesn't seem to work. Investigate why.
# If no more transfers are using the read stream, close it.
# if not self._read_transfers and self._read_stream:
# self._read_stream.cancel()
# self._read_stream = None
def _start_write_transfer(self, transfer: Transfer) -> None:
"""Begins a new write transfer, opening the stream if it isn't."""
self._write_transfers[transfer.resource_id] = transfer
if not self._write_stream:
self._open_write_stream()
_LOG.debug('Starting new write transfer %d', transfer.id)
self._loop.call_soon_threadsafe(
self._new_transfer_queue.put_nowait, transfer
)
def _end_write_transfer(self, transfer: Transfer) -> None:
"""Completes a write transfer."""
del self._write_transfers[transfer.resource_id]
if not transfer.status.ok():
_LOG.error(
'Write transfer %d terminated with status %s',
transfer.id,
transfer.status,
)
# TODO(frolv): This doesn't seem to work. Investigate why.
# If no more transfers are using the write stream, close it.
# if not self._write_transfers and self._write_stream:
# self._write_stream.cancel()
# self._write_stream = None
class Error(Exception):
"""Exception raised when a transfer fails.
Stores the ID of the failed transfer resource and the error that occurred.
"""
def __init__(self, resource_id: int, status: Status):
super().__init__(f'Transfer {resource_id} failed with status {status}')
self.resource_id = resource_id
self.status = status
|