Improve performance for checking hosts AZs

When updating aggregate metadata or when adding a host into aggregate,
a check is done to make sure that hosts can't be in multiple AZs.

In order to do the check, the routine was looping over each host to look
in the DB which AZs the host is part of.
As it's quite hungry in terms of DB accesses, here the proposal is to
extend the helper method for getting all AZs by also having their hosts,
and so only do memory iterations.

Partial-Bug: #1277230

Change-Id: I7c5bc980acb17f0a0e1f8c1403dac66c305866fb

nova/availability_zones.py

@@ -103,32 +103,50 @@ def update_host_availability_zone_cache(context, host, availability_zone=None):
     cache.set(cache_key, availability_zone, AZ_CACHE_SECONDS)
 
 
-def get_availability_zones(context, get_only_available=False):
+def get_availability_zones(context, get_only_available=False,
+                           with_hosts=False):
     """Return available and unavailable zones on demand.
 
-       :param get_only_available: flag to determine whether to return
-           available zones only, default False indicates return both
-           available zones and not available zones, True indicates return
-           available zones only
+        :param get_only_available: flag to determine whether to return
+            available zones only, default False indicates return both
+            available zones and not available zones, True indicates return
+            available zones only
+        :param with_hosts: whether to return hosts part of the AZs
+        :type with_hosts: bool
     """
     enabled_services = db.service_get_all(context, False)
     enabled_services = set_availability_zones(context, enabled_services)
 
     available_zones = []
-    for zone in [service['availability_zone'] for service
-                 in enabled_services]:
-        if zone not in available_zones:
+    for (zone, host) in [(service['availability_zone'], service['host'])
+                         for service in enabled_services]:
+        if not with_hosts and zone not in available_zones:
             available_zones.append(zone)
+        elif with_hosts:
+            _available_zones = dict(available_zones)
+            zone_hosts = _available_zones.setdefault(zone, set())
+            zone_hosts.add(host)
+            # .items() returns a view in Py3, casting it to list for Py2 compat
+            available_zones = list(_available_zones.items())
 
     if not get_only_available:
         disabled_services = db.service_get_all(context, True)
         disabled_services = set_availability_zones(context, disabled_services)
         not_available_zones = []
-        zones = [service['availability_zone'] for service in disabled_services
-                if service['availability_zone'] not in available_zones]
-        for zone in zones:
-            if zone not in not_available_zones:
+        azs = available_zones if not with_hosts else dict(available_zones)
+        zones = [(service['availability_zone'], service['host'])
+                 for service in disabled_services
+                 if service['availability_zone'] not in azs]
+        for (zone, host) in zones:
+            if not with_hosts and zone not in not_available_zones:
                 not_available_zones.append(zone)
+            elif with_hosts:
+                _not_available_zones = dict(not_available_zones)
+                zone_hosts = _not_available_zones.setdefault(zone, set())
+                zone_hosts.add(host)
+                # .items() returns a view in Py3, casting it to list for Py2
+                #   compat
+                not_available_zones = list(_not_available_zones.items())
         return (available_zones, not_available_zones)
     else:
         return available_zones

nova/compute/api.py

@@ -3362,9 +3362,17 @@ class AggregateAPI(base.Base):
         """
         if 'availability_zone' in metadata:
             _hosts = hosts or aggregate.hosts
+            zones, not_zones = availability_zones.get_availability_zones(
+                context, with_hosts=True)
             for host in _hosts:
-                host_az = availability_zones.get_host_availability_zone(
-                    context, host)
+                # NOTE(sbauza): Host can only be in one AZ, so let's take only
+                #               the first element
+                host_azs = [az for (az, az_hosts) in zones
+                            if host in az_hosts
+                            and az != CONF.internal_service_availability_zone]
+                host_az = host_azs.pop()
+                if host_azs:
+                    LOG.warning(_("More than 1 AZ for host %s"), host)
                 if host_az == CONF.default_availability_zone:
                     # NOTE(sbauza): Aggregate with AZ set to default AZ can
                     #               exist, we need to check

nova/tests/compute/test_compute.py

@@ -9495,8 +9495,20 @@ class ComputeAPIAggrTestCase(BaseTestCase):
 
         self.mox.StubOutWithMock(availability_zones,
                                  'update_host_availability_zone_cache')
-        availability_zones.update_host_availability_zone_cache(self.context,
-                                                               fake_host)
+
+        def _stub_update_host_avail_zone_cache(host, az=None):
+            if az is not None:
+                availability_zones.update_host_availability_zone_cache(
+                    self.context, host, az)
+            else:
+                availability_zones.update_host_availability_zone_cache(
+                    self.context, host)
+
+        for avail_zone, hosts in six.iteritems(values):
+            for host in hosts:
+                _stub_update_host_avail_zone_cache(
+                    host, CONF.default_availability_zone)
+        _stub_update_host_avail_zone_cache(fake_host)
         self.mox.ReplayAll()
 
         fake_notifier.NOTIFICATIONS = []

nova/tests/test_availability_zones.py

@@ -226,6 +226,14 @@ class AvailabilityZoneTestCases(test.TestCase):
 
         self.assertEqual(zones, ['nova-test', 'nova-test2'])
 
+        zones, not_zones = az.get_availability_zones(self.context,
+                                                     with_hosts=True)
+
+        self.assertEqual(zones, [(u'nova-test2', set([u'host3'])),
+                                 (u'nova-test', set([u'host1']))])
+        self.assertEqual(not_zones, [(u'nova-test3', set([u'host4'])),
+                                     (u'nova', set([u'host5']))])
+
     def test_get_instance_availability_zone_default_value(self):
         """Test get right availability zone by given an instance."""
         fake_inst_id = 162