[TS | LIP-164000] Reset Sweep Progress

atlasdb-impl-shared/src/main/java/com/palantir/atlasdb/sweep/queue/ShardProgress.java

@@ -135,6 +135,32 @@ private long increaseValueFromToAtLeast(ShardAndStrategy shardAndStrategy, long
         return currentValue;
     }
 
+    public void resetProgressForShard(ShardAndStrategy shardAndStrategy) {
+        // TODO (jkong): This is a bit crappy. Really we want this to be INITIAL_TIMESTAMP, but that doesn't
+        //  serialise because we require an unsigned integer.
+        byte[] colValZero = createColumnValue(SweepQueueUtils.RESET_TIMESTAMP);
+
+        long currentValue = getLastSweptTimestamp(shardAndStrategy);
+        while (currentValue > SweepQueueUtils.RESET_TIMESTAMP) {
+            CheckAndSetRequest casRequest = createRequest(shardAndStrategy, currentValue, colValZero);
+            try {
+                log.info(
+                        "Attempting to reset targeted sweep progress for a shard.",
+                        SafeArg.of("shardAndStrategy", shardAndStrategy));
+                kvs.checkAndSet(casRequest);
+                log.info(
+                        "Reset targeted sweep progress for a shard.", SafeArg.of("shardAndStrategy", shardAndStrategy));
+            } catch (CheckAndSetException e) {
+                log.info(
+                        "Failed to reset targeted sweep progress for a shard; trying again if someone changed it "
+                                + "(unless they also reset it).",
+                        SafeArg.of("shardAndStrategy", shardAndStrategy),
+                        e);
+                currentValue = rethrowIfUnchanged(shardAndStrategy, currentValue, e);
+            }
+        }
+    }
+
     static byte[] createColumnValue(long newVal) {
         return SweepShardProgressTable.Value.of(newVal).persistValue();
     }

atlasdb-impl-shared/src/main/java/com/palantir/atlasdb/sweep/queue/SweepQueue.java

@@ -167,6 +167,16 @@ public long sweepNextBatch(ShardAndStrategy shardStrategy, long sweepTs) {
         return sweepBatch.entriesRead();
     }
 
+    public void resetSweepProgress() {
+        int shards = getNumShards();
+        log.info("Now attempting to reset sweep progress for both strategies...", SafeArg.of("numShards", shards));
+        for (int shard = 0; shard < shards; shard++) {
+            progress.resetProgressForShard(ShardAndStrategy.conservative(shard));
+            progress.resetProgressForShard(ShardAndStrategy.thorough(shard));
+        }
+        log.info("Sweep progress was reset for shards for both strategies.", SafeArg.of("numShards", shards));
+    }
+
     /**
      * Returns the most recently known number of shards.
      */

atlasdb-impl-shared/src/main/java/com/palantir/atlasdb/sweep/queue/SweepQueueUtils.java

@@ -43,6 +43,7 @@ public final class SweepQueueUtils {
     public static final int BATCH_SIZE_KVS = 1000;
     public static final long READ_TS = Long.MAX_VALUE;
     public static final long INITIAL_TIMESTAMP = -1L;
+    public static final long RESET_TIMESTAMP = 0L;
     public static final ColumnRangeSelection ALL_COLUMNS = allPossibleColumns();
     public static final int MINIMUM_WRITE_INDEX = -TargetedSweepMetadata.MAX_DEDICATED_ROWS;
 

atlasdb-impl-shared/src/main/java/com/palantir/atlasdb/sweep/queue/TargetedSweeper.java

@@ -52,6 +52,7 @@
 public class TargetedSweeper implements MultiTableSweepQueueWriter, BackgroundSweeper {
     private static final Logger log = LoggerFactory.getLogger(TargetedSweeper.class);
 
+    private final boolean shouldResetAndStopSweep;
     private final Supplier<TargetedSweepRuntimeConfig> runtime;
     private final List<Follower> followers;
     private final MetricsManager metricsManager;
@@ -76,6 +77,7 @@ private TargetedSweeper(
         this.conservativeScheduler =
                 new BackgroundSweepScheduler(install.conservativeThreads(), SweeperStrategy.CONSERVATIVE);
         this.thoroughScheduler = new BackgroundSweepScheduler(install.thoroughThreads(), SweeperStrategy.THOROUGH);
+        this.shouldResetAndStopSweep = install.resetTargetedSweepQueueProgressAndStopSweep();
         this.followers = followers;
         this.metricsConfiguration = install.metricsConfiguration();
     }
@@ -175,8 +177,15 @@ public void initializeWithoutRunning(
     @Override
     public void runInBackground() {
         assertInitialized();
-        conservativeScheduler.scheduleBackgroundThreads();
-        thoroughScheduler.scheduleBackgroundThreads();
+        if (shouldResetAndStopSweep) {
+            log.warn("This AtlasDB node is operating in a mode where it is attempting to reset the progress of "
+                    + "targeted sweep. While in this mode, your data is not getting swept: please restart your node "
+                    + "once it is confirmed that sweep progress has been reset.");
+            queue.resetSweepProgress();
+        } else {
+            conservativeScheduler.scheduleBackgroundThreads();
+            thoroughScheduler.scheduleBackgroundThreads();
+        }
     }
 
     @Override

atlasdb-impl-shared/src/main/java/com/palantir/atlasdb/sweep/queue/config/TargetedSweepInstallConfig.java

@@ -83,6 +83,21 @@ public TargetedSweepMetrics.MetricsConfiguration metricsConfiguration() {
         return TargetedSweepMetricsConfigurations.DEFAULT;
     }
 
+    /**
+     * Specifies whether on startup we should reset progress in the targeted sweep queue. This may be useful to deal
+     * with circumstances where entries are written to the targeted sweep queue after the sweep timestamp has
+     * progressed past it - while the transaction in question will necessarily fail, there may still be cruft in the
+     * targeted sweep queue.
+     *
+     * If set to true, resets progress to zero for each shard and strategy on startup. This configuration can also only
+     * be safely used if nodes are not actively sweeping, and so if configured to be true will prevent targeted sweep
+     * from running.
+     */
+    @Value.Default
+    public boolean resetTargetedSweepQueueProgressAndStopSweep() {
+        return false;
+    }
+
     public static TargetedSweepInstallConfig defaultTargetedSweepConfig() {
         return ImmutableTargetedSweepInstallConfig.builder().build();
     }

atlasdb-impl-shared/src/test/java/com/palantir/atlasdb/sweep/queue/ShardProgressTest.java

@@ -16,6 +16,7 @@
 package com.palantir.atlasdb.sweep.queue;
 
 import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatCode;
 import static org.assertj.core.api.Assertions.assertThatThrownBy;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.anyMap;
@@ -37,6 +38,7 @@
 
 public class ShardProgressTest {
     private static final long INITIAL_TIMESTAMP = SweepQueueUtils.INITIAL_TIMESTAMP;
+    private static final long RESET_TIMESTAMP = SweepQueueUtils.RESET_TIMESTAMP;
 
     private ShardProgress progress;
     private KeyValueService kvs;
@@ -186,6 +188,46 @@ public void repeatedlyFailingCasThrows() {
                 .isInstanceOf(CheckAndSetException.class);
     }
 
+    @Test
+    public void canResetProgressForSpecificShards() {
+        progress.updateLastSweptTimestamp(CONSERVATIVE_TEN, 8888L);
+        progress.updateLastSweptTimestamp(CONSERVATIVE_TWENTY, 8888L);
+        assertThat(progress.getLastSweptTimestamp(CONSERVATIVE_TEN)).isEqualTo(8888L);
+        assertThat(progress.getLastSweptTimestamp(CONSERVATIVE_TWENTY)).isEqualTo(8888L);
+
+        progress.resetProgressForShard(CONSERVATIVE_TEN);
+        assertThat(progress.getLastSweptTimestamp(CONSERVATIVE_TEN)).isEqualTo(RESET_TIMESTAMP);
+        assertThat(progress.getLastSweptTimestamp(CONSERVATIVE_TWENTY)).isEqualTo(8888L);
+    }
+
+    @Test
+    public void stopsTryingToResetIfSomeoneElseDid() {
+        KeyValueService mockKvs = mock(KeyValueService.class);
+        when(mockKvs.get(any(), anyMap()))
+                .thenReturn(ImmutableMap.of(DUMMY, createValue(8L)))
+                .thenReturn(ImmutableMap.of(DUMMY, createValue(4L)))
+                .thenReturn(ImmutableMap.of(DUMMY, createValue(RESET_TIMESTAMP)));
+        doThrow(new CheckAndSetException("sadness")).when(mockKvs).checkAndSet(any());
+        ShardProgress instrumentedProgress = new ShardProgress(mockKvs);
+
+        assertThatCode(() -> instrumentedProgress.resetProgressForShard(CONSERVATIVE_TEN))
+                .doesNotThrowAnyException();
+    }
+
+    @Test
+    public void repeatedlyFailingCasThrowsForReset() {
+        KeyValueService mockKvs = mock(KeyValueService.class);
+        when(mockKvs.get(any(), anyMap()))
+                .thenReturn(ImmutableMap.of(DUMMY, createValue(8L)))
+                .thenReturn(ImmutableMap.of(DUMMY, createValue(9L)))
+                .thenReturn(ImmutableMap.of(DUMMY, createValue(10L)));
+        doThrow(new CheckAndSetException("sadness")).when(mockKvs).checkAndSet(any());
+        ShardProgress instrumentedProgress = new ShardProgress(mockKvs);
+
+        assertThatCode(() -> instrumentedProgress.resetProgressForShard(CONSERVATIVE_TEN))
+                .isInstanceOf(CheckAndSetException.class);
+    }
+
     private Value createValue(long num) {
         SweepShardProgressTable.Value value = SweepShardProgressTable.Value.of(num);
         return Value.create(value.persistValue(), 0L);

changelog/@unreleased/pr-5277.v2.yml

@@ -0,0 +1,8 @@
+type: improvement
+improvement:
+  description: |-
+    Targeted sweep progress may be reset with the `resetTargetedSweepQueueProgressAndStopSweep` flag in targeted sweep install configuration. This may be useful in cleaning up cruft in the targeted sweep queue that may have been written by failed transactions.
+
+    As the name suggests, this will prevent sweep from cleaning up old cells, so users should not run with this configuration in the steady state. If running your service in HA, once the last node rolls and reports that it has successfully reset the sweep progress table, we can be certain that progress has been reset to zero.
+  links:
+  - https://github.com/palantir/atlasdb/pull/5277