Replace snapshot-based detection with atomic cache lock

Author: mn-ram

openwisp_controller/connection/tasks.py

@@ -2,8 +2,9 @@
 import time
 
 import swapper
-from celery import current_app, shared_task
+from celery import shared_task
 from celery.exceptions import SoftTimeLimitExceeded
+from django.core.cache import cache
 from django.core.exceptions import ObjectDoesNotExist
 from django.utils.translation import gettext_lazy as _
 from swapper import load_model
@@ -13,23 +14,26 @@
 from .exceptions import NoWorkingDeviceConnectionError
 
 logger = logging.getLogger(__name__)
-_TASK_NAME = "openwisp_controller.connection.tasks.update_config"
+_UPDATE_CONFIG_LOCK_KEY = "ow_update_config_{device_id}"
+# Lock timeout (in seconds) acts as a safety net to release the lock
+# in case the task crashes without proper cleanup.
+_UPDATE_CONFIG_LOCK_TIMEOUT = 300
 
 
-def _is_update_in_progress(device_id, current_task_id=None):
-    active = current_app.control.inspect().active()
-    if not active:
-        return False
-    # check if there's any other running task before adding it
-    for task_list in active.values():
-        for task in task_list:
-            if (
-                task["name"] == _TASK_NAME
-                and str(device_id) in task["args"]
-                and task["id"] != current_task_id
-            ):
-                return True
-    return False
+def _acquire_update_config_lock(device_id):
+    """
+    Attempts to atomically acquire a per-device lock using the Django cache.
+    Returns True if the lock was acquired, False if another task already holds it.
+    """
+    lock_key = _UPDATE_CONFIG_LOCK_KEY.format(device_id=device_id)
+    # cache.add is atomic: returns True only if the key doesn't already exist
+    return cache.add(lock_key, True, timeout=_UPDATE_CONFIG_LOCK_TIMEOUT)
+
+
+def _release_update_config_lock(device_id):
+    """Releases the per-device update_config lock."""
+    lock_key = _UPDATE_CONFIG_LOCK_KEY.format(device_id=device_id)
+    cache.delete(lock_key)
 
 
 @shared_task(bind=True)
@@ -52,18 +56,21 @@ def update_config(self, device_id):
     except ObjectDoesNotExist as e:
         logger.warning(f'update_config("{device_id}") failed: {e}')
         return
-    if _is_update_in_progress(device_id, current_task_id=self.request.id):
+    if not _acquire_update_config_lock(device_id):
         logger.info(
             f"update_config for device {device_id} is already in progress, skipping"
         )
         return
     try:
-        device_conn = DeviceConnection.get_working_connection(device)
-    except NoWorkingDeviceConnectionError:
-        return
-    else:
-        logger.info(f"Updating {device} (pk: {device_id})")
-        device_conn.update_config()
+        try:
+            device_conn = DeviceConnection.get_working_connection(device)
+        except NoWorkingDeviceConnectionError:
+            return
+        else:
+            logger.info(f"Updating {device} (pk: {device_id})")
+            device_conn.update_config()
+    finally:
+        _release_update_config_lock(device_id)
 
 
 # task timeout is SSH_COMMAND_TIMEOUT plus a 20% margin

openwisp_controller/connection/tests/test_models.py

@@ -21,7 +21,11 @@
 )
 from ..exceptions import NoWorkingDeviceConnectionError
 from ..signals import is_working_changed
-from ..tasks import _TASK_NAME, _is_update_in_progress, update_config
+from ..tasks import (
+    _acquire_update_config_lock,
+    _release_update_config_lock,
+    update_config,
+)
 from .utils import CreateConnectionsMixin
 
 Config = load_model("config", "Config")
@@ -1026,26 +1030,19 @@ def _assert_applying_conf_test_command(mocked_exec):
     @mock.patch.object(DeviceConnection, "update_config")
     @mock.patch.object(DeviceConnection, "get_working_connection")
     def test_device_update_config_in_progress(
-        self, mocked_get_working_connection, update_config, mocked_sleep
+        self, mocked_get_working_connection, mocked_update_config, mocked_sleep
     ):
         conf = self._prepare_conf_object()
 
-        with mock.patch("celery.app.control.Inspect.active") as mocked_active:
-            mocked_active.return_value = {
-                "task": [
-                    {
-                        "name": _TASK_NAME,
-                        "args": [str(conf.device.pk)],
-                        "id": "other-task-id",
-                    }
-                ]
-            }
+        with mock.patch(
+            "openwisp_controller.connection.tasks._acquire_update_config_lock",
+            return_value=False,
+        ):
             conf.config = {"general": {"timezone": "UTC"}}
             conf.full_clean()
             conf.save()
-            mocked_active.assert_called_once()
             mocked_get_working_connection.assert_not_called()
-            update_config.assert_not_called()
+            mocked_update_config.assert_not_called()
 
     @mock.patch("time.sleep")
     @mock.patch.object(DeviceConnection, "update_config")
@@ -1058,53 +1055,30 @@ def test_device_update_config_not_in_progress(
             conf.device.deviceconnection_set.first()
         )
 
-        with mock.patch("celery.app.control.Inspect.active") as mocked_active:
-            mocked_active.return_value = {
-                "task": [{"name": _TASK_NAME, "args": ["..."], "id": "other-task-id"}]
-            }
+        with mock.patch(
+            "openwisp_controller.connection.tasks._acquire_update_config_lock",
+            return_value=True,
+        ), mock.patch(
+            "openwisp_controller.connection.tasks._release_update_config_lock",
+        ):
             conf.config = {"general": {"timezone": "UTC"}}
             conf.full_clean()
             conf.save()
-            mocked_active.assert_called_once()
             mocked_get_working_connection.assert_called_once()
             mocked_update_config.assert_called_once()
 
-    def test_is_update_in_progress_ignores_current_task(self):
-        """Regression test: _is_update_in_progress must not count
-        the calling task itself as a duplicate."""
+    def test_acquire_update_config_lock(self):
+        """Test that the lock can be acquired and prevents duplicate acquisition."""
         device_id = "test-device-id"
-        current_task_id = "current-task-id"
-        with mock.patch("celery.app.control.Inspect.active") as mocked_active:
-            mocked_active.return_value = {
-                "worker": [
-                    {
-                        "name": _TASK_NAME,
-                        "args": [device_id],
-                        "id": current_task_id,
-                    }
-                ]
-            }
-            result = _is_update_in_progress(device_id, current_task_id=current_task_id)
-            self.assertFalse(result)
-
-    def test_is_update_in_progress_detects_other_task(self):
-        """_is_update_in_progress returns True when another task
-        for the same device is active."""
-        device_id = "test-device-id"
-        with mock.patch("celery.app.control.Inspect.active") as mocked_active:
-            mocked_active.return_value = {
-                "worker": [
-                    {
-                        "name": _TASK_NAME,
-                        "args": [device_id],
-                        "id": "other-task-id",
-                    }
-                ]
-            }
-            result = _is_update_in_progress(
-                device_id, current_task_id="current-task-id"
-            )
-            self.assertTrue(result)
+        # First acquisition should succeed
+        self.assertTrue(_acquire_update_config_lock(device_id))
+        # Second acquisition should fail (lock already held)
+        self.assertFalse(_acquire_update_config_lock(device_id))
+        # After releasing, acquisition should succeed again
+        _release_update_config_lock(device_id)
+        self.assertTrue(_acquire_update_config_lock(device_id))
+        # Cleanup
+        _release_update_config_lock(device_id)
 
     @mock.patch(_connect_path)
     def test_schedule_command_called(self, connect_mocked):