get_hosts_mapped_with_segments add filter agt_type

Extend the get_hosts_mapped_with_segments method to add
optional filters to include/exclude based on agent type.
Uses a joined query, when both include and exclude
filtering is used togheter the exclude filter is most
significant.

Partial-Bug: #2040172

Change-Id: I2cfd52a2657fad989e24e974fda470ecd960262b
Signed-off-by: Harald Jensås <hjensas@redhat.com>

neutron/services/segments/db.py

@@ -31,6 +31,8 @@ from oslo_log import helpers as log_helpers
 from oslo_log import log as logging
 from oslo_utils import uuidutils
 
+from neutron.db.models import agent as agent_model
+from neutron.db.models import segment as segment_model
 from neutron.db import segments_db as db
 from neutron import manager
 from neutron.objects import base as base_obj
@@ -243,14 +245,43 @@ def update_segment_host_mapping(context, host, current_segment_ids):
                               entry.segment_id, entry.host)
 
 
-def get_hosts_mapped_with_segments(context):
+def get_hosts_mapped_with_segments(context, include_agent_types=None,
+                                   exclude_agent_types=None):
     """Get hosts that are mapped with segments.
 
     L2 providers can use this method to get an overview of SegmentHostMapping,
     and then delete the stale SegmentHostMapping.
+
+    When using both include_agent_types and exclude_agent_types,
+    exclude_agent_types is most significant.
+    All hosts without agent are excluded when using any agent_type filter.
+
+    :param context: current running context information
+    :param include_agent_types: (set) List of agent types, include hosts
+        with matching agents.
+    :param exclude_agent_types: (set) List of agent types, exclude hosts
+        with matching agents.
     """
-    segment_host_mapping = network.SegmentHostMapping.get_objects(context)
+    def add_filter_by_agent_types(qry, include, exclude):
-    return {row.host for row in segment_host_mapping}
+        qry = qry.join(
+            agent_model.Agent,
+            segment_model.SegmentHostMapping.host == agent_model.Agent.host)
+        if include:
+            qry = qry.filter(agent_model.Agent.agent_type.in_(include))
+        if exclude:
+            qry = qry.filter(agent_model.Agent.agent_type.not_in(exclude))
+
+        return qry
+
+    with db_api.CONTEXT_READER.using(context):
+        query = context.session.query(segment_model.SegmentHostMapping)
+        if include_agent_types or exclude_agent_types:
+            query = add_filter_by_agent_types(query, include_agent_types,
+                                              exclude_agent_types)
+
+        res = query.all()
+
+    return {row.host for row in res}
 
 
 def _get_phys_nets(agent):

neutron/tests/unit/extensions/test_segment.py

@@ -768,10 +768,37 @@ class TestMl2HostSegmentMappingNoAgent(HostSegmentMappingTestCase):
         actual_hosts = db.get_hosts_mapped_with_segments(ctx)
         self.assertEqual(hosts, actual_hosts)
 
+    def test_get_all_hosts_mapped_with_segments_agent_type_filter(self):
+        ctx = context.get_admin_context()
+        hosts = set()
+        with self.network() as network:
+            network_id = network['network']['id']
+        for i in range(1, 3):
+            host = "host%s" % i
+            segment = self._test_create_segment(
+                network_id=network_id, physical_network='physnet%s' % i,
+                segmentation_id=200 + i, network_type=constants.TYPE_VLAN)
+            db.update_segment_host_mapping(
+                ctx, host, {segment['segment']['id']})
+            hosts.add(host)
+
+        # Now they are 2 hosts with segment being mapped.
+        #  host1 does not have an agent
+        #  host2 does not have an agent
+        # Any agent_type filter excludes hosts that does not have an agent
+        actual_hosts = db.get_hosts_mapped_with_segments(
+            ctx, exclude_agent_types={'fake-agent-type'})
+        self.assertEqual(set(), actual_hosts)
+        actual_hosts = db.get_hosts_mapped_with_segments(
+            ctx, include_agent_types={'fake-agent-type'})
+        self.assertEqual(set(), actual_hosts)
+
 
 class TestMl2HostSegmentMappingOVS(HostSegmentMappingTestCase):
     _mechanism_drivers = ['openvswitch', 'logger']
     mock_path = 'neutron.services.segments.db.update_segment_host_mapping'
+    agent_type_a = constants.AGENT_TYPE_OVS
+    agent_type_b = constants.AGENT_TYPE_LINUXBRIDGE
 
     def test_new_agent(self):
         host = 'host1'
@@ -891,9 +918,118 @@ class TestMl2HostSegmentMappingOVS(HostSegmentMappingTestCase):
         self.assertFalse(segments_host_db)
         self.assertFalse(mock.mock_calls)
 
+    def test_get_all_hosts_mapped_with_segments(self):
+        ctx = context.get_admin_context()
+        hosts = set()
+        with self.network() as network:
+            network_id = network['network']['id']
+        for i in range(1, 3):
+            host = "host%s" % i
+            segment = self._test_create_segment(
+                network_id=network_id, physical_network='physnet%s' % i,
+                segmentation_id=200 + i, network_type=constants.TYPE_VLAN)
+            self._register_agent(host, mappings={'physnet%s' % i: 'br-eth-1'},
+                                 plugin=self.plugin)
+            db.update_segment_host_mapping(
+                ctx, host, {segment['segment']['id']})
+            hosts.add(host)
+
+        # Now they are 2 hosts with segment being mapped.
+        actual_hosts = db.get_hosts_mapped_with_segments(ctx)
+        self.assertEqual(hosts, actual_hosts)
+
+    def test_get_all_hosts_mapped_with_segments_agent_type_filters(self):
+        ctx = context.get_admin_context()
+        with self.network() as network:
+            network_id = network['network']['id']
+        for i in range(1, 3):
+            host = "host%s" % i
+            segment = self._test_create_segment(
+                network_id=network_id, physical_network='physnet%s' % i,
+                segmentation_id=200 + i, network_type=constants.TYPE_VLAN)
+            if i == 2:
+                agent_type = self.agent_type_a
+            else:
+                agent_type = self.agent_type_b
+            helpers.register_ovs_agent(
+                host, agent_type=agent_type,
+                bridge_mappings={'physnet%s' % i: 'br-eth-1'},
+                plugin=self.plugin, start_flag=True)
+            db.update_segment_host_mapping(
+                ctx, host, {segment['segment']['id']})
+
+        # Now they are 2 hosts with segment being mapped.
+        #   host1 is agent_type_b
+        #   host2 is agent_type_a
+        # get all hosts (host1 and host2) when not using any filtering
+        actual_hosts = db.get_hosts_mapped_with_segments(ctx)
+        self.assertEqual({"host1", "host2"}, actual_hosts)
+        # get host1 when exclude agent_type_a agents
+        actual_hosts = db.get_hosts_mapped_with_segments(
+            ctx, exclude_agent_types={self.agent_type_a})
+        self.assertEqual({"host1"}, actual_hosts)
+        # get host2 when exclude agent_type_b agents
+        actual_hosts = db.get_hosts_mapped_with_segments(
+            ctx, exclude_agent_types={self.agent_type_b})
+        self.assertEqual({"host2"}, actual_hosts)
+        # get host2 when include agent_type_a agents
+        actual_hosts = db.get_hosts_mapped_with_segments(
+            ctx, include_agent_types={self.agent_type_a})
+        self.assertEqual({"host2"}, actual_hosts)
+        # get host1 when include agent_type_b agents
+        actual_hosts = db.get_hosts_mapped_with_segments(
+            ctx, include_agent_types={self.agent_type_b})
+        self.assertEqual({"host1"}, actual_hosts)
+        # get host1 and host2 when include both agent_type_a and agent_type_b
+        actual_hosts = db.get_hosts_mapped_with_segments(
+            ctx, include_agent_types={self.agent_type_b, self.agent_type_a})
+        self.assertEqual({"host1", "host2"}, actual_hosts)
+        # When using both include and exclude, exclude is most significant
+        actual_hosts = db.get_hosts_mapped_with_segments(
+            ctx,
+            include_agent_types={self.agent_type_b, self.agent_type_a},
+            exclude_agent_types={self.agent_type_b}
+        )
+        self.assertEqual({"host2"}, actual_hosts)
+        # include and exclude both agent types - exclude is most significant
+        actual_hosts = db.get_hosts_mapped_with_segments(
+            ctx,
+            include_agent_types={self.agent_type_b, self.agent_type_a},
+            exclude_agent_types={self.agent_type_b, self.agent_type_a}
+        )
+        self.assertEqual(set(), actual_hosts)
+
+    def test_get_all_hosts_mapped_with_segments_agent_type_filter(self):
+        ctx = context.get_admin_context()
+        hosts = set()
+        with self.network() as network:
+            network_id = network['network']['id']
+        for i in range(1, 3):
+            host = "host%s" % i
+            segment = self._test_create_segment(
+                network_id=network_id, physical_network='physnet%s' % i,
+                segmentation_id=200 + i, network_type=constants.TYPE_VLAN)
+            self._register_agent(host, mappings={'physnet%s' % i: 'br-eth-1'},
+                                 plugin=self.plugin)
+            db.update_segment_host_mapping(
+                ctx, host, {segment['segment']['id']})
+            hosts.add(host)
+
+        # Now they are 2 hosts with segment being mapped.
+        #  host1 is agent_type_a
+        #  host2 is agent_type_a
+        actual_hosts = db.get_hosts_mapped_with_segments(
+            ctx, exclude_agent_types={self.agent_type_a})
+        self.assertEqual(set(), actual_hosts)
+        actual_hosts = db.get_hosts_mapped_with_segments(
+            ctx, include_agent_types={self.agent_type_a})
+        self.assertEqual(hosts, actual_hosts)
+
 
 class TestMl2HostSegmentMappingLinuxBridge(TestMl2HostSegmentMappingOVS):
     _mechanism_drivers = ['linuxbridge', 'logger']
+    agent_type_a = constants.AGENT_TYPE_LINUXBRIDGE
+    agent_type_b = constants.AGENT_TYPE_OVS
 
     def setUp(self, plugin=None):
         cfg.CONF.set_override(c_experimental.EXPERIMENTAL_LINUXBRIDGE, True,
@@ -908,6 +1044,8 @@ class TestMl2HostSegmentMappingLinuxBridge(TestMl2HostSegmentMappingOVS):
 
 class TestMl2HostSegmentMappingMacvtap(TestMl2HostSegmentMappingOVS):
     _mechanism_drivers = ['macvtap', 'logger']
+    agent_type_a = constants.AGENT_TYPE_MACVTAP
+    agent_type_b = constants.AGENT_TYPE_OVS
 
     def _register_agent(self, host, mappings=None, plugin=None):
         helpers.register_macvtap_agent(host=host, interface_mappings=mappings,
@@ -916,6 +1054,8 @@ class TestMl2HostSegmentMappingMacvtap(TestMl2HostSegmentMappingOVS):
 
 class TestMl2HostSegmentMappingSriovNicSwitch(TestMl2HostSegmentMappingOVS):
     _mechanism_drivers = ['sriovnicswitch', 'logger']
+    agent_type_a = constants.AGENT_TYPE_NIC_SWITCH
+    agent_type_b = constants.AGENT_TYPE_OVS
 
     def _register_agent(self, host, mappings=None, plugin=None):
         helpers.register_sriovnicswitch_agent(host=host,