Wire MessageBus reply traces through conditional Put pipeline

`CheckCondition::Outcome` now exposes the resulting trace (if any)
of the operations that were sent as part of the condition probe.
The outcome-accessor is changed to return a non-const reference to
allow for moving away the trace payload and into a higher-level
reply.

Turns out GetOperation did not aggregate reply traces as expected,
and PersistenceMessageTrackerImpl did not transfer trace state upon
failures (only success case). This makes this commit bigger than
initially expected, but trace coverage should now be improved on a
general basis, not just for conditional Put operations.

4 files changed, 104 insertions, 10 deletions

diff --git a/storage/src/tests/distributor/check_condition_test.cpp b/storage/src/tests/distributor/check_condition_test.cpp
index d4b7d7019d7..1b5cede8af6 100644
--- a/storage/src/tests/distributor/check_condition_test.cpp
+++ b/storage/src/tests/distributor/check_condition_test.cpp
@@ -9,7 +9,8 @@
 #include <vespa/storage/distributor/persistence_operation_metric_set.h>
 #include <vespa/storageapi/message/persistence.h>
 #include <tests/distributor/distributor_stripe_test_util.h>
-#include <vespa/vespalib/gtest/gtest.h>
+#include <gtest/gtest.h>
+#include <gmock/gmock.h>
 
 using namespace ::testing;
 
@@ -27,6 +28,7 @@ public:
     BucketId                      _bucket_id{16, 1234};
     TestAndSetCondition           _tas_cond{"foo or bar"};
     PersistenceOperationMetricSet _metrics{"dummy_metrics", nullptr};
+    uint32_t                      _trace_level{5};
 
     CheckConditionTest();
     ~CheckConditionTest() override;
@@ -52,7 +54,8 @@ public:
         auto bucket        = Bucket(FixedBucketSpaces::default_space(), _bucket_id);
         assert(_bucket_id.contains(doc_bucket));
         return CheckCondition::create_if_inconsistent_replicas(bucket, bucket_space, _doc_id, _tas_cond,
-                                                               node_context(), operation_context(), _metrics);
+                                                               node_context(), operation_context(), _metrics,
+                                                               _trace_level);
     }
 
     std::shared_ptr<api::GetCommand> sent_get_command(size_t idx) {
@@ -86,6 +89,12 @@ public:
         return make_reply(*sent_get_command(cmd_idx), ts, true, false);
     }
 
+    std::shared_ptr<api::GetReply> make_trace_reply(size_t cmd_idx, api::Timestamp ts, std::string trace_message) {
+        auto reply = make_reply(*sent_get_command(cmd_idx), ts, true, false);
+        MBUS_TRACE(reply->getTrace(), _trace_level, trace_message);
+        return reply;
+    }
+
     std::shared_ptr<api::GetReply> make_failed_reply(size_t cmd_idx) {
         auto reply = make_reply(*sent_get_command(cmd_idx), 0, false, false);
         reply->setResult(api::ReturnCode(api::ReturnCode::ABORTED, "did a bork"));
@@ -149,6 +158,7 @@ TEST_F(CheckConditionTest, starting_sends_condition_probe_gets) {
     EXPECT_EQ(cmd->condition(), _tas_cond);
     EXPECT_EQ(cmd->getFieldSet(), NoFields::NAME);
     EXPECT_EQ(cmd->internal_read_consistency(), api::InternalReadConsistency::Strong);
+    EXPECT_EQ(cmd->getTrace().getLevel(), _trace_level);
 }
 
 TEST_F(CheckConditionTest, condition_matching_completes_check_with_match_outcome) {
@@ -232,4 +242,15 @@ TEST_F(CheckConditionTest, check_fails_if_replica_set_changed_between_start_and_
     });
 }
 
+TEST_F(CheckConditionTest, nested_get_traces_are_propagated_to_outcome) {
+    test_cond_with_2_gets_sent([&](auto& cond) {
+        cond.handle_reply(_sender, make_trace_reply(0, 100, "hello"));
+        cond.handle_reply(_sender, make_trace_reply(1, 200, "world"));
+    }, [&](auto& outcome) {
+        auto trace_str = outcome.trace().toString();
+        EXPECT_THAT(trace_str, HasSubstr("hello"));
+        EXPECT_THAT(trace_str, HasSubstr("world"));
+    });
+}
+
 }
diff --git a/storage/src/tests/distributor/distributor_stripe_test_util.h b/storage/src/tests/distributor/distributor_stripe_test_util.h
index 707418512f9..9963b2c96b4 100644
--- a/storage/src/tests/distributor/distributor_stripe_test_util.h
+++ b/storage/src/tests/distributor/distributor_stripe_test_util.h
@@ -220,6 +220,15 @@ public:
         return cmd;
     }
 
+    template <typename ReplyType>
+    requires std::is_base_of_v<api::StorageReply, ReplyType>
+    [[nodiscard]] std::shared_ptr<ReplyType> sent_reply(size_t idx) {
+        assert(idx < _sender.replies().size());
+        auto reply = std::dynamic_pointer_cast<ReplyType>(_sender.reply(idx));
+        assert(reply != nullptr);
+        return reply;
+    }
+
     void config_enable_condition_probing(bool enable);
     void tag_content_node_supports_condition_probing(uint16_t index, bool supported);
 
diff --git a/storage/src/tests/distributor/getoperationtest.cpp b/storage/src/tests/distributor/getoperationtest.cpp
index 6601fcabbeb..600a4de462e 100644
--- a/storage/src/tests/distributor/getoperationtest.cpp
+++ b/storage/src/tests/distributor/getoperationtest.cpp
@@ -16,8 +16,9 @@
 #include <vespa/storage/distributor/externaloperationhandler.h>
 #include <vespa/storage/distributor/operations/external/getoperation.h>
 #include <vespa/storageapi/message/persistence.h>
-#include <vespa/vespalib/gtest/gtest.h>
 #include <iomanip>
+#include <gtest/gtest.h>
+#include <gmock/gmock.h>
 
 using config::ConfigGetter;
 using config::FileSpec;
@@ -75,7 +76,8 @@ struct GetOperationTest : Test, DistributorStripeTestUtil {
                std::string authorVal,
                uint32_t timestamp,
                bool is_tombstone = false,
-               bool condition_matched = false)
+               bool condition_matched = false,
+               std::string trace_msg = "")
     {
         if (idx == LastCommand) {
             idx = _sender.commands().size() - 1;
@@ -98,6 +100,9 @@ struct GetOperationTest : Test, DistributorStripeTestUtil {
 
         auto reply = std::make_shared<api::GetReply>(*tmp, doc, timestamp, false, is_tombstone, condition_matched);
         reply->setResult(result);
+        if (!trace_msg.empty()) {
+            MBUS_TRACE(reply->getTrace(), 1, trace_msg);
+        }
 
         op->receive(_sender, reply);
     }
@@ -110,6 +115,10 @@ struct GetOperationTest : Test, DistributorStripeTestUtil {
         sendReply(idx, api::ReturnCode::OK, "", timestamp, false, condition_match);
     }
 
+    void reply_with_trace(uint32_t idx, uint32_t timestamp, std::string trace_message) {
+        sendReply(idx, api::ReturnCode::OK, "", timestamp, false, true, std::move(trace_message));
+    }
+
     void replyWithFailure() {
         sendReply(LastCommand, api::ReturnCode::IO_FAILURE, "", 0);
     }
@@ -682,6 +691,7 @@ void GetOperationTest::set_up_condition_match_get_operation() {
     TestAndSetCondition my_cond("my_cool_condition");
     auto msg = std::make_shared<api::GetCommand>(makeDocumentBucket(BucketId(0)), docId, document::NoFields::NAME);
     msg->set_condition(my_cond);
+    msg->getTrace().setLevel(9); // FIXME a very tiny bit dirty to set this here >_>
     start_operation(std::move(msg), api::InternalReadConsistency::Strong);
 
     ASSERT_EQ("Get => 0,Get => 2,Get => 1", _sender.getCommands(true));
@@ -742,4 +752,21 @@ TEST_F(GetOperationTest, condition_match_result_is_aggregated_for_newest_replica
     EXPECT_EQ(replica_of(api::Timestamp(500), bucketId, 2, true, false), *op->newest_replica());
 }
 
+TEST_F(GetOperationTest, trace_is_aggregated_from_all_sub_replies_and_propagated_to_operation_reply) {
+    ASSERT_NO_FATAL_FAILURE(set_up_condition_match_get_operation());
+
+    ASSERT_NO_FATAL_FAILURE(reply_with_trace(0, 400, "foo"));
+    ASSERT_NO_FATAL_FAILURE(reply_with_trace(1, 500, "bar"));
+    ASSERT_NO_FATAL_FAILURE(reply_with_trace(2, 300, "baz"));
+
+    ASSERT_EQ(_sender.replies().size(), 1);
+    auto get_reply = sent_reply<api::GetReply>(0);
+
+    auto trace_str = get_reply->getTrace().toString();
+    EXPECT_THAT(trace_str, HasSubstr("foo"));
+    EXPECT_THAT(trace_str, HasSubstr("bar"));
+    EXPECT_THAT(trace_str, HasSubstr("baz"));
+}
+
+
 }
diff --git a/storage/src/tests/distributor/putoperationtest.cpp b/storage/src/tests/distributor/putoperationtest.cpp
index 405f25f9359..ff375e5b902 100644
--- a/storage/src/tests/distributor/putoperationtest.cpp
+++ b/storage/src/tests/distributor/putoperationtest.cpp
@@ -10,20 +10,18 @@
 #include <vespa/storageapi/message/bucket.h>
 #include <vespa/storageapi/message/persistence.h>
 #include <vespa/storageapi/message/state.h>
-#include <vespa/vespalib/gtest/gtest.h>
 #include <vespa/vespalib/text/stringtokenizer.h>
 #include <vespa/config/helper/configgetter.h>
+#include <gtest/gtest.h>
+#include <gmock/gmock.h>
 
-using std::shared_ptr;
 using config::ConfigGetter;
 using config::FileSpec;
 using vespalib::string;
 using namespace document;
-using namespace storage;
-using namespace storage::api;
-using namespace storage::lib;
 using namespace std::literals::string_literals;
 using document::test::makeDocumentBucket;
+using storage::api::TestAndSetCondition;
 
 using namespace ::testing;
 
@@ -493,7 +491,8 @@ TEST_F(PutOperationTest, update_correct_bucket_on_remapped_put) {
     {
         std::shared_ptr<api::StorageCommand> msg2  = _sender.command(0);
         std::shared_ptr<api::StorageReply> reply(msg2->makeReply().release());
-        PutReply* sreply = (PutReply*)reply.get();
+        auto* sreply = dynamic_cast<api::PutReply*>(reply.get());
+        ASSERT_TRUE(sreply);
         sreply->remapBucketId(document::BucketId(17, 13));
         sreply->setBucketInfo(api::BucketInfo(1,2,3,4,5));
         op->receive(_sender, reply);
@@ -690,6 +689,7 @@ void PutOperationTest::set_up_tas_put_with_2_inconsistent_replica_nodes(bool cre
     auto put = createPut(doc);
     put->setCondition(TestAndSetCondition("test.foo"));
     put->set_create_if_non_existent(create);
+    put->getTrace().setLevel(9);
     sendPut(std::move(put));
     ASSERT_EQ("Get => 1,Get => 0", _sender.getCommands(true));
 }
@@ -821,4 +821,41 @@ TEST_F(PutOperationTest, conditional_put_no_replicas_case_with_create_set_acts_a
     EXPECT_TRUE(put_n1->get_create_if_non_existent());
 }
 
+TEST_F(PutOperationTest, trace_is_propagated_from_condition_probe_gets_ok_probe_case) {
+    ASSERT_NO_FATAL_FAILURE(set_up_tas_put_with_2_inconsistent_replica_nodes());
+
+    ASSERT_EQ(sent_get_command(0)->getTrace().getLevel(), 9);
+    auto get_reply = make_get_reply(*sent_get_command(0), 50, false, true);
+    MBUS_TRACE(get_reply->getTrace(), 1, "a foo walks into a bar");
+
+    op->receive(_sender, get_reply);
+    op->receive(_sender, make_get_reply(*sent_get_command(1), 50, false, true));
+
+    ASSERT_EQ("Get => 1,Get => 0,Put => 1,Put => 0", _sender.getCommands(true));
+    sendReply(2);
+    sendReply(3);
+    ASSERT_EQ(_sender.replies().size(), 1);
+    auto put_reply = sent_reply<api::PutReply>(0);
+
+    auto trace_str = put_reply->getTrace().toString();
+    EXPECT_THAT(trace_str, HasSubstr("a foo walks into a bar"));
+}
+
+TEST_F(PutOperationTest, trace_is_propagated_from_condition_probe_gets_failed_probe_case) {
+    ASSERT_NO_FATAL_FAILURE(set_up_tas_put_with_2_inconsistent_replica_nodes());
+
+    auto get_reply = make_get_reply(*sent_get_command(0), 50, false, false);
+    MBUS_TRACE(get_reply->getTrace(), 1, "a foo walks into a zoo");
+
+    op->receive(_sender, get_reply);
+    op->receive(_sender, make_get_reply(*sent_get_command(1), 50, false, false));
+
+    ASSERT_EQ("Get => 1,Get => 0", _sender.getCommands(true));
+    ASSERT_EQ(_sender.replies().size(), 1);
+    auto put_reply = sent_reply<api::PutReply>(0);
+
+    auto trace_str = put_reply->getTrace().toString();
+    EXPECT_THAT(trace_str, HasSubstr("a foo walks into a zoo"));
+}
+
 }