def process(messages):
"""Processes the given Pub/Sub messages.
Args:
messages: A list of message dicts where each message refers to an Instance
with the same parent InstanceGroupManager.
"""
logging.info('Processing messages: %s', len(messages))
ack_ids = [message['ackId'] for message in messages]
# Since it can take some time to handle the Pub/Sub messages and commit the
# large transaction, extend the ack deadline to match the cron interval with
# which we pull Pub/Sub messages. We should not extend the deadline longer
# than that, otherwise if the transaction fails and the messages are not
# acknowledged then we won't receive those messages again during the next
# run of the cron job.
yield pubsub.modify_ack_deadline_async(
pubsub.full_subscription_name(
get_machine_provider_topic_project(),
get_machine_provider_subscription(),
), 60, *ack_ids)
yield _process(messages)
yield pubsub.ack_async(
pubsub.full_subscription_name(
get_machine_provider_topic_project(),
get_machine_provider_subscription(),
), *ack_ids)
def post(self):
"""Subscribe a machine to a Cloud Pub/Sub topic.
Params:
backend_project: If specified, project that the machine subscription
topic is contained in for the backend.
backend_attributes: If specified, JSON-encoded dict of attributes to
include in the machine subscription message for the backend.
backend_topic: If specified, topic that the machine subscription should
be published to for the backend.
hostname: Hostname being reclaimed.
machine_id: ID of the CatalogMachineEntry being reclaimed.
machine_service_account: Service account to authorize to consume the
subscription.
machine_subscription: Cloud Pub/Sub subscription to create for the
machine.
machine_subscription_project: Project the Cloud Pub/Sub subscription
should live in.
machine_topic: Cloud Pub/Sub topic to create for the machine.
machine_topic_project: Project the Cloud Pub/Sub topic should live in.
"""
backend_attributes = json.loads(
self.request.get('backend_attributes', {}))
backend_project = self.request.get('backend_project')
backend_topic = self.request.get('backend_topic')
hostname = self.request.get('hostname')
machine_id = self.request.get('machine_id')
machine_service_account = self.request.get('machine_service_account')
machine_subscription = self.request.get('machine_subscription')
machine_subscription_project = self.request.get(
'machine_subscription_project')
machine_topic = self.request.get('machine_topic')
machine_topic_project = self.request.get('machine_topic_project')
topic = pubsub.full_topic_name(machine_topic_project, machine_topic)
subscription = pubsub.full_subscription_name(
machine_subscription_project, machine_subscription)
pubsub.ensure_subscription_deleted(subscription)
pubsub.ensure_subscription_exists(subscription, topic)
with pubsub.iam_policy(subscription) as policy:
policy.add_member(
'roles/pubsub.subscriber',
'serviceAccount:%s' % machine_service_account,
)
if backend_topic:
attributes = backend_attributes.copy()
attributes['hostname'] = hostname
attributes['subscription'] = machine_subscription
attributes['subscription_project'] = machine_subscription_project
attributes['topic'] = machine_topic
attributes['topic_project'] = machine_topic_project
pubsub.publish(
pubsub.full_topic_name(backend_project, backend_topic),
'SUBSCRIBED',
attributes,
)
set_available(ndb.Key(models.CatalogMachineEntry, machine_id))
def poll():
"""Polls and processes Pub/Sub messages."""
response = pubsub.pull(pubsub.full_subscription_name(
get_machine_provider_topic_project(),
get_machine_provider_subscription(),
))
utilities.batch_process_async(response.get('receivedMessages', []), process)
def poll():
"""Polls and processes Pub/Sub messages."""
response = pubsub.pull(
pubsub.full_subscription_name(
get_machine_provider_topic_project(),
get_machine_provider_subscription(),
),
max_messages=400,
)
if response.get('receivedMessages', []):
logging.info('Messages received: %s',
len(response['receivedMessages']))
messages = split_by_entity_group(response['receivedMessages'])
utilities.batch_process_async(messages, process)
def add_subscription_metadata(key, subscription_project, subscription):
"""Queues the addition of subscription metadata.
Args:
key: ndb.Key for a models.Instance entity.
subscription_project: Project containing the Pub/Sub subscription.
subscription: Name of the Pub/Sub subscription that Machine Provider will
communicate with the instance on.
"""
entity = key.get()
if not entity:
logging.warning('Instance does not exist: %s', key)
return
if entity.pubsub_subscription:
return
parent = key.parent().get()
if not parent:
logging.warning('InstanceGroupManager does not exist: %s', key.parent())
return
grandparent = parent.key.parent().get()
if not grandparent:
logging.warning(
'InstanceTemplateRevision does not exist: %s', parent.key.parent())
return
if not grandparent.service_accounts:
logging.warning(
'InstanceTemplateRevision service account unspecified: %s',
parent.key.parent(),
)
return
logging.info('Instance Pub/Sub subscription received: %s', key)
entity.pending_metadata_updates.append(models.MetadataUpdate(
metadata={
'pubsub_service_account': grandparent.service_accounts[0].name,
'pubsub_subscription': subscription,
'pubsub_subscription_project': subscription_project,
},
))
entity.pubsub_subscription = pubsub.full_subscription_name(
subscription_project, subscription)
entity.put()
def post(self):
"""Reclaim a machine.
Params:
hostname: Hostname of the machine being reclaimed.
machine_key: URL-safe ndb.Key for a models.CatalogMachineEntry.
machine_subscription: Subscription created for the machine to listen
for instructions on.
machine_subscription_project: Project that the machine subscription is
contained in.
machine_topic: Topic that the machine communication should occur on.
machine_topic_project: Project that the machine communication topic is
contained in.
policies: JSON-encoded string representation of the
rpc_messages.Policies governing this machine.
request_json: JSON-encoded string representation of the
rpc_messages.LeaseRequest being fulfilled.
response_json: JSON-encoded string representation of the
rpc_messages.LeaseResponse being delivered.
"""
hostname = self.request.get('hostname')
machine_key = ndb.Key(urlsafe=self.request.get('machine_key'))
machine_subscription = self.request.get('machine_subscription')
machine_subscription_project = self.request.get(
'machine_subscription_project')
machine_topic = self.request.get('machine_topic')
machine_topic_project = self.request.get('machine_topic_project')
policies = json.loads(self.request.get('policies'))
request = json.loads(self.request.get('request_json'))
response = json.loads(self.request.get('response_json'))
assert machine_key.kind() == 'CatalogMachineEntry', machine_key
maybe_notify_backend('RECLAIMED', hostname, policies)
maybe_notify_lessee(request, response)
# Delete machine Pub/Sub channel.
pubsub.ensure_subscription_deleted(
pubsub.full_subscription_name(machine_subscription_project,
machine_subscription))
pubsub.ensure_topic_deleted(
pubsub.full_topic_name(machine_topic_project, machine_topic))
reclaim(machine_key)
metrics.lease_requests_expired.increment()
def process(message):
"""Processes the given Pub/Sub message.
Args:
message: A message dict.
"""
attributes = message.get('message', {}).get('attributes', {})
data = base64.b64decode(message.get('message', {}).get('data', ''))
logging.info(
'Received Pub/Sub message: %s\nMessage: %s\nAttributes: %s',
message['ackId'],
data,
json.dumps(attributes, indent=2),
)
key = ndb.Key(urlsafe=attributes['key'])
if key.kind() == 'Instance':
if data == 'LEASED':
logging.info('Instance leased: %s', key)
elif data == 'RECLAIMED':
yield instances.mark_for_deletion(key)
elif data == 'SUBSCRIBED':
yield instances.add_subscription_metadata(
key,
attributes['subscription_project'],
attributes['subscription'],
)
else:
logging.error('Unexpected key: %s', key)
yield pubsub.ack_async(
pubsub.full_subscription_name(
get_machine_provider_topic_project(), get_machine_provider_topic()),
message['ackId'],
)
def test_full_subscription_name(self):
self.assertEqual(
'projects/abc/subscriptions/def',
pubsub.full_subscription_name('abc', 'def'))
