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'))