def batched_lease(self, request):
"""Handles an incoming BatchedLeaseRequest.
Batches are intended to save on RPCs only. The batched requests will not
execute transactionally.
"""
# To avoid having large batches timed out by AppEngine after 60 seconds
# when some requests have been processed and others haven't, enforce a
# smaller deadline on ourselves to process the entire batch.
DEADLINE_SECS = 30
start_time = utils.utcnow()
user = auth.get_current_identity().to_bytes()
logging.info('Received BatchedLeaseRequest:\nUser: %s\n%s', user,
request)
responses = []
for request in request.requests:
request_hash = models.LeaseRequest.generate_key(user, request).id()
logging.info(
'Processing LeaseRequest:\nRequest hash: %s\n%s',
request_hash,
request,
)
if (utils.utcnow() - start_time).seconds > DEADLINE_SECS:
logging.warning(
'BatchedLeaseRequest exceeded enforced deadline: %s',
DEADLINE_SECS)
responses.append(
rpc_messages.LeaseResponse(
client_request_id=request.request_id,
error=rpc_messages.LeaseRequestError.DEADLINE_EXCEEDED,
request_hash=request_hash,
))
else:
try:
responses.append(self._lease(request, user, request_hash))
except (
datastore_errors.NotSavedError,
datastore_errors.Timeout,
runtime.apiproxy_errors.CancelledError,
runtime.apiproxy_errors.DeadlineExceededError,
runtime.apiproxy_errors.OverQuotaError,
) as e:
logging.warning('Exception processing LeaseRequest:\n%s',
e)
responses.append(
rpc_messages.LeaseResponse(
client_request_id=request.request_id,
error=rpc_messages.LeaseRequestError.
TRANSIENT_ERROR,
request_hash=request_hash,
))
return rpc_messages.BatchedLeaseResponse(responses=responses)
def test_releases(self):
self.mock(handlers_cron, 'release_lease', lambda *args, **kwargs: True)
request = rpc_messages.LeaseRequest(
dimensions=rpc_messages.Dimensions(
os_family=rpc_messages.OSFamily.LINUX,
),
duration=1,
request_id='fake-id',
)
key = models.LeaseRequest(
deduplication_checksum=
models.LeaseRequest.compute_deduplication_checksum(request),
key=models.LeaseRequest.generate_key(
auth_testing.DEFAULT_MOCKED_IDENTITY.to_bytes(),
request,
),
owner=auth_testing.DEFAULT_MOCKED_IDENTITY,
released=True,
request=request,
response=rpc_messages.LeaseResponse(
client_request_id='fake-id',
state=rpc_messages.LeaseRequestState.UNTRIAGED,
),
).put()
self.app.get(
'/internal/cron/process-lease-releases',
headers={'X-AppEngine-Cron': 'true'},
)
def test_no_machine_id(self):
request = rpc_messages.LeaseRequest(
dimensions=rpc_messages.Dimensions(
os_family=rpc_messages.OSFamily.LINUX,
),
duration=1,
request_id='fake-id',
)
lease_key = models.LeaseRequest(
deduplication_checksum=
models.LeaseRequest.compute_deduplication_checksum(request),
owner=auth_testing.DEFAULT_MOCKED_IDENTITY,
released=True,
request=request,
response=rpc_messages.LeaseResponse(
client_request_id='fake-id',
),
).put()
machine_key = models.CatalogMachineEntry(
id='id',
dimensions=rpc_messages.Dimensions(
os_family=rpc_messages.OSFamily.LINUX,
),
).put()
handlers_cron.release_lease(lease_key)
self.assertFalse(lease_key.get().response.lease_expiration_ts)
self.assertFalse(lease_key.get().released)
self.assertFalse(machine_key.get().lease_expiration_ts)
def test_no_expiration_ts(self):
self.mock(utils, 'enqueue_task', lambda *args, **kwargs: True)
request = rpc_messages.LeaseRequest(
dimensions=rpc_messages.Dimensions(
hostname='fake-host',
os_family=rpc_messages.OSFamily.LINUX,
),
duration=1,
request_id='fake-id',
)
lease_request_key = models.LeaseRequest(
id='id',
deduplication_checksum=
models.LeaseRequest.compute_deduplication_checksum(request),
owner=auth_testing.DEFAULT_MOCKED_IDENTITY,
request=request,
response=rpc_messages.LeaseResponse(
client_request_id='fake-id',
hostname='fake-host',
),
).put()
machine_key = models.CatalogMachineEntry(
dimensions=rpc_messages.Dimensions(
os_family=rpc_messages.OSFamily.LINUX,
),
lease_id=lease_request_key.id(),
policies=rpc_messages.Policies(
machine_service_account='fake-service-account',
),
).put()
handlers_cron.reclaim_machine(machine_key, utils.utcnow())
self.assertTrue(lease_request_key.get().response.hostname)
def test_one_request_one_matching_machine_entry(self):
request = rpc_messages.LeaseRequest(
dimensions=rpc_messages.Dimensions(
os_family=rpc_messages.OSFamily.LINUX, ),
duration=1,
request_id='fake-id',
)
models.LeaseRequest(
deduplication_checksum=models.LeaseRequest.
compute_deduplication_checksum(request, ),
key=models.LeaseRequest.generate_key(
auth_testing.DEFAULT_MOCKED_IDENTITY.to_bytes(),
request,
),
owner=auth_testing.DEFAULT_MOCKED_IDENTITY,
request=request,
response=rpc_messages.LeaseResponse(client_request_id='fake-id', ),
).put()
models.CatalogMachineEntry.create_and_put(
rpc_messages.Dimensions(
backend=rpc_messages.Backend.DUMMY,
hostname='fake-host',
os_family=rpc_messages.OSFamily.LINUX,
),
rpc_messages.Policies(
backend_project='fake-project',
backend_topic='fake-topic',
),
models.CatalogMachineEntryStates.AVAILABLE,
)
self.app.get(
'/internal/cron/process-lease-requests',
headers={'X-AppEngine-Cron': 'true'},
)
def test_leased_task_failed(self):
self.mock(utils, 'enqueue_task', lambda *args, **kwargs: False)
machine_key = models.CatalogMachineEntry(
id='machine-id',
dimensions=rpc_messages.Dimensions(
os_family=rpc_messages.OSFamily.LINUX,
),
policies=rpc_messages.Policies(
machine_service_account='service-account',
),
state=models.CatalogMachineEntryStates.AVAILABLE,
).put()
request = rpc_messages.LeaseRequest(
dimensions=rpc_messages.Dimensions(
os_family=rpc_messages.OSFamily.LINUX,
),
duration=1,
request_id='request-id',
)
lease_request_key = models.LeaseRequest(
id='lease-id',
deduplication_checksum=
models.LeaseRequest.compute_deduplication_checksum(request),
owner=auth_testing.DEFAULT_MOCKED_IDENTITY,
request=request,
response=rpc_messages.LeaseResponse(
client_request_id='client-request-id',
state=rpc_messages.LeaseRequestState.UNTRIAGED,
),
).put()
with self.assertRaises(handlers_cron.TaskEnqueuingError):
handlers_cron.lease_machine(machine_key, lease_request_key.get())
self.assertFalse(lease_request_key.get().machine_id)
self.assertEqual(
lease_request_key.get().response.state,
rpc_messages.LeaseRequestState.UNTRIAGED,
)
self.assertFalse(machine_key.get().lease_id)
self.assertEqual(
machine_key.get().state,
models.CatalogMachineEntryStates.AVAILABLE,
)
def test_reclaim_immediately(self):
request = rpc_messages.LeaseRequest(
dimensions=rpc_messages.Dimensions(
os_family=rpc_messages.OSFamily.LINUX,
),
duration=0,
request_id='fake-id',
)
lease = models.LeaseRequest(
deduplication_checksum=
models.LeaseRequest.compute_deduplication_checksum(request),
key=models.LeaseRequest.generate_key(
auth_testing.DEFAULT_MOCKED_IDENTITY.to_bytes(),
request,
),
owner=auth_testing.DEFAULT_MOCKED_IDENTITY,
request=request,
response=rpc_messages.LeaseResponse(
client_request_id='fake-id',
),
)
dimensions = rpc_messages.Dimensions(
backend=rpc_messages.Backend.DUMMY,
hostname='fake-host',
os_family=rpc_messages.OSFamily.LINUX,
)
machine = models.CatalogMachineEntry(
dimensions=dimensions,
key=models.CatalogMachineEntry.generate_key(dimensions),
lease_id=lease.key.id(),
lease_expiration_ts=datetime.datetime.utcfromtimestamp(1),
policies=rpc_messages.Policies(
machine_service_account='fake-service-account',
),
state=models.CatalogMachineEntryStates.AVAILABLE,
).put()
lease.machine_id = machine.id()
lease.put()
self.app.get(
'/internal/cron/process-machine-reclamations',
headers={'X-AppEngine-Cron': 'true'},
)
def test_one_request_one_matching_machine_entry_lease_expiration_ts(self):
ts = int(utils.time_time())
request = rpc_messages.LeaseRequest(
dimensions=rpc_messages.Dimensions(
os_family=rpc_messages.OSFamily.LINUX,
),
lease_expiration_ts=ts,
request_id='fake-id',
)
key = models.LeaseRequest(
deduplication_checksum=
models.LeaseRequest.compute_deduplication_checksum(request),
key=models.LeaseRequest.generate_key(
auth_testing.DEFAULT_MOCKED_IDENTITY.to_bytes(),
request,
),
owner=auth_testing.DEFAULT_MOCKED_IDENTITY,
request=request,
response=rpc_messages.LeaseResponse(
client_request_id='fake-id',
state=rpc_messages.LeaseRequestState.UNTRIAGED,
),
).put()
dimensions = rpc_messages.Dimensions(
backend=rpc_messages.Backend.DUMMY,
hostname='fake-host',
os_family=rpc_messages.OSFamily.LINUX,
)
models.CatalogMachineEntry(
key=models.CatalogMachineEntry.generate_key(dimensions),
dimensions=dimensions,
policies=rpc_messages.Policies(
machine_service_account='fake-service-account',
),
state=models.CatalogMachineEntryStates.AVAILABLE,
).put()
self.app.get(
'/internal/cron/process-lease-requests',
headers={'X-AppEngine-Cron': 'true'},
)
self.assertEqual(key.get().response.lease_expiration_ts, ts)
def test_reclaimed(self):
lease_key = models.LeaseRequest(
id='fake-id',
deduplication_checksum='checksum',
machine_id='fake-host',
owner=auth_testing.DEFAULT_MOCKED_IDENTITY,
request=rpc_messages.LeaseRequest(
dimensions=rpc_messages.Dimensions(),
request_id='request-id',
),
response=rpc_messages.LeaseResponse(
client_request_id='request-id',
hostname='fake-host',
),
).put()
machine_key = models.CatalogMachineEntry(
dimensions=rpc_messages.Dimensions(),
lease_id=lease_key.id(),
).put()
handlers_queues.reclaim(machine_key)
self.assertFalse(lease_key.get().machine_id)
self.assertFalse(lease_key.get().response.hostname)
self.assertFalse(machine_key.get())
def _lease(self, request, user, request_hash):
"""Handles an incoming LeaseRequest."""
if request.pubsub_topic:
if not pubsub.validate_topic(request.pubsub_topic):
logging.warning(
'Invalid topic for Cloud Pub/Sub: %s',
request.pubsub_topic,
)
return rpc_messages.LeaseResponse(
client_request_id=request.request_id,
error=rpc_messages.LeaseRequestError.INVALID_TOPIC,
)
if not request.pubsub_project:
logging.info(
'Cloud Pub/Sub project unspecified, using default: %s',
PUBSUB_DEFAULT_PROJECT,
)
request.pubsub_project = PUBSUB_DEFAULT_PROJECT
if request.pubsub_project:
if not pubsub.validate_project(request.pubsub_project):
logging.warning(
'Invalid project for Cloud Pub/Sub: %s',
request.pubsub_topic,
)
return rpc_messages.LeaseResponse(
client_request_id=request.request_id,
error=rpc_messages.LeaseRequestError.INVALID_PROJECT,
)
elif not request.pubsub_topic:
logging.warning(
'Cloud Pub/Sub project specified without specifying topic: %s',
request.pubsub_project,
)
return rpc_messages.LeaseResponse(
client_request_id=request.request_id,
error=rpc_messages.LeaseRequestError.UNSPECIFIED_TOPIC,
)
duplicate = models.LeaseRequest.get_by_id(request_hash)
deduplication_checksum = models.LeaseRequest.compute_deduplication_checksum(
request, )
if duplicate:
# Found a duplicate request ID from the same user. Attempt deduplication.
if deduplication_checksum == duplicate.deduplication_checksum:
# The LeaseRequest RPC we just received matches the original.
# We're safe to dedupe.
logging.info(
'Dropped duplicate LeaseRequest:\n%s',
duplicate.response,
)
return duplicate.response
else:
logging.warning('Request ID reuse:\nOriginally used for:\n%s',
duplicate.request)
return rpc_messages.LeaseResponse(
client_request_id=request.request_id,
error=rpc_messages.LeaseRequestError.REQUEST_ID_REUSE,
)
else:
logging.info('Storing LeaseRequest')
response = rpc_messages.LeaseResponse(
client_request_id=request.request_id,
request_hash=request_hash,
state=rpc_messages.LeaseRequestState.UNTRIAGED,
)
models.LeaseRequest(
deduplication_checksum=deduplication_checksum,
id=request_hash,
owner=auth.get_current_identity(),
request=request,
response=response,
).put()
logging.info('Sending LeaseResponse:\n%s', response)
return response
def _lease(self, request, user, request_hash):
"""Handles an incoming LeaseRequest."""
# Arbitrary limit. Increase if necessary.
MAX_LEASE_DURATION = 60 * 60 * 24 * 2
now = utils.time_time()
max_lease_expiration_ts = now + MAX_LEASE_DURATION
metrics.lease_requests_received.increment()
if request.duration:
if request.lease_expiration_ts:
return rpc_messages.LeaseResponse(
client_request_id=request.request_id,
error=rpc_messages.LeaseRequestError.
MUTUAL_EXCLUSION_ERROR,
)
if request.duration < 1:
return rpc_messages.LeaseResponse(
client_request_id=request.request_id,
error=rpc_messages.LeaseRequestError.NONPOSITIVE_DEADLINE,
)
if request.duration > MAX_LEASE_DURATION:
return rpc_messages.LeaseResponse(
client_request_id=request.request_id,
error=rpc_messages.LeaseRequestError.LEASE_TOO_LONG,
)
elif request.lease_expiration_ts:
if request.lease_expiration_ts <= now:
return rpc_messages.LeaseResponse(
client_request_id=request.request_id,
error=rpc_messages.LeaseRequestError.
LEASE_EXPIRATION_TS_ERROR,
)
if request.lease_expiration_ts > max_lease_expiration_ts:
return rpc_messages.LeaseResponse(
client_request_id=request.request_id,
error=rpc_messages.LeaseRequestError.LEASE_TOO_LONG,
)
else:
return rpc_messages.LeaseResponse(
client_request_id=request.request_id,
error=rpc_messages.LeaseRequestError.LEASE_LENGTH_UNSPECIFIED,
)
if request.pubsub_topic:
if not pubsub.validate_topic(request.pubsub_topic):
logging.warning(
'Invalid topic for Cloud Pub/Sub: %s',
request.pubsub_topic,
)
return rpc_messages.LeaseResponse(
client_request_id=request.request_id,
error=rpc_messages.LeaseRequestError.INVALID_TOPIC,
)
if not request.pubsub_project:
logging.info(
'Cloud Pub/Sub project unspecified, using default: %s',
PUBSUB_DEFAULT_PROJECT,
)
request.pubsub_project = PUBSUB_DEFAULT_PROJECT
if request.pubsub_project:
if not pubsub.validate_project(request.pubsub_project):
logging.warning(
'Invalid project for Cloud Pub/Sub: %s',
request.pubsub_topic,
)
return rpc_messages.LeaseResponse(
client_request_id=request.request_id,
error=rpc_messages.LeaseRequestError.INVALID_PROJECT,
)
elif not request.pubsub_topic:
logging.warning(
'Cloud Pub/Sub project specified without specifying topic: %s',
request.pubsub_project,
)
return rpc_messages.LeaseResponse(
client_request_id=request.request_id,
error=rpc_messages.LeaseRequestError.UNSPECIFIED_TOPIC,
)
duplicate = models.LeaseRequest.get_by_id(request_hash)
deduplication_checksum = models.LeaseRequest.compute_deduplication_checksum(
request, )
if duplicate:
# Found a duplicate request ID from the same user. Attempt deduplication.
if deduplication_checksum == duplicate.deduplication_checksum:
metrics.lease_requests_deduped.increment()
# The LeaseRequest RPC we just received matches the original.
# We're safe to dedupe.
logging.info(
'Dropped duplicate LeaseRequest:\n%s',
duplicate.response,
)
return duplicate.response
else:
logging.warning('Request ID reuse:\nOriginally used for:\n%s',
duplicate.request)
return rpc_messages.LeaseResponse(
client_request_id=request.request_id,
error=rpc_messages.LeaseRequestError.REQUEST_ID_REUSE,
)
else:
logging.info('Storing LeaseRequest')
response = rpc_messages.LeaseResponse(
client_request_id=request.request_id,
request_hash=request_hash,
state=rpc_messages.LeaseRequestState.UNTRIAGED,
)
models.LeaseRequest(
deduplication_checksum=deduplication_checksum,
id=request_hash,
owner=auth.get_current_identity(),
request=request,
response=response,
).put()
logging.info('Sending LeaseResponse:\n%s', response)
return response
