Create a BotEvent when a bot is leased from Machine Provider

appengine/swarming/server/lease_management.py

 # Copyright 2016 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.
 
 """Lease management for machines leased from the Machine Provider.
 
 Keeps a list of machine types which should be leased from the Machine Provider
 and the list of machines of each type currently leased.
 
 Swarming integration with Machine Provider
@@ skipping 53 matching lines @@
   # Request ID used to generate this request.
   client_request_id = ndb.StringProperty(indexed=True)
   # Whether or not this MachineLease should issue lease requests.
   drained = ndb.BooleanProperty(indexed=True)
   # Hostname of the machine currently allocated for this request.
   hostname = ndb.StringProperty()
   # Duration to lease for.
   lease_duration_secs = ndb.IntegerProperty(indexed=False)
   # DateTime indicating lease expiration time.
   lease_expiration_ts = ndb.DateTimeProperty()
+  # Lease ID assigned by Machine Provider.
+  lease_id = ndb.StringProperty(indexed=False)
   # ndb.Key for the MachineType this MachineLease is created for.
   machine_type = ndb.KeyProperty()
   # machine_provider.Dimensions describing the machine.
   mp_dimensions = msgprop.MessageProperty(
       machine_provider.Dimensions, indexed=False)
   # Last request number used.
   request_count = ndb.IntegerProperty(default=0, required=True)
 
 
 class MachineType(ndb.Model):
@@ skipping 140 matching lines @@
     )
     return
 
   machine_lease.client_request_id = None
   machine_lease.hostname = None
   machine_lease.lease_expiration_ts = None
   machine_lease.put()
 
 
 @ndb.transactional
-def log_lease_fulfillment(key, request_id, hostname, lease_expiration_ts):
+def log_lease_fulfillment(
+    key, request_id, hostname, lease_expiration_ts, lease_id):
   """Logs lease fulfillment.
 
   Args:
     request_id: ID of the request being fulfilled.
     hostname: Hostname of the machine fulfilling the request.
     lease_expiration_ts: UTC seconds since epoch when the lease expires.
+    lease_id: ID of the lease assigned by Machine Provider.
   """
   machine_lease = key.get()
   if not machine_lease:
     logging.error('MachineLease does not exist\nKey: %s', key)
     return
 
   if request_id != machine_lease.client_request_id:
     logging.error(
         'Request ID mismatch\nKey: %s\nExpected: %s\nActual: %s',
         key,
         machine_lease.client_request_id,
         request_id,
     )
     return
 
   if (hostname == machine_lease.hostname
-      and lease_expiration_ts == machine_lease.lease_expiration_ts):
+      and lease_expiration_ts == machine_lease.lease_expiration_ts
+      and lease_id == machine_lease.lease_id):
     return
 
   machine_lease.hostname = hostname
   machine_lease.lease_expiration_ts = datetime.datetime.utcfromtimestamp(
       lease_expiration_ts)
+  machine_lease.lease_id = lease_id
   machine_lease.put()
 
 
 @ndb.transactional
 def update_client_request_id(key):
   """Sets the client request ID used to lease a machine.
 
   Args:
     key: ndb.Key for a MachineLease entity.
   """
@@ skipping 90 matching lines @@
           machine_lease.key,
           machine_lease.client_request_id,
           response['hostname'],
           response['lease_expiration_ts'],
       )
       log_lease_fulfillment(
           machine_lease.key,
           machine_lease.client_request_id,
           response['hostname'],
           int(response['lease_expiration_ts']),
+          response['request_hash'],
       )
-      # TODO(smut): Create BotInfo, associate lease information.
   elif state == machine_provider.LeaseRequestState.DENIED:
     logging.warning(
         'Request denied: %s\nRequest ID: %s',
         machine_lease.key,
         machine_lease.client_request_id,
     )
     clear_lease_request(machine_lease.key, machine_lease.client_request_id)
 
 
 def manage_pending_lease_request(machine_lease):
@@ skipping 28 matching lines @@
 
   Args:
     key: ndb.Key for a MachineLease entity.
   """
   machine_lease = key.get()
   if not machine_lease:
     return
 
   # Manage a leased machine.
   if machine_lease.lease_expiration_ts:
+    assert machine_lease.hostname, key
+    bot_info = bot_management.get_info_key(machine_lease.hostname).get()
+    if not (bot_info and bot_info.lease_id and bot_info.lease_expiration_ts):
+      bot_management.bot_event(
+          event_type='bot_leased',
+          bot_id=machine_lease.hostname,
+          external_ip=None,
+          authenticated_as=None,
+          dimensions=None,
+          state=None,
+          version=None,
+          quarantined=False,
+          task_id='',
+          task_name=None,
+          lease_id=machine_lease.lease_id,
+          lease_expiration_ts=machine_lease.lease_expiration_ts,
+      )
+
     if machine_lease.lease_expiration_ts <= utils.utcnow():
       logging.info('MachineLease expired: %s', key)
-      assert machine_lease.hostname, key
       clear_lease_request(key, machine_lease.client_request_id)
     return
 
   # Lease expiration time is unknown, so there must be no leased machine.
   assert not machine_lease.hostname, key
 
   # Manage a pending lease request.
   if machine_lease.client_request_id:
     manage_pending_lease_request(machine_lease)
     return
 
   # Manage an uninitiated lease request.
   if not machine_lease.drained:
     update_client_request_id(key)
     return
 
   # Manage an uninitiated, drained lease request.
   delete_machine_lease(key)