Fix NPE caused by txn commit race condition (#5940)

Fixed a race condition where if a transaction was just committed by another thread, we may return null values from TransactionService.get.

atlasdb-impl-shared/src/main/java/com/palantir/atlasdb/pue/ResilientCommitTimestampPutUnlessExistsTable.java

@@ -91,6 +91,7 @@ private Map<Long, Long> processReads(Map<Cell, byte[]> reads, Map<Long, Cell> st
             // if we reach here, actual is guaranteed to be a staging value
             try {
                 store.checkAndTouch(cell, actual);
+                checkAndTouch.put(cell, actual);
             } catch (CheckAndSetException e) {
                 PutUnlessExistsValue<Long> kvsValue = encodingStrategy.decodeValueAsCommitTimestamp(
                         startTs, Iterables.getOnlyElement(e.getActualValues()));
@@ -100,10 +101,11 @@ private Map<Long, Long> processReads(Map<Cell, byte[]> reads, Map<Long, Cell> st
                                 + "was found in the KVS",
                         SafeArg.of("kvsValue", kvsValue),
                         SafeArg.of("stagingValue", currentValue));
-                continue;
+            } finally {
+                // If we got here after catching CheckAndSetException, then some other thread committed this
+                // transaction, and we must therefore return the commit timestamp.
+                resultBuilder.put(startTs, commitTs);
             }
-            checkAndTouch.put(cell, actual);
-            resultBuilder.put(startTs, commitTs);
         }
         store.put(KeyedStream.stream(checkAndTouch)
                 .map(encodingStrategy::transformStagingToCommitted)

atlasdb-impl-shared/src/test/java/com/palantir/atlasdb/pue/ResilientCommitTimestampPutUnlessExistsTableTest.java

@@ -21,13 +21,15 @@
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.anyMap;
 import static org.mockito.Mockito.atLeastOnce;
+import static org.mockito.Mockito.doThrow;
 import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.spy;
 import static org.mockito.Mockito.verify;
 
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.palantir.atlasdb.keyvalue.api.Cell;
+import com.palantir.atlasdb.keyvalue.api.CheckAndSetException;
 import com.palantir.atlasdb.keyvalue.api.KeyAlreadyExistsException;
 import com.palantir.atlasdb.keyvalue.api.TableReference;
 import com.palantir.atlasdb.keyvalue.impl.InMemoryKeyValueService;
@@ -86,6 +88,29 @@ public void pueThatThrowsIsCorrectedOnGet() throws ExecutionException, Interrupt
         verify(spiedStore).put(anyMap());
     }
 
+    @Test
+    public void getReturnsStagingValuesThatWereCommittedBySomeoneElse()
+            throws ExecutionException, InterruptedException {
+        TwoPhaseEncodingStrategy strategy = TwoPhaseEncodingStrategy.INSTANCE;
+
+        long startTimestamp = 1L;
+        long commitTimestamp = 2L;
+        Cell timestampAsCell = strategy.encodeStartTimestampAsCell(startTimestamp);
+        byte[] stagingValue =
+                strategy.encodeCommitTimestampAsValue(startTimestamp, PutUnlessExistsValue.staging(commitTimestamp));
+        byte[] committedValue =
+                strategy.encodeCommitTimestampAsValue(startTimestamp, PutUnlessExistsValue.committed(commitTimestamp));
+        spiedStore.putUnlessExists(timestampAsCell, stagingValue);
+
+        List<byte[]> actualValues = ImmutableList.of(committedValue);
+
+        doThrow(new CheckAndSetException("done elsewhere", timestampAsCell, stagingValue, actualValues))
+                .when(spiedStore)
+                .checkAndTouch(timestampAsCell, stagingValue);
+
+        assertThat(pueTable.get(startTimestamp).get()).isEqualTo(commitTimestamp);
+    }
+
     @Test
     public void onceNonNullValueIsReturnedItIsAlwaysReturned() {
         PutUnlessExistsTable<Long, Long> putUnlessExistsTable = new ResilientCommitTimestampPutUnlessExistsTable(

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

@@ -0,0 +1,8 @@
+type: fix
+fix:
+  description: 'Fixed a race condition where if a transaction was just committed by
+    another thread, users would receive null values from TransactionService.get. This
+    caused users to observe that a completed transaction was still running, but did
+    not lead to any correctness issue. '
+  links:
+  - https://github.com/palantir/atlasdb/pull/5940