[Fixes #13382] Harvesting scheduler: Fixing expiration issue of the discovery task (#13383)

[Fixes #13382] adding dynamic expiration time to the discovery subtasks

Author: Gpetrak

geonode/harvesting/models.py

@@ -368,7 +368,7 @@ def initiate(self, harvestable_resource_ids: typing.Optional[typing.List[int]] =
             raise RuntimeError("Invalid selection")
         self.status = self.STATUS_PENDING
         self.save()
-        task_signature.apply_async(args=(), expiration=30)
+        task_signature.apply_async(args=())
 
     def abort(self):
         """Abort a pending or on-going session."""

geonode/harvesting/tasks.py

@@ -567,7 +567,6 @@ def check_harvester_available(self, harvester_id: int):
 @app.task(
     bind=True,
     queue="geonode",
-    expires=30,
     time_limit=600,
     acks_late=False,
     ignore_result=False,
@@ -578,49 +577,75 @@ def update_harvestable_resources(self, refresh_session_id: int):
     # of resources because these have potentially been individually selected by the
     # user, which means we are not interested in all of them
     session = models.AsynchronousHarvestingSession.objects.get(pk=refresh_session_id)
-    if session.status != session.STATUS_ABORTED:
-        session.status = session.STATUS_ON_GOING
-        session.save()
-        harvester = session.harvester
-        if harvester.update_availability():
-            harvester.status = harvester.STATUS_UPDATING_HARVESTABLE_RESOURCES
-            harvester.save()
-            worker = harvester.get_harvester_worker()
-            try:
-                num_resources = worker.get_num_available_resources()
-            except (NotImplementedError, base.HarvestingException) as exc:
-                _handle_harvestable_resources_update_error(
-                    self.request.id, refresh_session_id=refresh_session_id, raised_exception=exc
-                )
-            else:
-                harvester.num_harvestable_resources = num_resources
-                harvester.save()
-                session.total_records_to_process = num_resources
-                session.save()
-                page_size = 10
-                total_pages = math.ceil(num_resources / page_size)
-                batches = []
-                for page in range(total_pages):
-                    batches.append(
-                        _update_harvestable_resources_batch.signature(
-                            args=(refresh_session_id, page, page_size),
-                        )
-                    )
-                update_finalizer = _finish_harvestable_resources_update.signature(
-                    args=(refresh_session_id,), immutable=True
-                ).on_error(
-                    _handle_harvestable_resources_update_error.signature(
-                        kwargs={"refresh_session_id": refresh_session_id}
-                    )
-                )
-                update_workflow = chord(batches, body=update_finalizer)
-                update_workflow.apply_async(args=(), expiration=30)
-        else:
+
+    if session.status == session.STATUS_ABORTED:
+        logger.debug("Session has been aborted, skipping...")
+        return
+
+    session.status = session.STATUS_ON_GOING
+    session.save()
+
+    harvester = session.harvester
+
+    try:
+
+        if not harvester.update_availability():
             finish_asynchronous_session(
                 refresh_session_id, session.STATUS_FINISHED_ALL_FAILED, final_details="Harvester is not available"
             )
-    else:
-        logger.debug("Session has been aborted, skipping...")
+            return
+
+        harvester.status = harvester.STATUS_UPDATING_HARVESTABLE_RESOURCES
+        harvester.save()
+
+        worker = harvester.get_harvester_worker()
+        try:
+            num_resources = worker.get_num_available_resources()
+            # definition of the dynamic expiration time
+            task_expiration_time = calculate_dynamic_expiration(num_resources)
+        except (NotImplementedError, base.HarvestingException) as exc:
+            logger.exception("Failed to get number of available resources for harvester %s", harvester.id)
+            details = f"Failed to get number of available resources: {exc}"
+            finish_asynchronous_session(
+                refresh_session_id,
+                models.AsynchronousHarvestingSession.STATUS_FINISHED_ALL_FAILED,
+                final_details=details,
+            )
+            return
+
+        harvester.num_harvestable_resources = num_resources
+        harvester.save()
+        session.total_records_to_process = num_resources
+        session.save()
+        page_size = 10
+        total_pages = math.ceil(num_resources / page_size)
+        batches = []
+        for page in range(total_pages):
+            batches.append(
+                _update_harvestable_resources_batch.signature(
+                    args=(refresh_session_id, page, page_size),
+                ).set(expires=task_expiration_time)
+            )
+        update_finalizer = (
+            _finish_harvestable_resources_update.signature(args=(refresh_session_id,), immutable=True)
+            .on_error(
+                _handle_harvestable_resources_update_error.signature(kwargs={"refresh_session_id": refresh_session_id})
+            )
+            .set(expires=task_expiration_time)
+        )
+
+        update_workflow = chord(batches, body=update_finalizer)
+        update_workflow.apply_async(args=(), expires=task_expiration_time)
+        logger.info(f"Applying harvesting chord with expiration in {task_expiration_time} seconds")
+
+    except Exception as exc:
+        logger.exception("Unexpected error during update_harvestable_resources for harvester %s", harvester.id)
+        details = f"Unexpected error: {exc}"
+        finish_asynchronous_session(
+            refresh_session_id,
+            models.AsynchronousHarvestingSession.STATUS_FINISHED_ALL_FAILED,
+            final_details=details,
+        )
 
 
 @app.task(

geonode/harvesting/tests/test_tasks.py

@@ -188,28 +188,59 @@ def test_check_harvester_available(self, mock_harvester_model):
         tasks.check_harvester_available(1000)
         mock_harvester.update_availability.assert_called()
 
+    @mock.patch("geonode.harvesting.tasks.calculate_dynamic_expiration")
     @mock.patch("geonode.harvesting.tasks._handle_harvestable_resources_update_error")
     @mock.patch("geonode.harvesting.tasks._finish_harvestable_resources_update")
     @mock.patch("geonode.harvesting.tasks._update_harvestable_resources_batch")
     @mock.patch("geonode.harvesting.tasks.chord")
     @mock.patch("geonode.harvesting.tasks.models")
     def test_update_harvestable_resources_sends_batched_requests(
-        self, mock_models, mock_chord, mock_batch, mock_finalizer, mock_error_handler
+        self,
+        mock_models,
+        mock_chord,
+        mock_batch,
+        mock_finalizer,
+        mock_error_handler,
+        mock_calculate_expiration,
     ):
-        """Verify that the `update_harvestable_resources` task creates a celery chord with the batched task, a finalizer and an error handler."""
+        """Verify that `update_harvestable_resources` creates a celery chord with expiration."""
+
+        # Set up mock harvester and session
         mock_worker = mock.MagicMock()
         mock_worker.get_num_available_resources.return_value = 1
-        mock_harvester = mock.MagicMock(models.Harvester)
-        mock_models.Harvester.objects.get.return_value = mock_harvester
+
+        mock_harvester = mock.MagicMock()
         mock_harvester.get_harvester_worker.return_value = mock_worker
+        mock_harvester.update_availability.return_value = True
+
+        mock_session = mock.MagicMock()
+        mock_session.status = models.AsynchronousHarvestingSession.STATUS_ON_GOING
+        mock_session.harvester = mock_harvester
+
+        mock_models.AsynchronousHarvestingSession.objects.get.return_value = mock_session
+
+        mock_calculate_expiration.return_value = 123
+
+        # Setup the finalizer signature and its chained call
+        mock_finalizer_sig = mock.MagicMock(name="finalizer_signature")
+        mock_finalizer_on_error = mock.MagicMock(name="finalizer_on_error")
+        mock_finalizer_sig.on_error.return_value = mock_finalizer_on_error
+        mock_finalizer.signature.return_value = mock_finalizer_sig
+
+        # Run the task
+        tasks.update_harvestable_resources("fake_session_id")
+
+        # Assert expiration was calculated
+        mock_calculate_expiration.assert_called_once_with(1)
+
+        # Assert batch task had expiration set
+        mock_batch.signature.return_value.set.assert_any_call(expires=123)
 
-        tasks.update_harvestable_resources("fake harvester id")
+        # Assert finalizer had expiration set via the on_error chain
+        mock_finalizer_on_error.set.assert_called_once_with(expires=123)
 
-        mock_batch.signature.assert_called()
-        mock_finalizer.signature.assert_called()
-        mock_error_handler.signature.assert_called()
-        mock_chord.assert_called()
-        mock_chord.return_value.apply_async.assert_called()
+        # Assert chord was called with expiration
+        mock_chord.return_value.apply_async.assert_called_once_with(args=(), expires=123)
 
     def test_harvesting_scheduler(self):
         mock_harvester = mock.MagicMock(spec=models.Harvester).return_value