Merge "Use callback to update SegmentHostMapping"

neutron/db/agents_db.py

@@ -43,7 +43,6 @@ from neutron.db import model_base
 from neutron.extensions import agent as ext_agent
 from neutron.extensions import availability_zone as az_ext
 from neutron import manager
-from neutron.services.segments import db as segments_db
 
 LOG = logging.getLogger(__name__)
 
@@ -385,6 +384,7 @@ class AgentDbMixin(ext_agent.AgentPluginBase, AgentAvailabilityZoneMixin):
                 greenthread.sleep(0)
                 self._log_heartbeat(agent_state, agent_db, configurations_dict)
                 agent_db.update(res)
+                event_type = events.AFTER_UPDATE
             except ext_agent.AgentNotFoundByTypeHost:
                 greenthread.sleep(0)
                 res['created_at'] = current_time
@@ -394,11 +394,14 @@ class AgentDbMixin(ext_agent.AgentPluginBase, AgentAvailabilityZoneMixin):
                 agent_db = Agent(**res)
                 greenthread.sleep(0)
                 context.session.add(agent_db)
+                event_type = events.AFTER_CREATE
                 self._log_heartbeat(agent_state, agent_db, configurations_dict)
                 status = n_const.AGENT_NEW
             greenthread.sleep(0)
-            segments_db.update_segment_host_mapping_for_agent(
-                context, agent_state['host'], self, agent_state)
+
+        registry.notify(resources.AGENT, event_type, self, context=context,
+                        host=agent_state['host'], plugin=self,
+                        agent=agent_state)
         return status
 
     def create_or_update_agent(self, context, agent):

neutron/db/migration/models/head.py

@@ -58,6 +58,7 @@ from neutron.plugins.ml2.drivers import type_vlan  # noqa
 from neutron.plugins.ml2.drivers import type_vxlan  # noqa
 from neutron.plugins.ml2 import models  # noqa
 from neutron.services.auto_allocate import models  # noqa
+from neutron.services.segments import db # noqa
 from neutron.services.trunk import models  # noqa
 
 

neutron/services/segments/db.py

@@ -26,6 +26,10 @@ from sqlalchemy import orm
 from sqlalchemy.orm import exc
 
 from neutron.api.v2 import attributes
+from neutron.callbacks import events
+from neutron.callbacks import exceptions as c_exc
+from neutron.callbacks import registry
+from neutron.callbacks import resources
 from neutron.db import common_db_mixin
 from neutron.db import model_base
 from neutron.db import segments_db as db
@@ -181,7 +185,8 @@ def _get_phys_nets(agent):
 reported_hosts = set()
 
 
-def update_segment_host_mapping_for_agent(context, host, plugin, agent):
+def _update_segment_host_mapping_for_agent(resource, event, trigger,
+                                           context, host, plugin, agent):
     check_segment_for_agent = getattr(plugin, 'check_segment_for_agent', None)
     if not check_segment_for_agent:
         return
@@ -199,3 +204,14 @@ def update_segment_host_mapping_for_agent(context, host, plugin, agent):
             segment['id'] for segment in segments
             if check_segment_for_agent(segment, agent)}
         update_segment_host_mapping(context, host, current_segment_ids)
+
+
+def subscribe():
+    registry.subscribe(_update_segment_host_mapping_for_agent,
+                       resources.AGENT,
+                       events.AFTER_CREATE)
+    registry.subscribe(_update_segment_host_mapping_for_agent,
+                       resources.AGENT,
+                       events.AFTER_UPDATE)
+
+subscribe()

neutron/tests/unit/extensions/test_segment.py

@@ -283,6 +283,7 @@ class HostSegmentMappingTestCase(SegmentTestCase):
         if not plugin:
             plugin = 'neutron.plugins.ml2.plugin.Ml2Plugin'
         super(HostSegmentMappingTestCase, self).setUp(plugin=plugin)
+        db.subscribe()
 
     def _get_segments_for_host(self, host):
         ctx = context.get_admin_context()