From 246b9776ba223e7603e0ff8648495d75500df9e3 Mon Sep 17 00:00:00 2001
From: RobinTF <83676088+RobinTF@users.noreply.github.com>
Date: Fri, 22 Nov 2024 18:48:32 +0100
Subject: [PATCH] Check for cancellation after mutex

---
 src/index/CompressedRelation.cpp | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/src/index/CompressedRelation.cpp b/src/index/CompressedRelation.cpp
index 8e41809c5f..700c1862d9 100644
--- a/src/index/CompressedRelation.cpp
+++ b/src/index/CompressedRelation.cpp
@@ -77,8 +77,8 @@ CompressedRelationReader::asyncParallelBlockGenerator(
   auto readAndDecompressBlock = [&]()
       -> std::optional<
           std::pair<size_t, std::optional<DecompressedBlockAndMetadata>>> {
-    cancellationHandle->throwIfCancelled();
     std::unique_lock lock{blockIteratorMutex};
+    cancellationHandle->throwIfCancelled();
     if (blockMetadataIterator == endBlock) {
       return std::nullopt;
     }
@@ -102,8 +102,10 @@ CompressedRelationReader::asyncParallelBlockGenerator(
     CompressedBlock compressedBlock =
         readCompressedBlockFromFile(blockMetadata, columnIndices);
     lock.unlock();
+    cancellationHandle->throwIfCancelled();
     auto decompressedBlockAndMetadata = decompressAndPostprocessBlock(
         compressedBlock, blockMetadata.numRows_, scanConfig, blockMetadata);
+    cancellationHandle->throwIfCancelled();
     return std::pair{myIndex,
                      std::optional{std::move(decompressedBlockAndMetadata)}};
   };