async def watch(self,
key_range,
start_revision=None,
noput=False,
nodelete=False,
prev_kv=False,
always_reconnect=False,
ignore_compact=False,
batch_events=False,
create_event=False):
filters = []
if noput:
filters.append(rpc.WatchCreateRequest.NOPUT)
if nodelete:
filters.append(rpc.WatchCreateRequest.NODELETE)
reconnect_revision = start_revision
done_future = None
try:
while True:
watch_request = rpc.WatchCreateRequest(
start_revision=reconnect_revision,
filters=filters,
prev_kv=prev_kv)
put_key_range(watch_request, key_range)
self._ensure_watch_task()
output_queue = asyncio.Queue(loop=self._loop)
done_future = self._loop.create_future()
await self._create_request_queue.put(
(watch_request, output_queue, done_future))
try:
await done_future
if create_event:
yield None
create_event = False
while True:
is_event, result, revision = await output_queue.get()
if not is_event:
if revision is not None:
reconnect_revision = revision
if result is None:
break
else:
# When an exception is raised in multiple positions
# the traceback will mix up, so clone the exception
# for each raise
if isinstance(result, WatchException):
raise result._clone() from result
else:
raise WatchException(
"Watch failed with server exception"
) from result
else:
reconnect_revision = revision + 1
if batch_events:
yield tuple(result)
else:
for e in result:
yield e
except CompactRevisonException:
if ignore_compact:
continue
else:
raise
except CancelledError:
raise
except Exception:
if always_reconnect:
continue
else:
raise
else:
break
finally:
if done_future is not None and not done_future.done():
done_future.cancel()
if self._watch_task_running:
done_future = self._loop.create_future()
await self._cancel_request_queue.put(
(output_queue, done_future))
if self._watch_task_running:
try:
await done_future
except Exception:
pass
async def _watch_task(self, reconnect_event):
# Queue for WatchRequest
async def input_iterator(input_queue):
while True:
n = await input_queue.get()
if n is None:
break
yield n
async def watch_call(input_queue, watch_stub, output_pipe):
async with watch_stub.Watch.with_scope(
input_iterator(input_queue)) as response_iter:
async for r in response_iter:
await output_pipe.put(r)
last_received_revision = None
# watch_id -> revision
last_watches_revision = {}
# watch_id -> (WatchCreateRequest, output_queue)
registered_watches = {}
# output_queue -> watch_id
registered_queues = {}
# A tuple (WatchCreateRequest, output_queue, done_future, cancel_future)
pending_create_request = None
# watch_id -> done_future
pending_cancel_requests = {}
# output_queue -> (WatchCreateRequest, done_future)
restore_creates = {}
quitting = False
def _reconnect_revision(watch_id):
if last_received_revision is None:
return None
else:
if watch_id in last_watches_revision:
last_revision = last_watches_revision[watch_id]
return max(last_revision + 1, last_received_revision)
else:
return None
try:
while not quitting: # Auto reconnect when failed or channel updated
reconnect_event.clear()
output_pipe = _Pipe(loop=self._loop)
input_queue = asyncio.Queue(loop=self._loop)
call_task = asyncio.ensure_future(watch_call(
input_queue, self._watch_stub, output_pipe),
loop=self._loop)
try:
# Restore registered watches
for watch_id, (create_request,
output_queue) in registered_watches.items():
if watch_id in pending_cancel_requests:
# Already cancelling
fut = pending_cancel_requests.pop(watch_id)
if not fut.done():
fut.set_result(True)
continue
r = rpc.WatchCreateRequest()
r.CopyFrom(create_request)
restore_revision = _reconnect_revision(watch_id)
if restore_revision is not None:
r.start_revision = restore_revision
restore_creates[output_queue] = (r, None)
registered_watches.clear()
registered_queues.clear()
# Restore pending cancels - should already be processed though
for watch_id, fut in pending_cancel_requests.items():
fut.set_result(True)
pending_cancel_requests.clear()
# Restore pending create request
if pending_create_request is not None:
if pending_create_request[3] is not None: # Cancelled
pending_create_request[1].put_nowait(
(False, None, None))
if pending_create_request[
2] is not None and not pending_create_request[
2].done():
pending_create_request[2].set_result(True)
if pending_create_request[
3] is not None and not pending_create_request[
3].done():
pending_create_request[3].set_result(True)
else:
restore_creates[pending_create_request[1]] = (
pending_create_request[0],
pending_create_request[2])
pending_create_request = None
while True:
if pending_create_request is None:
if restore_creates:
q, (req, fut) = restore_creates.popitem()
# Send create request
pending_create_request = (req, q, fut, None)
input_queue.put_nowait(
rpc.WatchRequest(create_request=req))
if pending_create_request is None:
select_pipes = [
output_pipe, self._create_request_queue,
self._cancel_request_queue
]
else:
select_pipes = [
output_pipe, self._cancel_request_queue
]
reconn_wait = asyncio.ensure_future(
reconnect_event.wait(), loop=self._loop)
select_futs = [reconn_wait, call_task]
if not pending_create_request and not registered_watches and \
not restore_creates:
select_futs.append(
asyncio.sleep(2, loop=self._loop))
pipes, _ = await _select(select_pipes,
select_futs,
loop=self._loop)
reconn_wait.cancel()
if not pipes and not reconnect_event.is_set(
) and not call_task.done():
# No watch, stop the task
quitting = True
break
# Process cancel requests first
if self._cancel_request_queue in pipes:
cancel_requests = self._cancel_request_queue.read_nowait(
)
for output_queue, done_fut in cancel_requests:
if output_queue in pending_cancel_requests:
# Chain this future
pending_cancel_requests[
output_queue].add_done_callback(
lambda f, done_fut=done_fut:
done_fut.set_result(True))
elif output_queue in restore_creates:
# Cancel a request which is not started
_, fut = restore_creates.pop(output_queue)
output_queue.put_nowait(
(False, None, None))
if fut is not None and not fut.done():
fut.set_result(True)
if done_fut is not None and not done_fut.done(
):
done_fut.set_result(True)
elif pending_create_request is not None and \
pending_create_request[1] == output_queue:
# Cancel the pending create watch
if pending_create_request[3] is None:
pending_create_request = pending_create_request[:3] + (
done_fut, )
else:
pending_create_request[
3].add_done_callback(
lambda f, done_fut=done_fut:
done_fut.set_result(True))
else:
watch_id = registered_queues.get(
output_queue)
if watch_id is None:
done_fut.set_result(True)
else:
# Send cancel request and save it to pending requests
input_queue.put_nowait(
rpc.WatchRequest(
cancel_request=rpc.
WatchCancelRequest(
watch_id=watch_id)))
pending_cancel_requests[
watch_id] = done_fut
# Process received events
if output_pipe in pipes:
outputs = output_pipe.read_nowait()
for response in outputs:
if response.created:
assert pending_create_request is not None
if response.compact_revision > 0:
# Cancelled (Is it possible?)
exc = CompactRevisonException(
response.compact_revision)
pending_create_request[1].put_nowait(
(False, exc,
response.compact_revision))
if pending_create_request[2] is not None and not \
pending_create_request[2].done():
pending_create_request[
2].set_exception(exc)
if pending_create_request[3] is not None and not \
pending_create_request[3].done():
pending_create_request[
3].set_result(True)
else:
registered_watches[
response.
watch_id] = pending_create_request[
0:2]
registered_queues[
pending_create_request[
1]] = response.watch_id
if pending_create_request[2] is not None and not \
pending_create_request[2].done():
pending_create_request[
2].set_result(True)
if pending_create_request[
3] is not None:
# Immediately cancel the watch
input_queue.put_nowait(
rpc.WatchRequest(
cancel_request=rpc.
WatchCancelRequest(
watch_id=response.
watch_id)))
pending_cancel_requests[
response.
watch_id] = pending_create_request[
3]
pending_create_request = None
if response.events:
last_received_revision = response.header.revision
last_watches_revision[
response.
watch_id] = last_received_revision
if response.watch_id in registered_watches:
_, output_queue = registered_watches[
response.watch_id]
output_queue.put_nowait((True, [
Event(e, last_received_revision)
for e in response.events
], last_received_revision))
if response.compact_revision > 0:
if response.watch_id in registered_watches:
_, output_queue = registered_watches.pop(
response.watch_id)
exc = CompactRevisonException(
response.compact_revision)
output_queue.put_nowait(
(False, exc,
response.compact_revision))
del registered_queues[output_queue]
if response.watch_id in pending_cancel_requests:
if not pending_cancel_requests[
response.watch_id].done():
pending_cancel_requests[
response.watch_id].set_result(
True)
del pending_cancel_requests[
response.watch_id]
if response.canceled:
# Cancel response
if response.watch_id in registered_watches:
_, output_queue = registered_watches.pop(
response.watch_id)
if response.watch_id in pending_cancel_requests:
# Normal cancel
output_queue.put_nowait(
(False, None, None))
else:
output_queue.put_nowait(
(False,
ServerCancelException(
response.cancel_reason),
_reconnect_revision(
response.watch_id)))
del registered_queues[output_queue]
if response.watch_id in pending_cancel_requests:
if not pending_cancel_requests[
response.watch_id].done():
pending_cancel_requests[
response.watch_id].set_result(
True)
del pending_cancel_requests[
response.watch_id]
if self._create_request_queue in pipes:
while pending_create_request is None and not self._create_request_queue.is_empty(
):
create_req, output_queue, done_fut = self._create_request_queue.get_nowait(
)
if done_fut.done():
# Ignore cancelled create requests
output_queue.put_nowait(
(False, None, None))
continue
# Send create request
pending_create_request = (create_req,
output_queue,
done_fut, None)
input_queue.put_nowait(
rpc.WatchRequest(
create_request=create_req))
if reconnect_event.is_set():
# Reconnected
break
if call_task.done():
# Maybe not available
if call_task.exception() is not None:
await call_task
else:
break
finally:
input_queue.put_nowait(None)
call_task.cancel()
if quitting:
self._watch_task_running = None
try:
await call_task
except Exception:
pass
except Exception as exc:
if registered_queues:
for q, watch_id in registered_queues.items():
q.put_nowait((False, exc, _reconnect_revision(watch_id)))
if pending_create_request is not None:
pending_create_request[1].put_nowait((False, exc, None))
if pending_create_request[
2] is not None and not pending_create_request[2].done(
):
pending_create_request[2].set_exception(exc)
if pending_create_request[
3] is not None and not pending_create_request[3].done(
):
pending_create_request[3].set_result(True)
if pending_cancel_requests:
for _, fut in pending_cancel_requests.items():
if not fut.done():
fut.set_result(True)
if restore_creates:
for q, (_, fut) in restore_creates.items():
if fut is not None and not fut.done():
fut.set_result(exc)
q.put_nowait((False, exc, _reconnect_revision(watch_id)))
if not self._create_request_queue.is_empty():
create_requests = self._create_request_queue.read_nowait()
for r in create_requests:
r[1].put_nowait((False, exc, None))
if r[2] is not None and not r[2].done():
r[2].set_exception(exc)
if not self._cancel_request_queue.is_empty():
cancel_requests = self._cancel_request_queue.read_nowait()
for _, fut in cancel_requests:
if fut is not None and not fut.done():
fut.set_result(True)
if exc is CancelledError:
raise
finally:
self._watch_task_running = None