Create a BotEvent when a bot is leased from Machine Provider

appengine/swarming/server/bot_management.py

 # Copyright 2014 The LUCI Authors. All rights reserved.
 # Use of this source code is governed under the Apache License, Version 2.0
 # that can be found in the LICENSE file.
 
 """Swarming bot management, e.g. list of known bots and their state.
 
           +---------+
           |BotRoot  |
           |id=bot_id|
           +---------+
@@ skipping 66 matching lines @@
   #   dimensions) or via self-health check.
   # - dimensions['id'] is not exactly one item.
   # - invalid HTTP POST request keys.
   # - BotSettings.quarantined was set at that moment.
   quarantined = ndb.BooleanProperty(default=False)
 
   # Affected by event_type == 'request_task', 'task_canceled', 'task_completed',
   # 'task_error'.
   task_id = ndb.StringProperty(indexed=False)
 
+  # Machine Provider lease ID, for bots acquired from Machine Provider.
+  lease_id = ndb.StringProperty(indexed=False)
+
+  # UTC seconds from epoch when bot will be reclaimed by Machine Provider.
+  lease_expiration_ts = ndb.DateTimeProperty(indexed=False)
+
   @property
   def dimensions(self):
     """Returns a dict representation of self.dimensions_flat."""
     if self.dimensions_old:
       return self.dimensions_old
     out = {}
     for i in self.dimensions_flat:
       k, v = i.split(':', 1)
       out.setdefault(k, []).append(v)
     return out
@@ skipping 71 matching lines @@
 
 class BotEvent(_BotCommon):
   """This entity is immutable.
 
   Parent is BotRoot. Key id is monotonically decreasing with
   datastore_utils.store_new_version().
 
   This entity is created on each bot state transition.
   """
   ALLOWED_EVENTS = {
-    'bot_connected', 'bot_error', 'bot_log', 'bot_rebooting', 'bot_shutdown',
-    'bot_terminate',
+    'bot_connected', 'bot_error', 'bot_leased', 'bot_log', 'bot_rebooting',
+    'bot_shutdown', 'bot_terminate',
     'request_restart', 'request_update', 'request_sleep', 'request_task',
     'task_completed', 'task_error', 'task_update',
   }
   # Dimensions are used for task selection. They are encoded as a list of
   # key:value. Keep in mind that the same key can be used multiple times. The
   # list must be sorted.
   # It is NOT indexed because this is not needed for events.
   dimensions_flat = ndb.StringProperty(repeated=True, indexed=False)
 
   # Common properties for all events (which includes everything in _BotCommon).
@@ skipping 115 matching lines @@
   - dimensions: Bot's dimensions as self-reported. If not provided, keep
         previous value.
   - state: ephemeral state of the bot. It is expected to change constantly. If
         not provided, keep previous value.
   - version: swarming_bot.zip version as self-reported. Used to spot if a bot
         failed to update promptly. If not provided, keep previous value.
   - quarantined: bool to determine if the bot was declared quarantined.
   - task_id: packed task id if relevant. Set to '' to zap the stored value.
   - task_name: task name if relevant. Zapped when task_id is zapped.
   - kwargs: optional values to add to BotEvent relevant to event_type.
+  - lease_id (in kwargs): ID assigned by Machine Provider for this bot.
+  - lease_expiration_ts (in kwargs): UTC seconds from epoch when Machine
+        Provider lease expires.
   """
   if not bot_id:
     return
 
   # Retrieve the previous BotInfo and update it.
   info_key = get_info_key(bot_id)
   bot_info = info_key.get() or BotInfo(key=info_key)
   bot_info.last_seen_ts = utils.utcnow()
   bot_info.external_ip = external_ip
   bot_info.authenticated_as = authenticated_as
   if dimensions:
     bot_info.dimensions_flat = dimensions_to_flat(dimensions)
   if state:
     bot_info.state = state
   if quarantined is not None:
     bot_info.quarantined = quarantined
   if task_id is not None:
     bot_info.task_id = task_id
   if task_name:
     bot_info.task_name = task_name
   if version is not None:
     bot_info.version = version
+  if kwargs.get('lease_id') is not None:
+    bot_info.lease_id = kwargs['lease_id']
+  if kwargs.get('lease_expiration_ts') is not None:
+    bot_info.lease_expiration_ts = kwargs['lease_expiration_ts']
 
   if event_type in ('request_sleep', 'task_update'):
     # Handle this specifically. It's not much of an even worth saving a BotEvent
     # for but it's worth updating BotInfo. The only reason BotInfo is GET is to
     # keep first_seen_ts. It's not necessary to use a transaction here since no
     # BotEvent is being added, only last_seen_ts is really updated.
     bot_info.put()
     return
 
   event = BotEvent(
@@ skipping 46 matching lines @@
   Returns:
     Tuple (True to restart, text message explaining the reason).
   """
   # Periodically reboot bots to workaround OS level leaks (especially on Win).
   running_time = state.get('running_time', 0)
   assert isinstance(running_time, (int, float))
   period = get_bot_reboot_period(bot_id, state)
   if period and running_time > period:
     return True, 'Periodic reboot: running longer than %ds' % period
   return False, ''