Cache update worker: send broadcast WebSockets message after update

This can be disabled by setting SPATIAL_UPDATE_CLIENT_NOTIFICATIONS =
False in the settings.

Author: tomka

CHANGELOG.md

@@ -12,6 +12,10 @@ Miscellaneous:
 
 - The skeleton reroot API now also returns the new root node's edition time.
 
+- Grid cache worker: send broadcast message after updates using WebSockets. This
+  can be disabled by settin `SPATIAL_UPDATE_CLIENT_NOTIFICATIONS = False` in
+  `settings.py`.
+
 ### Bug fixes
 
 

django/applications/catmaid/consumers.py

@@ -11,6 +11,9 @@
 logger = logging.getLogger(__name__)
 
 
+broadcast_group_name = 'broadcast'
+
+
 def get_user_group_name(user_id):
     return f"updates-{user_id}"
 
@@ -24,8 +27,9 @@ def connect(self):
         if not self.channel_layer:
             logger.error("UpdateConsumer: can't handle WebSockets connection, no channels layer")
             return
-        # Add user to the matching user group
+        # Add user to the matching user group and to broadcast group
         user = self.scope["user"]
+        async_to_sync(self.channel_layer.group_add)(broadcast_group_name, self.channel_name)
         async_to_sync(self.channel_layer.group_add)(get_user_group_name(user.id), self.channel_name)
         self.accept()
 
@@ -37,6 +41,7 @@ def disconnect(self, message):
             logger.error("UpdateConsumer: can't handle WebSockets disconnect, no channels layer")
             return
         user = self.scope["user"]
+        async_to_sync(self.channel_layer.group_discard)(broadcast_group_name, self.channel_name)
         async_to_sync(self.channel_layer.group_discard)(get_user_group_name(user.id), self.channel_name)
 
     def receive(self, *, text_data):
@@ -125,3 +130,43 @@ def publish_message_to_broker(message, routing_key):
             routing_key=routing_key,
             retry=False,
         )
+
+
+def msg_all(event_name, data:str="", data_type:str="text", is_raw_data:bool=False,
+        ignore_missing:bool=True) -> None:
+    """Send a message to a user. This message will contain a dictionary with the
+    field <data_type> with content <data> if raw data is requested, otherwise
+    with a dictionary. Its fields are "event" for the <event_name> and "payload"
+    for <data>."""
+    if is_raw_data:
+        payload = data
+    else:
+        payload = json.dumps({
+            "event": event_name,
+            "payload": data
+        })
+    # Broadcast to listening sockets
+    try:
+        logger.info('attempting send broadcast msg')
+        if is_in_celery_task():
+            logger.info('sending broadcast msg from celery task')
+            publish_message_to_broker({
+                'type': 'user.message',
+                'data': payload,
+            }, broadcast_group_name)
+        else:
+            channel_layer = get_channel_layer()
+            # Without any channel layer, there is no point in trying to send a
+            # message.
+            if not channel_layer:
+                return
+            logger.info('sending msg')
+            async_to_sync(channel_layer.group_send)(broadcast_group_name, {
+                'type': 'user.message',
+                'data': payload,
+            })
+    except KeyError as e:
+        if ignore_missing:
+            pass
+        else:
+            raise e

django/applications/catmaid/management/commands/catmaid_cache_update_worker.py

@@ -11,6 +11,7 @@
 from django.core.management.base import BaseCommand
 from django.db import connection
 
+from catmaid.consumers import msg_all
 from catmaid.control.edge import get_intersected_grid_cells
 from catmaid.control.node import (get_configured_node_providers,
         GridCachedNodeProvider, Postgis3dNodeProvider, update_grid_cell)
@@ -45,6 +46,7 @@ def update(self, updates, cursor):
         # Batch of grids to update during one run
         updated_cells = 0
         updated_grids = set()
+        updated_cell_refs = []
         for update in updates:
             g = grid_map[update['grid_id']]
             w_i, h_i, d_i = update['x'], update['y'], update['z']
@@ -61,21 +63,34 @@ def update(self, updates, cursor):
             if g.hidden_last_editor_id:
                 params['hidden_last_editor_id'] = int(g.hidden_last_editor_id)
 
-            added = update_grid_cell(g.project_id, g.id, w_i, h_i, d_i,
+            updated = update_grid_cell(g.project_id, g.id, w_i, h_i, d_i,
                     g.cell_width, g.cell_height, g.cell_depth,
                     params, g.allow_empty, g.n_lod_levels, g.lod_min_bucket_size,
                     g.lod_strategy, g.has_json_data, g.has_json_text_data,
                     g.has_msgpack_data, provider=provider, cursor=cursor,
                     delete_empty=True)
 
-            if added:
+            if updated:
                 updated_cells += 1
                 updated_grids.add(g.id)
+                updated_cell_refs.add([{
+                    'project_id': g.project_id,
+                    'grid_id', g.id,
+                    'x': w_i,
+                    'y': h_i,
+                    'z': d_i,
+                })
 
                 # TODO: delete in batches
                 DirtyNodeGridCacheCell.objects.filter(grid_id=g.id, x_index=w_i,
                         y_index=h_i, z_index=d_i).delete()
 
+        # Optionally, let users know about updated cells
+        if settings.SPATIAL_UPDATE_CLIENT_NOTIFICATIONS:
+            msg_all('grid-cache-update', {
+                'updated_cells': updated_cell_refs,
+            })
+
         logger.debug(f'Updated {updated_cells} grid cell(s) in {len(updated_grids)} grid cache(s)')
 
 

django/applications/catmaid/static/js/init.js

@@ -1570,6 +1570,11 @@ var project;
           CATMAID.msg("New message", payload.message_title);
           client.check_messages(true);
         }
+      ], [
+        'grid-cache-update',
+        function (client, payload) {
+          // Ignore for now
+        }
       ], [
         'unknown',
         function(message) {

django/projects/mysite/settings_base.py

@@ -535,6 +535,10 @@
 # connector links).
 SPATIAL_UPDATE_NOTIFICATIONS = False
 
+# Specifies whether grid cache update workser will notify all clients about
+# recomputed dirty cells using WebSockets.
+SPATIAL_UPDATE_CLIENT_NOTIFICATIONS = True
+
 # On statup, the default client instance settings can be populated based on a
 # JSON string, representing a list of objects with a "key" field and a "value"
 # field. These settings will only be applied if they exist already.