Repository: cortexproject/cortex (Go)

Issue:
Compactor: Race condition between cleaner and compactor of partitioning compaction could cause data loss
**Describe the bug**
A race condition could happen when cleaner deleting previous version of partitioned group info file and its visit marers while another compactor picked up compaction job from previous version of partitioned group info file. This race condition would result in extra result block being compacted based on previous version of partitioned group info file. If there would be a new version of partitioned group info file being generated later on, that extra result block would be treated as result block from new version. Since the result block from previous version always contains few samples than new version, data loss would occur under such situation.

**To Reproduce**
It is hard to setup a controllable reproduce steps. Will provide the timeline the race condition happens:

1. compactor x: Generate partitioned group version 1 with total 8 partitions.
1. compactor n..m: Finished compacting all 8 partitions for partitioned group version 1.
1. compactor 1: Load partitioned group version 1 from storage along with all existing and new partitioned groups. And iterating through to pick a compaction job out of them.
1. cleaner: Found partitioned group version 1 got all partitions completed. Delete partitioned group version 1 from bucket store along with all visit markers associated with this partitioned group.
1. compactor 1: Picked partition 0 from partitioned group version 1 for compaction and checked if there was visit marker for this partition. It started compaction since the visit marker was deleted by cleaner.
1. compactor 2: Found there was new lower level block from the same time range got uploaded and generated partitioned group version 2 with total 4 partitions. The partitioned group would contain newly uploaded blocks and all 8 result blocks from partitioned group version 1.
1. compactor n..m: Started compacting partition 1, 2, 3 from partitioned group version 2 since partition 0 got visit marker that was held by compactor 1.
1. compactor 1: Finished compacting partition 0 from partitioned group version 1. Visit marker got updated to complete state for partition 0.
1. compactor n..m: Finished compacting partition 1, 2, 3 from partitioned group version 2. Visit markers got updated to complete state for partition 1, 2, 3.
1. cleaner: Found all 4 partitions were completed based on visit markers. It considered partitioned group version 2 was complete. It marked all parent blocks in partitioned group version 2 for deletion. Those blocks contains newly uploaded blocks and 8 result blocks from partitioned group version 1. 

At the end, blocks got deleted:
- block1: partition 0 from partitioned group version 1
- block2: partition 1 from partitioned group version 1
- block3: partition 2 from partitioned group version 1
- block4: partition 3 from partitioned group version 1
- block5: partition 4 from partitioned group version 1
- block6: partition 5 from partitioned group version 1
- block7: partition 6 from partitioned group version 1
- block8: partition 7 from partitioned group version 1

Blocks were active:
- block9:  partition 0 from partitioned group version 1
- block10: partition 1 from partitioned group version 2
- block11: partition 2 from partitioned group version 2
- block12: partition 3 from partitioned group version 2

There was no data loss in block10, block11, and block12.

There was data loss in block9 because:
- Partition 0 from partitioned group version 2 should contain:
  -  1/4 of the data from newly uploaded block
  - all data from block1 and block5
- block9 only contained all data from block1

**Expected behavior**
No data loss

**Environment:**
 - Infrastructure: Any distributed environment that has multiple compactors running at same time

**Possible Fix**
- Short term: 
  - Use `CreationTime` from `PartitionedGroupInfo` as a version tracker to make sure the picked compaction job matches the latest version of partitioned group in bucket store.
  - Planner could do such check. After planner passed check for visit marker, it could also make sure partitioned group info still exists in bucket store. If it exists, planner could check `CreationTime` matches.
- Long term:
  - Use leader/worker architecture for compactor.
  - Leader should deal with all the coordinations and cleanup to avoid race condition.
  - Worker should purely do actual compaction only.




Fix (patch):
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 8f609adeaca..8fc7022eea1 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -2,6 +2,9 @@
 
 ## master / unreleased
 
+* [BUGFIX] Compactor: Avoid race condition which allow a grouper to not compact all partitions. #7082
+
+
 ## 1.20.0 in progress
 
 * [CHANGE] StoreGateway/Alertmanager: Add default 5s connection timeout on client. #6603
diff --git a/pkg/compactor/partition_compaction_grouper.go b/pkg/compactor/partition_compaction_grouper.go
index 2d308fb6360..4f1d955bfc6 100644
--- a/pkg/compactor/partition_compaction_grouper.go
+++ b/pkg/compactor/partition_compaction_grouper.go
@@ -639,6 +639,19 @@ func (g *PartitionCompactionGrouper) pickPartitionCompactionJob(partitionCompact
 			level.Info(partitionedGroupLogger).Log("msg", "skipping group because partition is visited")
 			continue
 		}
+
+		// Validate that the partition group still exists before creating a visit marker
+		// This prevents the race condition where the cleaner deletes the partition group
+		// between the visit marker check and the visit marker creation
+		if _, err := ReadPartitionedGroupInfo(g.ctx, g.bkt, g.logger, partitionedGroupID); err != nil {
+			if errors.Is(err, ErrorPartitionedGroupInfoNotFound) {
+				level.Info(partitionedGroupLogger).Log("msg", "skipping group because partition group was deleted by cleaner", "partitioned_group_id", partitionedGroupID)
+			} else {
+				level.Warn(partitionedGroupLogger).Log("msg", "unable to read partition group info", "err", err, "partitioned_group_id", partitionedGroupID)
+			}
+			continue
+		}
+
 		partitionedGroupKey := createGroupKeyWithPartitionID(groupHash, partitionID, *partitionedGroup)
 
 		level.Info(partitionedGroupLogger).Log("msg", "found compactable group for user", "group", partitionedGroup.String())
diff --git a/pkg/compactor/partition_compaction_planner.go b/pkg/compactor/partition_compaction_planner.go
index d9750a86e40..566a62e87e3 100644
--- a/pkg/compactor/partition_compaction_planner.go
+++ b/pkg/compactor/partition_compaction_planner.go
@@ -109,6 +109,34 @@ func (p *PartitionCompactionPlanner) PlanWithPartition(_ context.Context, metasB
 		}
 	}
 
+	// Double-check that the partition group still exists and is the same one we started with
+	// to prevent race condition with cleaner. If the cleaner deleted the partition group
+	// after we created the visit marker in the grouper, we should abort the compaction
+	// to avoid orphaned visit markers.
+	currentPartitionedGroupInfo, err := ReadPartitionedGroupInfo(p.ctx, p.bkt, p.logger, partitionedGroupID)
+	if err != nil {
+		if errors.Is(err, ErrorPartitionedGroupInfoNotFound) {
+			level.Warn(p.logger).Log("msg", "partition group was deleted by cleaner, aborting compaction", "partitioned_group_id", partitionedGroupID, "partition_id", partitionID)
+			return nil, plannerCompletedPartitionError
+		} else {
+			level.Warn(p.logger).Log("msg", "unable to read partition group info during planning", "err", err, "partitioned_group_id", partitionedGroupID, "partition_id", partitionID)
+			return nil, fmt.Errorf("unable to read partition group info for partition ID %d, partitioned group ID %d: %s", partitionID, partitionedGroupID, err.Error())
+		}
+	}
+
+	// Verify that this is the same partition group that the grouper created the visit marker for
+	// by comparing creation times. If they don't match, it means the cleaner deleted the old
+	// partition group and a new one was created with the same ID.
+	expectedCreationTime := partitionInfo.PartitionedGroupCreationTime
+	if currentPartitionedGroupInfo.CreationTime != expectedCreationTime {
+		level.Warn(p.logger).Log("msg", "partition group creation time mismatch, cleaner deleted old group and new one was created, aborting compaction",
+			"partitioned_group_id", partitionedGroupID,
+			"partition_id", partitionID,
+			"expected_creation_time", expectedCreationTime,
+			"current_creation_time", currentPartitionedGroupInfo.CreationTime)
+		return nil, plannerCompletedPartitionError
+	}
+
 	// Ensure all blocks fits within the largest range. This is a double check
 	// to ensure there's no bug in the previous blocks grouping, given this Plan()
 	// is just a pass-through.