async def test_add_lease_prolong_delete(taskqueue):
# Initialize tasks
queue_str = 'pull-queue-b'
queue_bytes = bytes(queue_str, 'utf8')
tasks = []
for n in range(10):
add_task = taskqueue_service_pb2.TaskQueueAddRequest()
add_task.app_id = bytes(TEST_PROJECT, 'utf8')
add_task.queue_name = queue_bytes
add_task.mode = taskqueue_service_pb2.TaskQueueMode.PULL
add_task.task_name = b'task-%d' % n
add_task.body = b'some-payload-%d' % n
add_task.eta_usec = 0
tasks.append(add_task)
bulk_add = taskqueue_service_pb2.TaskQueueBulkAddRequest()
bulk_add.add_request.extend(tasks[:5])
# Add tasks using single and bulk add
await taskqueue.protobuf('BulkAdd', bulk_add)
asyncio_tasks = []
for task in tasks[5:]:
asyncio_tasks.append(taskqueue.protobuf('Add', task))
await asyncio.wait(asyncio_tasks)
remote_time = await taskqueue.remote_time_usec()
# Lease 10 tasks for 5 second
lease_req = taskqueue_service_pb2.TaskQueueQueryAndOwnTasksRequest()
lease_req.queue_name = queue_bytes
lease_req.lease_seconds = 5
lease_req.max_tasks = 10
leased = await taskqueue.protobuf('QueryAndOwnTasks', lease_req)
# Make sure eta_usec is ~5s greater than current remote_time
assert all(
task.eta_usec == pytest.approx(remote_time + 5_000_000, abs=600_000)
for task in leased.task)
def lease(self):
""" Locust action performed by worker.
Sends lease request to TaskQueue protobuffer API.
Reports leased tasks to validation log. Deletes tasks which
suppose to be done according to scenario.
"""
# Lease tasks
lease_req = taskqueue_service_pb2.TaskQueueQueryAndOwnTasksRequest(
)
lease_req.queue_name = bytes(PULL_QUEUE, 'utf8')
lease_req.lease_seconds = LEASE_SECONDS
lease_req.max_tasks = MAX_TASKS
with self.client.protobuf('QueryAndOwnTasks', lease_req) as leased:
# Report what tasks were leased
for task in leased.task:
self.locust.log_action('LEASED',
str(task.task_name, 'utf8'),
int(task.eta_usec / 1000))
# Parse payload to see how long should take task execution
if not leased.task:
return
Worker.LAST_LEASE_TIME = time.time()
if PULL_QUEUES_BACKEND == 'cassandra':
for task in leased.task:
task.retry_count += 1
tasks_info = [self.get_task_info(task) for task in leased.task]
# Assuming that virtual tasks can be run in parallel
work_time_ms = max(
work_time for task, work_time, should_fail in tasks_info)
gevent.sleep(work_time_ms / 1000)
# Determine tasks that should not be retried
tasks_to_delete = [
task for task, work_time, should_fail in tasks_info
if not should_fail
]
# Delete tasks
delete_req = taskqueue_service_pb2.TaskQueueDeleteRequest()
delete_req.queue_name = bytes(PULL_QUEUE, 'utf8')
delete_req.task_name.extend(
[task.task_name for task in tasks_to_delete])
with self.client.protobuf('Delete', delete_req):
# Report what tasks were deleted
for task in tasks_to_delete:
self.locust.log_action('DELETED',
str(task.task_name, 'utf8'),
f'{task.retry_count}')
async def test_add_lease_retry_retry_delete_pg(taskqueue, pull_queues_backend):
if pull_queues_backend == CASSANDRA:
pytest.skip('Skipped for Cassandra backend')
# Initialize tasks
queue_str = 'pull-queue-5-retry'
queue_bytes = bytes(queue_str, 'utf8')
add_tasks = []
for n in range(4):
add_task = taskqueue_service_pb2.TaskQueueAddRequest()
add_task.app_id = bytes(TEST_PROJECT, 'utf8')
add_task.queue_name = queue_bytes
add_task.mode = taskqueue_service_pb2.TaskQueueMode.PULL
add_task.task_name = b'task-%d' % n
add_task.body = b'some-payload-%d' % n
add_task.eta_usec = 0
add_tasks.append(add_task)
bulk_add = taskqueue_service_pb2.TaskQueueBulkAddRequest()
bulk_add.add_request.extend(add_tasks)
start_time = await taskqueue.remote_time_usec()
total_delay = 0
# Add tasks using bulk add
_, delay = await taskqueue.timed_protobuf('BulkAdd', bulk_add)
total_delay += delay
# Lease 4 tasks for 2 seconds
lease_req = taskqueue_service_pb2.TaskQueueQueryAndOwnTasksRequest()
lease_req.queue_name = queue_bytes
lease_req.lease_seconds = 2
lease_req.max_tasks = 4
leased, delay = await taskqueue.timed_protobuf('QueryAndOwnTasks',
lease_req)
total_delay += delay
assert [task.retry_count for task in leased.task] == [1, 1, 1, 1]
# Try to lease 4 tasks for 2 seconds
leased, delay = await taskqueue.timed_protobuf('QueryAndOwnTasks',
lease_req)
total_delay += delay
assert [task.retry_count for task in leased.task] == []
# Give 3 seconds for lease to expire
time.sleep(3)
# Lease 2 tasks for 2 seconds (retry)
lease_req.max_tasks = 2
leased, delay = await taskqueue.timed_protobuf('QueryAndOwnTasks',
lease_req)
total_delay += delay
assert [task.retry_count for task in leased.task] == [2, 2]
# Lease 2 tasks for 2 seconds (retry)
leased, delay = await taskqueue.timed_protobuf('QueryAndOwnTasks',
lease_req)
total_delay += delay
assert [task.retry_count for task in leased.task] == [2, 2]
# Try to lease 2 tasks for 2 seconds
leased, delay = await taskqueue.timed_protobuf('QueryAndOwnTasks',
lease_req)
total_delay += delay
assert [task.retry_count for task in leased.task] == []
# Give 3 seconds for lease to expire
time.sleep(3)
# Try to lease 3 tasks for 2 seconds
lease_req.max_tasks = 3
leased, delay = await taskqueue.timed_protobuf('QueryAndOwnTasks',
lease_req)
total_delay += delay
assert [task.retry_count for task in leased.task] == [3, 3, 3]
# Give 3 seconds for lease to expire
time.sleep(3)
# Try to lease 3 tasks for 2 seconds
leased, delay = await taskqueue.timed_protobuf('QueryAndOwnTasks',
lease_req)
total_delay += delay
assert [task.retry_count for task in leased.task] == [3, 4, 4]
# Verify listed tasks
listed = await taskqueue.rest('GET', path_suffix=f'/{queue_str}/tasks')
sorting_key = lambda item: (item['retry_count'], int(item['leaseTimestamp']
))
tasks = sorted(listed.json['items'], key=sorting_key)
actual = [(task['retry_count'], int(task['leaseTimestamp']))
for task in tasks]
expected = [
(3, pytest.approx(start_time + 8_000_000, abs=300_000 + total_delay)),
(3, pytest.approx(start_time + 11_000_000, abs=300_000 + total_delay)),
(4, pytest.approx(start_time + 11_000_000, abs=300_000 + total_delay)),
(4, pytest.approx(start_time + 11_000_000, abs=300_000 + total_delay)),
]
assert actual == expected
# Delete tasks
req = taskqueue_service_pb2.TaskQueueDeleteRequest()
req.queue_name = queue_bytes
req.task_name.extend([task.task_name for task in add_tasks])
await taskqueue.protobuf('Delete', req)
# Verify that queue is empty
listed = await taskqueue.rest('GET', path_suffix=f'/{queue_str}/tasks')
assert listed.json == {
'kind': 'taskqueues#tasks'
} # items should be missing