Enforce document timestamp requirements for updates in backend

The document API has long since had a special field for update operations
where an optional expected _existing_ backend timestamp can be specified,
and where the update should only go through iff there is a timestamp
match.

This has been supported on the distributor all along, but only when
write-repair is taking place (i.e. rarely), but the actual backend
support has been lacking. No one has complained yet since this is
very much not an advertised feature, but if we want to e.g. use this
feature for improvements to batch updates we should ensure that it
works as expected.

With this commit, a non-zero "old timestamp" field is cross-checked
against the existing document, and the update is only applied if the
actual and expected timestamps match.

3 files changed, 60 insertions, 10 deletions

diff --git a/storage/src/tests/persistence/testandsettest.cpp b/storage/src/tests/persistence/testandsettest.cpp
index 3bb4b0e6345..b8012b9152b 100644
--- a/storage/src/tests/persistence/testandsettest.cpp
+++ b/storage/src/tests/persistence/testandsettest.cpp
@@ -65,7 +65,7 @@ struct TestAndSetTest : PersistenceTestUtils {
         PersistenceTestUtils::TearDown();
     }
 
-    std::shared_ptr<api::UpdateCommand> conditional_update_test(
+    std::shared_ptr<api::UpdateCommand> make_conditional_update(
         bool createIfMissing,
         api::Timestamp updateTimestamp);
 
@@ -184,7 +184,7 @@ TEST_F(TestAndSetTest, conditional_remove_executed_on_condition_match) {
 }
 
 std::shared_ptr<api::UpdateCommand>
-TestAndSetTest::conditional_update_test(bool createIfMissing, api::Timestamp updateTimestamp)
+TestAndSetTest::make_conditional_update(bool createIfMissing, api::Timestamp updateTimestamp)
 {
     auto docUpdate = std::make_shared<document::DocumentUpdate>(_env->_testDocMan.getTypeRepo(), testDoc->getType(), testDocId);
     docUpdate->addUpdate(document::FieldUpdate(testDoc->getField("content")).addUpdate(std::make_unique<document::AssignValueUpdate>(std::make_unique<StringFieldValue>(NEW_CONTENT))));
@@ -199,7 +199,7 @@ TEST_F(TestAndSetTest, conditional_update_not_executed_on_condition_mismatch) {
     api::Timestamp timestampOne = 0;
     api::Timestamp timestampTwo = 1;
     putTestDocument(false, timestampOne);
-    auto updateUp = conditional_update_test(false, timestampTwo);
+    auto updateUp = make_conditional_update(false, timestampTwo);
 
     ASSERT_EQ(fetchResult(asyncHandler->handleUpdate(*updateUp, createTracker(updateUp, BUCKET))).getResult(),
               api::ReturnCode::Result::TEST_AND_SET_CONDITION_FAILED);
@@ -212,7 +212,7 @@ TEST_F(TestAndSetTest, conditional_update_executed_on_condition_match) {
     api::Timestamp timestampOne = 0;
     api::Timestamp timestampTwo = 1;
     putTestDocument(true, timestampOne);
-    auto updateUp = conditional_update_test(false, timestampTwo);
+    auto updateUp = make_conditional_update(false, timestampTwo);
 
     ASSERT_EQ(fetchResult(asyncHandler->handleUpdate(*updateUp, createTracker(updateUp, BUCKET))).getResult(), api::ReturnCode::Result::OK);
     EXPECT_EQ(expectedDocEntryString(timestampOne, testDocId) +
@@ -224,7 +224,7 @@ TEST_F(TestAndSetTest, conditional_update_executed_on_condition_match) {
 
 TEST_F(TestAndSetTest, conditional_update_not_executed_when_no_document_and_no_auto_create) {
     api::Timestamp updateTimestamp = 200;
-    auto updateUp = conditional_update_test(false, updateTimestamp);
+    auto updateUp = make_conditional_update(false, updateTimestamp);
 
     ASSERT_EQ(fetchResult(asyncHandler->handleUpdate(*updateUp, createTracker(updateUp, BUCKET))).getResult(),
               api::ReturnCode::Result::TEST_AND_SET_CONDITION_FAILED);
@@ -233,13 +233,46 @@ TEST_F(TestAndSetTest, conditional_update_not_executed_when_no_document_and_no_a
 
 TEST_F(TestAndSetTest, conditional_update_executed_when_no_document_but_auto_create_is_enabled) {
     api::Timestamp updateTimestamp = 200;
-    auto updateUp = conditional_update_test(true, updateTimestamp);
+    auto updateUp = make_conditional_update(true, updateTimestamp);
 
     ASSERT_EQ(fetchResult(asyncHandler->handleUpdate(*updateUp, createTracker(updateUp, BUCKET))).getResult(), api::ReturnCode::Result::OK);
     EXPECT_EQ(expectedDocEntryString(updateTimestamp, testDocId), dumpBucket(BUCKET_ID));
     assertTestDocumentFoundAndMatchesContent(NEW_CONTENT);
 }
 
+// Although it's not a TaS _selection_ condition, we consider an update with a timestamp predicate
+// to be a "kind of" test-and-set operation, thus we test it here.
+TEST_F(TestAndSetTest, timestamp_predicated_update_should_not_apply_if_no_existing_document) {
+    auto update_cmd = make_conditional_update(true, api::Timestamp(200));
+    update_cmd->setCondition(documentapi::TestAndSetCondition()); // No condition; it has precedence
+    update_cmd->setOldTimestamp(api::Timestamp(150));
+    EXPECT_EQ(fetchResult(asyncHandler->handleUpdate(*update_cmd, createTracker(update_cmd, BUCKET))),
+              api::ReturnCode(api::ReturnCode::OK, "No document with requested timestamp found"));
+    EXPECT_EQ("", dumpBucket(BUCKET_ID));
+}
+
+TEST_F(TestAndSetTest, timestamp_predicated_update_should_not_apply_if_existing_document_has_unexpected_timestamp) {
+    putTestDocument(true, api::Timestamp(180));
+    auto update_cmd = make_conditional_update(false, api::Timestamp(200));
+    update_cmd->setCondition(documentapi::TestAndSetCondition());
+    update_cmd->setOldTimestamp(api::Timestamp(150)); // != 180, should fail
+    EXPECT_EQ(fetchResult(asyncHandler->handleUpdate(*update_cmd, createTracker(update_cmd, BUCKET))),
+              api::ReturnCode(api::ReturnCode::OK, "No document with requested timestamp found"));
+    EXPECT_EQ(expectedDocEntryString(api::Timestamp(180), testDocId), dumpBucket(BUCKET_ID));
+}
+
+TEST_F(TestAndSetTest, timestamp_predicated_update_is_applied_if_existing_timestamp_matches) {
+    putTestDocument(true, api::Timestamp(180));
+    auto update_cmd = make_conditional_update(false, api::Timestamp(200));
+    update_cmd->setCondition(documentapi::TestAndSetCondition());
+    update_cmd->setOldTimestamp(api::Timestamp(180));
+    EXPECT_EQ(fetchResult(asyncHandler->handleUpdate(*update_cmd, createTracker(update_cmd, BUCKET))),
+              api::ReturnCode(api::ReturnCode::OK, ""));
+    EXPECT_EQ(expectedDocEntryString(api::Timestamp(180), testDocId) +
+              expectedDocEntryString(api::Timestamp(200), testDocId),
+              dumpBucket(BUCKET_ID));
+}
+
 TEST_F(TestAndSetTest, invalid_document_selection_should_fail) {
     // Conditionally replace nonexisting document
     // Fail early since document selection is invalid 
diff --git a/storage/src/vespa/storage/persistence/asynchandler.cpp b/storage/src/vespa/storage/persistence/asynchandler.cpp
index 725cf2c7511..59e0853cc21 100644
--- a/storage/src/vespa/storage/persistence/asynchandler.cpp
+++ b/storage/src/vespa/storage/persistence/asynchandler.cpp
@@ -349,9 +349,17 @@ AsyncHandler::handleUpdate(api::UpdateCommand& cmd, MessageTracker::UP trackerUP
         metrics.test_and_set_failed.inc();
         return trackerUP;
     }
-
     spi::Bucket bucket = _env.getBucket(cmd.getDocumentId(), cmd.getBucket());
 
+    if ((cmd.getOldTimestamp() != 0) &&
+        (fetch_existing_document_timestamp(cmd.getDocumentId(), bucket, tracker.context()) != cmd.getOldTimestamp()))
+    {
+        metrics.notFound.inc();
+        // It's debatable if this should be OK or a TaS failure, but OK is the legacy behavior, i.e. what
+        // the distributor already does as part of a multiphase update where the expected timestamp is set.
+        tracker.fail(api::ReturnCode::OK, "No document with requested timestamp found");
+        return trackerUP;
+    }
     // Note that the &cmd capture is OK since its lifetime is guaranteed by the tracker
     auto task = makeResultTask([&cmd, tracker = std::move(trackerUP)](spi::Result::UP responseUP) {
         auto & response = dynamic_cast<const spi::UpdateResult &>(*responseUP);
@@ -398,6 +406,12 @@ AsyncHandler::handleRemove(api::RemoveCommand& cmd, MessageTracker::UP trackerUP
     return trackerUP;
 }
 
+api::Timestamp
+AsyncHandler::fetch_existing_document_timestamp(const document::DocumentId& id, const spi::Bucket& bucket, spi::Context& ctx) const
+{
+    return _spi.get(bucket, document::NoFields(), id, ctx).getTimestamp();
+}
+
 bool
 AsyncHandler::is_async_unconditional_message(const api::StorageMessage & cmd) noexcept
 {
diff --git a/storage/src/vespa/storage/persistence/asynchandler.h b/storage/src/vespa/storage/persistence/asynchandler.h
index c78dfe6282d..df47f4af537 100644
--- a/storage/src/vespa/storage/persistence/asynchandler.h
+++ b/storage/src/vespa/storage/persistence/asynchandler.h
@@ -37,9 +37,12 @@ public:
     static bool is_async_unconditional_message(const api::StorageMessage& cmd) noexcept;
 private:
     [[nodiscard]] bool checkProviderBucketInfoMatches(const spi::Bucket&, const api::BucketInfo&) const;
-    static bool tasConditionExists(const api::TestAndSetCommand& cmd);
-    bool tasConditionMatches(const api::TestAndSetCommand& cmd, MessageTracker& tracker,
-                             spi::Context& context, bool missingDocumentImpliesMatch = false) const;
+    [[nodiscard]] static bool tasConditionExists(const api::TestAndSetCommand& cmd);
+    [[nodiscard]] bool tasConditionMatches(const api::TestAndSetCommand& cmd, MessageTracker& tracker,
+                                           spi::Context& context, bool missingDocumentImpliesMatch = false) const;
+    [[nodiscard]] api::Timestamp fetch_existing_document_timestamp(const document::DocumentId& id,
+                                                                   const spi::Bucket& bucket,
+                                                                   spi::Context& ctx) const;
     void on_delete_bucket_complete(const document::Bucket& bucket) const;
     const PersistenceUtil            & _env;
     spi::PersistenceProvider         & _spi;