fix(redis): Prevent counter corruption from concurrent mark handled in Redis RQ (#1878)

### Description

- Fixes counter corruption on concurrent `mark_request_as_handled` calls
for the same request. When two coroutines concurrently called
`mark_request_as_handled` for the same request, both could pass the
`hexists` check before either executed the pipeline, causing counters to
be updated multiple times for the same request, breaking the queue
operation.

### Issues

- Closes: #1873

### Testing

- Added a test to verify that counters are updated correctly during
concurrent execution.
- Added a test to verify that the request is correctly restored to
`in_progress` after failure.

Author: Mantisus

src/crawlee/storage_clients/_redis/_request_queue_client.py

@@ -396,41 +396,56 @@ async def get_request(self, unique_key: str) -> Request | None:
     @retry_on_error(RedisError)
     @override
     async def mark_request_as_handled(self, request: Request) -> ProcessedRequest | None:
-        # Check if the request is in progress.
-        check_in_progress = await await_redis_response(self._redis.hexists(self._in_progress_key, request.unique_key))
-        if not check_in_progress:
+        # `hdel` is an atomic operation, so we can be sure that if it returns 1, the request was in progress and is now
+        # removed from `in_progress`.
+        deleted = await await_redis_response(self._redis.hdel(self._in_progress_key, request.unique_key))
+
+        # If not deleted, the request was not in progress.
+        if not deleted:
             logger.warning(f'Marking request {request.unique_key} as handled that is not in progress.')
             return None
 
         # Update the request's handled_at timestamp.
         if request.handled_at is None:
             request.handled_at = datetime.now(timezone.utc)
 
-        async with self._get_pipeline() as pipe:
-            if self._dedup_strategy == 'default':
-                await await_redis_response(pipe.sadd(self._handled_set_key, request.unique_key))
-                await await_redis_response(pipe.srem(self._pending_set_key, request.unique_key))
-            elif self._dedup_strategy == 'bloom':
-                await await_redis_response(pipe.bf().add(self._handled_filter_key, request.unique_key))
+        try:
+            async with self._get_pipeline() as pipe:
+                if self._dedup_strategy == 'default':
+                    await await_redis_response(pipe.sadd(self._handled_set_key, request.unique_key))
+                    await await_redis_response(pipe.srem(self._pending_set_key, request.unique_key))
+                elif self._dedup_strategy == 'bloom':
+                    await await_redis_response(pipe.bf().add(self._handled_filter_key, request.unique_key))
 
-            await await_redis_response(pipe.hdel(self._in_progress_key, request.unique_key))
-            await await_redis_response(pipe.hset(self._data_key, request.unique_key, request.model_dump_json()))
+                await await_redis_response(pipe.hset(self._data_key, request.unique_key, request.model_dump_json()))
 
-            await self._update_metadata(
-                pipe,
-                **_QueueMetadataUpdateParams(
-                    update_accessed_at=True,
-                    update_modified_at=True,
-                    delta_handled_request_count=1,
-                    delta_pending_request_count=-1,
-                ),
-            )
+                await self._update_metadata(
+                    pipe,
+                    **_QueueMetadataUpdateParams(
+                        update_accessed_at=True,
+                        update_modified_at=True,
+                        delta_handled_request_count=1,
+                        delta_pending_request_count=-1,
+                    ),
+                )
 
-        return ProcessedRequest(
-            unique_key=request.unique_key,
-            was_already_present=True,
-            was_already_handled=True,
-        )
+            return ProcessedRequest(
+                unique_key=request.unique_key,
+                was_already_present=True,
+                was_already_handled=True,
+            )
+        except Exception:
+            blocked_until = int(datetime.now(tz=timezone.utc).timestamp() * 1000) + self._BLOCK_REQUEST_TIME
+            # If we fail to mark the request as handled after removing it from in_progress, we restore request in
+            # `in_progress` hash.
+            await await_redis_response(
+                self._redis.hset(
+                    self._in_progress_key,
+                    request.unique_key,
+                    json.dumps({'client_id': self.client_key, 'blocked_until_timestamp': blocked_until}),
+                )
+            )
+            raise
 
     @retry_on_error(RedisError)
     @override

tests/unit/storage_clients/_redis/test_redis_rq_client.py

@@ -292,3 +292,61 @@ async def test_get_metadata_does_not_retry_on_unexpected_exception(rq_client: Re
 
     # Verify that retry logic was not attempted
     assert mock_sleep.call_count == 0
+
+
+async def test_mark_request_as_handled_concurrent_no_double_decrement(
+    rq_client: RedisRequestQueueClient,
+) -> None:
+    """Test that concurrent calls to mark_request_as_handled decrement pending_request_count exactly once."""
+    request = Request.from_url('https://example.com/concurrent')
+    await rq_client.add_batch_of_requests([request])
+
+    fetched = await rq_client.fetch_next_request()
+    assert fetched is not None
+
+    results = await asyncio.gather(
+        rq_client.mark_request_as_handled(fetched),
+        rq_client.mark_request_as_handled(fetched),
+        rq_client.mark_request_as_handled(fetched),
+        rq_client.mark_request_as_handled(fetched),
+        rq_client.mark_request_as_handled(fetched),
+    )
+
+    successful = [result for result in results if result is not None]
+    assert len(successful) == 1
+
+    metadata = await rq_client.get_metadata()
+    assert metadata.pending_request_count == 0
+    assert metadata.handled_request_count == 1
+
+
+async def test_mark_request_as_handled_restores_in_progress_on_pipeline_failure(
+    rq_client: RedisRequestQueueClient,
+) -> None:
+    """Test that 'request' is restored to 'in_progress' when the pipeline fails after 'hdel'."""
+    request = Request.from_url('https://example.com/restore')
+    await rq_client.add_batch_of_requests([request])
+
+    fetched = await rq_client.fetch_next_request()
+    assert fetched is not None
+
+    mock_pipe = MagicMock()
+    mock_pipe.execute = AsyncMock(side_effect=RedisError('connection lost'))
+
+    mock_pipeline_ctx = MagicMock()
+    mock_pipeline_ctx.__aenter__ = AsyncMock(return_value=mock_pipe)
+    mock_pipeline_ctx.__aexit__ = AsyncMock(return_value=None)
+
+    with (
+        patch('crawlee._utils.retry._retry_sleep', new_callable=AsyncMock),
+        patch.object(rq_client.redis, 'pipeline', return_value=mock_pipeline_ctx),
+        pytest.raises(RedisError),
+    ):
+        await rq_client.mark_request_as_handled(fetched)
+
+    in_progress = await await_redis_response(rq_client.redis.hexists(rq_client._in_progress_key, fetched.unique_key))
+    assert in_progress
+
+    metadata = await rq_client.get_metadata()
+    assert metadata.pending_request_count == 1
+    assert metadata.handled_request_count == 0