Do not return TaS error code when doc is not found for conditional updatevekterli/change-test-and-set-update-not-found-semantics

Instead, return an OK response with `wasFound() == false`. This mirrors
the behavior of existing non-conditional updates.

9 files changed, 148 insertions, 79 deletions

diff --git a/storage/src/tests/distributor/twophaseupdateoperationtest.cpp b/storage/src/tests/distributor/twophaseupdateoperationtest.cpp
index da32225cde3..d4344ee0255 100644
--- a/storage/src/tests/distributor/twophaseupdateoperationtest.cpp
+++ b/storage/src/tests/distributor/twophaseupdateoperationtest.cpp
@@ -939,7 +939,7 @@ TEST_F(TwoPhaseUpdateOperationTest, safe_path_condition_unknown_doc_type_fails_w
               _sender.getLastReply(true));
 }
 
-TEST_F(TwoPhaseUpdateOperationTest, safe_path_condition_with_missing_doc_and_no_auto_create_fails_with_tas_error) {
+TEST_F(TwoPhaseUpdateOperationTest, safe_path_condition_with_missing_doc_and_no_auto_create_returns_ok_but_tagged_not_found) {
     setup_stripe(2, 2, "storage:2 distributor:1");
     auto cb = sendUpdate("0=1/2/3,1=2/3/4", UpdateOptions().condition("testdoctype1.headerval==120"));
 
@@ -947,15 +947,12 @@ TEST_F(TwoPhaseUpdateOperationTest, safe_path_condition_with_missing_doc_and_no_
     // Both Gets return nothing at all, nothing at all.
     replyToGet(*cb, _sender, 0, 100, false);
     replyToGet(*cb, _sender, 1, 110, false);
-    EXPECT_EQ("UpdateReply(id:ns:testdoctype1::1, "
-                          "BucketId(0x0000000000000000), "
-                          "timestamp 0, timestamp of updated doc: 0) "
-                          "ReturnCode(TEST_AND_SET_CONDITION_FAILED, "
-                                     "Document did not exist)",
+    EXPECT_EQ("UpdateReply(id:ns:testdoctype1::1, BucketId(0x0000000000000000), "
+              "timestamp 0, timestamp of updated doc: 0) ReturnCode(NONE)",
               _sender.getLastReply(true));
 
-    EXPECT_EQ(metrics().updates.failures.notfound.getValue(), 0); // Not counted as "not found" failure when TaS is present
-    EXPECT_EQ(metrics().updates.failures.test_and_set_failed.getValue(), 1);
+    EXPECT_EQ(metrics().updates.failures.notfound.getValue(), 1);
+    EXPECT_EQ(metrics().updates.failures.test_and_set_failed.getValue(), 0);
 }
 
 TEST_F(TwoPhaseUpdateOperationTest, safe_path_condition_with_missing_doc_and_auto_create_sends_puts) {
diff --git a/storage/src/tests/persistence/persistencetestutils.h b/storage/src/tests/persistence/persistencetestutils.h
index 94ae7b9fb53..e60260f3ee8 100644
--- a/storage/src/tests/persistence/persistencetestutils.h
+++ b/storage/src/tests/persistence/persistencetestutils.h
@@ -150,6 +150,18 @@ public:
                                                 _replySender, MockBucketLock::make(bucket, _mock_bucket_locks), std::move(cmd));
     }
 
+    template <typename T>
+    requires std::is_base_of_v<api::StorageReply, T>
+    [[nodiscard]] std::shared_ptr<T>
+    fetch_single_reply(MessageTracker::UP tracker) {
+        if (tracker && tracker->hasReply()) {
+            tracker->sendReply(); // Forward to queue so we can fetch it below
+        }
+        std::shared_ptr<api::StorageMessage> msg;
+        _replySender.queue.getNext(msg, 60s);
+        return std::dynamic_pointer_cast<T>(msg);
+    }
+
     api::ReturnCode
     fetchResult(const MessageTracker::UP & tracker) {
         if (tracker) {
diff --git a/storage/src/tests/persistence/testandsettest.cpp b/storage/src/tests/persistence/testandsettest.cpp
index 5be1c7cd92a..02be3e4d2c0 100644
--- a/storage/src/tests/persistence/testandsettest.cpp
+++ b/storage/src/tests/persistence/testandsettest.cpp
@@ -66,7 +66,7 @@ struct TestAndSetTest : PersistenceTestUtils {
 
     document::Document::SP createTestDocument();
     document::Document::SP retrieveTestDocument();
-    void setTestCondition(api::TestAndSetCommand & command);
+    static void setTestCondition(api::TestAndSetCommand& command);
     void putTestDocument(bool matchingHeader, api::Timestamp timestamp);
     void assertTestDocumentFoundAndMatchesContent(const document::FieldValue & value);
 
@@ -154,6 +154,16 @@ TEST_F(TestAndSetTest, conditional_remove_executed_on_condition_match) {
               dumpBucket(BUCKET_ID));
 }
 
+TEST_F(TestAndSetTest, conditional_remove_to_non_existing_document_fails_with_tas_error) {
+    api::Timestamp timestamp = 0;
+    auto remove = std::make_shared<api::RemoveCommand>(BUCKET, testDocId, timestamp);
+    setTestCondition(*remove);
+
+    ASSERT_EQ(fetchResult(asyncHandler->handleRemove(*remove, createTracker(remove, BUCKET))).getResult(),
+              api::ReturnCode::Result::TEST_AND_SET_CONDITION_FAILED);
+    EXPECT_EQ("", dumpBucket(BUCKET_ID));
+}
+
 std::shared_ptr<api::UpdateCommand>
 TestAndSetTest::conditional_update_test(bool createIfMissing, api::Timestamp updateTimestamp)
 {
@@ -197,8 +207,10 @@ TEST_F(TestAndSetTest, conditional_update_not_executed_when_no_document_and_no_a
     api::Timestamp updateTimestamp = 200;
     auto updateUp = conditional_update_test(false, updateTimestamp);
 
-    ASSERT_EQ(fetchResult(asyncHandler->handleUpdate(*updateUp, createTracker(updateUp, BUCKET))).getResult(),
-              api::ReturnCode::Result::TEST_AND_SET_CONDITION_FAILED);
+    auto reply = fetch_single_reply<api::UpdateReply>(asyncHandler->handleUpdate(*updateUp, createTracker(updateUp, BUCKET)));
+    ASSERT_TRUE(reply);
+    EXPECT_EQ(reply->getResult(), api::ReturnCode());
+    EXPECT_FALSE(reply->wasFound());
     EXPECT_EQ("", dumpBucket(BUCKET_ID));
 }
 
@@ -270,7 +282,7 @@ TestAndSetTest::retrieveTestDocument()
     auto tracker = _persistenceHandler->simpleMessageHandler().handleGet(*get, createTracker(get, BUCKET));
     assert(tracker->getResult() == api::ReturnCode::Result::OK);
 
-    auto & reply = static_cast<api::GetReply &>(tracker->getReply());
+    auto& reply = dynamic_cast<api::GetReply&>(tracker->getReply());
     assert(reply.wasFound());
 
     return reply.getDocument();
diff --git a/storage/src/vespa/storage/distributor/operations/external/twophaseupdateoperation.cpp b/storage/src/vespa/storage/distributor/operations/external/twophaseupdateoperation.cpp
index 515b72520ec..302845b2b8d 100644
--- a/storage/src/vespa/storage/distributor/operations/external/twophaseupdateoperation.cpp
+++ b/storage/src/vespa/storage/distributor/operations/external/twophaseupdateoperation.cpp
@@ -568,12 +568,8 @@ TwoPhaseUpdateOperation::handleSafePathReceivedGet(DistributorStripeMessageSende
         }
         docToUpdate = reply.getDocument();
         setUpdatedForTimestamp(receivedTimestamp);
-    } else if (hasTasCondition() && !shouldCreateIfNonExistent()) {
-        replyWithTasFailure(sender, "Document did not exist");
-        return;
     } else if (shouldCreateIfNonExistent()) {
-        LOG(debug, "No existing documents found for %s, creating blank document to update",
-            update_doc_id().c_str());
+        LOG(debug, "No existing documents found for %s, creating blank document to update", update_doc_id().c_str());
         docToUpdate = createBlankDocument();
         setUpdatedForTimestamp(putTimestamp);
     } else {
diff --git a/storage/src/vespa/storage/persistence/asynchandler.cpp b/storage/src/vespa/storage/persistence/asynchandler.cpp
index e20c0475556..85a13f5f0db 100644
--- a/storage/src/vespa/storage/persistence/asynchandler.cpp
+++ b/storage/src/vespa/storage/persistence/asynchandler.cpp
@@ -115,7 +115,7 @@ bucketStatesAreSemanticallyEqual(const api::BucketInfo& a, const api::BucketInfo
 class UnrevertableRemoveEntryProcessor : public BucketProcessor::EntryProcessor {
 public:
     using DocumentIdsAndTimeStamps = std::vector<spi::IdAndTimestamp>;
-    UnrevertableRemoveEntryProcessor(DocumentIdsAndTimeStamps & to_remove)
+    explicit UnrevertableRemoveEntryProcessor(DocumentIdsAndTimeStamps& to_remove) noexcept
         : _to_remove(to_remove)
     {}
 
@@ -159,11 +159,15 @@ AsyncHandler::handlePut(api::PutCommand& cmd, MessageTracker::UP trackerUP) cons
     tracker.setMetric(metrics);
     metrics.request_size.addValue(cmd.getApproxByteSize());
 
-    if (tasConditionExists(cmd) && !tasConditionMatches(cmd, tracker, tracker.context())) {
-        // Will also count condition parse failures etc as TaS failures, but
-        // those results _will_ increase the error metrics as well.
-        metrics.test_and_set_failed.inc();
-        return trackerUP;
+    if (tasConditionExists(cmd)) {
+        auto tas_res = tasConditionMatches(cmd, tracker, tracker.context(), DocNotFoundPolicy::ReturnTaSError);
+        if (tas_res == TasResult::FailedAndSetInTracker) {
+            // Will also count condition parse failures etc. as TaS failures, but
+            // those results _will_ increase the error metrics as well.
+            metrics.test_and_set_failed.inc();
+            return trackerUP;
+        }
+        assert(tas_res == TasResult::Matched); // NotFound treated as failure for Puts
     }
 
     spi::Bucket bucket = _env.getBucket(cmd.getDocumentId(), cmd.getBucket());
@@ -283,9 +287,19 @@ AsyncHandler::handleUpdate(api::UpdateCommand& cmd, MessageTracker::UP trackerUP
     tracker.setMetric(metrics);
     metrics.request_size.addValue(cmd.getApproxByteSize());
 
-    if (tasConditionExists(cmd) && !tasConditionMatches(cmd, tracker, tracker.context(), cmd.getUpdate()->getCreateIfNonExistent())) {
-        metrics.test_and_set_failed.inc();
-        return trackerUP;
+    if (tasConditionExists(cmd)) {
+        const bool create_missing = cmd.getUpdate()->getCreateIfNonExistent();
+        const auto not_found_policy = (create_missing ? DocNotFoundPolicy::TreatAsMatch : DocNotFoundPolicy::ReturnNotFound);
+        auto tas_res = tasConditionMatches(cmd, tracker, tracker.context(), not_found_policy);
+        if (tas_res == TasResult::FailedAndSetInTracker) {
+            metrics.test_and_set_failed.inc();
+            return trackerUP;
+        } else if (tas_res == TasResult::NotFound) {
+            metrics.notFound.inc();
+            tracker.setReply(std::make_shared<api::UpdateReply>(cmd, 0)); // Zero-timestamp == not found
+            return trackerUP;
+        }
+        assert(tas_res == TasResult::Matched);
     }
 
     spi::Bucket bucket = _env.getBucket(cmd.getDocumentId(), cmd.getBucket());
@@ -313,9 +327,13 @@ AsyncHandler::handleRemove(api::RemoveCommand& cmd, MessageTracker::UP trackerUP
     tracker.setMetric(metrics);
     metrics.request_size.addValue(cmd.getApproxByteSize());
 
-    if (tasConditionExists(cmd) && !tasConditionMatches(cmd, tracker, tracker.context())) {
-        metrics.test_and_set_failed.inc();
-        return trackerUP;
+    if (tasConditionExists(cmd)) {
+        auto tas_res = tasConditionMatches(cmd, tracker, tracker.context(), DocNotFoundPolicy::ReturnTaSError);
+        if (tas_res == TasResult::FailedAndSetInTracker) {
+            metrics.test_and_set_failed.inc();
+            return trackerUP;
+        }
+        assert(tas_res == TasResult::Matched); // NotFound treated as failure for Removes
     }
 
     spi::Bucket bucket = _env.getBucket(cmd.getDocumentId(), cmd.getBucket());
@@ -354,24 +372,26 @@ AsyncHandler::tasConditionExists(const api::TestAndSetCommand & cmd) {
     return cmd.getCondition().isPresent();
 }
 
-bool
-AsyncHandler::tasConditionMatches(const api::TestAndSetCommand & cmd, MessageTracker & tracker,
-                                  spi::Context & context, bool missingDocumentImpliesMatch) const {
+AsyncHandler::TasResult
+AsyncHandler::tasConditionMatches(const api::TestAndSetCommand& cmd, MessageTracker& tracker,
+                                  spi::Context& context, DocNotFoundPolicy doc_not_found_policy) const {
     try {
-        TestAndSetHelper helper(_env, _spi, _bucketIdFactory, cmd, missingDocumentImpliesMatch);
-
-        auto code = helper.retrieveAndMatch(context);
-        if (code.failed()) {
-            tracker.fail(code.getResult(), code.getMessage());
-            return false;
+        TestAndSetHelper helper(_env, _spi, _bucketIdFactory, cmd, doc_not_found_policy);
+
+        auto maybe_code = helper.retrieveAndMatch(context);
+        if (maybe_code && maybe_code->failed()) {
+            tracker.fail(maybe_code->getResult(), maybe_code->getMessage());
+            return TasResult::FailedAndSetInTracker;
+        } else if (!maybe_code) {
+            return TasResult::NotFound;
         }
     } catch (const TestAndSetException & e) {
         auto code = e.getCode();
         tracker.fail(code.getResult(), code.getMessage());
-        return false;
+        return TasResult::FailedAndSetInTracker;
     }
 
-    return true;
+    return TasResult::Matched;
 }
 
 bool
diff --git a/storage/src/vespa/storage/persistence/asynchandler.h b/storage/src/vespa/storage/persistence/asynchandler.h
index c5122647caa..c68fa962361 100644
--- a/storage/src/vespa/storage/persistence/asynchandler.h
+++ b/storage/src/vespa/storage/persistence/asynchandler.h
@@ -1,6 +1,7 @@
 // Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
 #pragma once
 
+#include "doc_not_found_policy.h"
 #include "messages.h"
 #include <vespa/storageapi/message/persistence.h>
 #include <vespa/storageapi/message/removelocation.h>
@@ -38,8 +39,15 @@ public:
 private:
     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;
+
+    enum class TasResult {
+        Matched,
+        FailedAndSetInTracker,
+        NotFound
+    };
+
+    TasResult tasConditionMatches(const api::TestAndSetCommand& cmd, MessageTracker& tracker,
+                                  spi::Context& context, DocNotFoundPolicy doc_not_found_policy) const;
     const PersistenceUtil            & _env;
     spi::PersistenceProvider         & _spi;
     BucketOwnershipNotifier          & _bucketOwnershipNotifier;
diff --git a/storage/src/vespa/storage/persistence/doc_not_found_policy.h b/storage/src/vespa/storage/persistence/doc_not_found_policy.h
new file mode 100644
index 00000000000..27a96ecc179
--- /dev/null
+++ b/storage/src/vespa/storage/persistence/doc_not_found_policy.h
@@ -0,0 +1,14 @@
+// Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+#pragma once
+
+namespace storage {
+
+// Specifies the semantics of a test-and-set comparison where the underlying
+// document does not exist in the backing store.
+enum class DocNotFoundPolicy {
+    ReturnTaSError, // Return an explicit test-and-set failure return code (_not_ a not-found error)
+    TreatAsMatch,   // Act as if the document _did_ exist and the condition matched it
+    ReturnNotFound  // Explicitly propagate up the fact that the document did not exist
+};
+
+}
diff --git a/storage/src/vespa/storage/persistence/testandsethelper.cpp b/storage/src/vespa/storage/persistence/testandsethelper.cpp
index 393dac09f72..1f5661a7813 100644
--- a/storage/src/vespa/storage/persistence/testandsethelper.cpp
+++ b/storage/src/vespa/storage/persistence/testandsethelper.cpp
@@ -15,13 +15,15 @@ using namespace std::string_literals;
 namespace storage {
 
 void TestAndSetHelper::resolveDocumentType(const document::DocumentTypeRepo & documentTypeRepo) {
-    if (_docTypePtr != nullptr) return;
-    if (!_docId.hasDocType()) {
+    if (_doc_type_ptr != nullptr) {
+        return;
+    }
+    if (!_doc_id.hasDocType()) {
         throw TestAndSetException(api::ReturnCode(api::ReturnCode::ILLEGAL_PARAMETERS, "Document id has no doctype"));
     }
 
-    _docTypePtr = documentTypeRepo.getDocumentType(_docId.getDocType());
-    if (_docTypePtr == nullptr) {
+    _doc_type_ptr = documentTypeRepo.getDocumentType(_doc_id.getDocType());
+    if (_doc_type_ptr == nullptr) {
         throw TestAndSetException(api::ReturnCode(api::ReturnCode::ILLEGAL_PARAMETERS, "Document type does not exist"));
     }
 }
@@ -31,39 +33,41 @@ void TestAndSetHelper::parseDocumentSelection(const document::DocumentTypeRepo &
     document::select::Parser parser(documentTypeRepo, bucketIdFactory);
 
     try {
-        _docSelectionUp = parser.parse(_cmd.getCondition().getSelection());
+        _doc_selection_up = parser.parse(_cmd.getCondition().getSelection());
     } catch (const document::select::ParsingFailedException & e) {
         throw TestAndSetException(api::ReturnCode(api::ReturnCode::ILLEGAL_PARAMETERS, "Failed to parse test and set condition: "s + e.getMessage()));
     }
 }
 
 spi::GetResult TestAndSetHelper::retrieveDocument(const document::FieldSet & fieldSet, spi::Context & context) {
-    return _spi.get(_env.getBucket(_docId, _cmd.getBucket()), fieldSet, _cmd.getDocumentId(), context);
+    return _spi.get(_env.getBucket(_doc_id, _cmd.getBucket()), fieldSet, _cmd.getDocumentId(), context);
 }
 
-TestAndSetHelper::TestAndSetHelper(const PersistenceUtil & env, const spi::PersistenceProvider  & spi,
-                                   const document::BucketIdFactory & bucketFactory,
-                                   const api::TestAndSetCommand & cmd, bool missingDocumentImpliesMatch)
+TestAndSetHelper::TestAndSetHelper(const PersistenceUtil& env,
+                                   const spi::PersistenceProvider& spi,
+                                   const document::BucketIdFactory& bucket_id_factory,
+                                   const api::TestAndSetCommand& cmd,
+                                   DocNotFoundPolicy doc_not_found_policy)
     : _env(env),
       _spi(spi),
       _cmd(cmd),
-      _docId(cmd.getDocumentId()),
-      _docTypePtr(_cmd.getDocumentType()),
-      _missingDocumentImpliesMatch(missingDocumentImpliesMatch)
+      _doc_id(cmd.getDocumentId()),
+      _doc_type_ptr(_cmd.getDocumentType()),
+      _doc_not_found_policy(doc_not_found_policy)
 {
-    const auto & repo = _env.getDocumentTypeRepo();
+    const auto& repo = _env.getDocumentTypeRepo();
     resolveDocumentType(repo);
-    parseDocumentSelection(repo, bucketFactory);
+    parseDocumentSelection(repo, bucket_id_factory);
 }
 
 TestAndSetHelper::~TestAndSetHelper() = default;
 
-api::ReturnCode
-TestAndSetHelper::retrieveAndMatch(spi::Context & context) {
+std::optional<api::ReturnCode>
+TestAndSetHelper::retrieveAndMatch(spi::Context& context) {
     // Walk document selection tree to build a minimal field set 
-    FieldVisitor fieldVisitor(*_docTypePtr);
+    FieldVisitor fieldVisitor(*_doc_type_ptr);
     try {
-        _docSelectionUp->visit(fieldVisitor);
+        _doc_selection_up->visit(fieldVisitor);
     } catch (const document::FieldNotFoundException& e) {
         return api::ReturnCode(api::ReturnCode::ILLEGAL_PARAMETERS,
                                vespalib::make_string("Condition field '%s' could not be found, or is an imported field. "
@@ -71,13 +75,12 @@ TestAndSetHelper::retrieveAndMatch(spi::Context & context) {
                                                      e.getFieldName().c_str()));
     }
 
-    // Retrieve document
     auto result = retrieveDocument(fieldVisitor.getFieldSet(), context);
 
     // If document exists, match it with selection
     if (result.hasDocument()) {
         auto docPtr = result.getDocumentPtr();
-        if (_docSelectionUp->contains(*docPtr) != document::select::Result::True) {
+        if (_doc_selection_up->contains(*docPtr) != document::select::Result::True) {
             return api::ReturnCode(api::ReturnCode::TEST_AND_SET_CONDITION_FAILED,
                                    vespalib::make_string("Condition did not match document nodeIndex=%d bucket=%" PRIx64 " %s",
                                                          _env._nodeIndex, _cmd.getBucketId().getRawId(),
@@ -86,14 +89,15 @@ TestAndSetHelper::retrieveAndMatch(spi::Context & context) {
 
         // Document matches
         return api::ReturnCode();
-    } else if (_missingDocumentImpliesMatch) {
+    } else if (_doc_not_found_policy == DocNotFoundPolicy::TreatAsMatch) {
         return api::ReturnCode();
+    } else if (_doc_not_found_policy == DocNotFoundPolicy::ReturnTaSError) {
+        return api::ReturnCode(api::ReturnCode::TEST_AND_SET_CONDITION_FAILED,
+                               vespalib::make_string("Document does not exist nodeIndex=%d bucket=%" PRIx64 " %s",
+                                                     _env._nodeIndex, _cmd.getBucketId().getRawId(),
+                                                     _cmd.hasBeenRemapped() ? "remapped" : ""));
     }
-
-    return api::ReturnCode(api::ReturnCode::TEST_AND_SET_CONDITION_FAILED,
-                           vespalib::make_string("Document does not exist nodeIndex=%d bucket=%" PRIx64 " %s",
-                                                 _env._nodeIndex, _cmd.getBucketId().getRawId(),
-                                                 _cmd.hasBeenRemapped() ? "remapped" : ""));
+    return std::nullopt; // Not found
 }
 
 } // storage
diff --git a/storage/src/vespa/storage/persistence/testandsethelper.h b/storage/src/vespa/storage/persistence/testandsethelper.h
index 82710e523c4..a4fcf9f9c5f 100644
--- a/storage/src/vespa/storage/persistence/testandsethelper.h
+++ b/storage/src/vespa/storage/persistence/testandsethelper.h
@@ -3,9 +3,11 @@
 
 #pragma once
 
+#include "doc_not_found_policy.h"
 #include <vespa/storageapi/message/persistence.h>
 #include <vespa/persistence/spi/result.h>
 #include <stdexcept>
+#include <optional>
 
 namespace document::select { class Node; }
 namespace document {
@@ -36,13 +38,13 @@ public:
 };
 
 class TestAndSetHelper {
-    const PersistenceUtil                  &_env;
-    const spi::PersistenceProvider         &_spi;
-    const api::TestAndSetCommand           &_cmd;
-    const document::DocumentId              _docId;
-    const document::DocumentType *          _docTypePtr;
-    std::unique_ptr<document::select::Node> _docSelectionUp;
-    bool                                    _missingDocumentImpliesMatch;
+    const PersistenceUtil&                  _env;
+    const spi::PersistenceProvider&         _spi;
+    const api::TestAndSetCommand&           _cmd;
+    const document::DocumentId              _doc_id;
+    const document::DocumentType*           _doc_type_ptr;
+    std::unique_ptr<document::select::Node> _doc_selection_up;
+    DocNotFoundPolicy                       _doc_not_found_policy;
 
     void resolveDocumentType(const document::DocumentTypeRepo & documentTypeRepo);
     void parseDocumentSelection(const document::DocumentTypeRepo & documentTypeRepo,
@@ -50,11 +52,15 @@ class TestAndSetHelper {
     spi::GetResult retrieveDocument(const document::FieldSet & fieldSet, spi::Context & context);
 
 public:
-    TestAndSetHelper(const PersistenceUtil & env, const spi::PersistenceProvider & _spi,
-                     const document::BucketIdFactory & bucketIdFactory,
-                     const api::TestAndSetCommand & cmd, bool missingDocumentImpliesMatch = false);
+    TestAndSetHelper(const PersistenceUtil& env, const spi::PersistenceProvider& _spi,
+                     const document::BucketIdFactory& bucket_id_factory,
+                     const api::TestAndSetCommand& cmd,
+                     DocNotFoundPolicy doc_not_found_policy = DocNotFoundPolicy::ReturnTaSError);
     ~TestAndSetHelper();
-    api::ReturnCode retrieveAndMatch(spi::Context & context);
+    // If document was not found:
+    //   iff doc_not_found_returns_tas_error == true, returns nullopt
+    //   otherwise, returns a TaS error return code
+    std::optional<api::ReturnCode> retrieveAndMatch(spi::Context& context);
 };
 
 } // storage