Handle MessageBus trace propagation for conditional removes

Pretty much a mirror image of the logic in `PutOperation`.

4 files changed, 49 insertions, 7 deletions

diff --git a/storage/src/tests/distributor/removeoperationtest.cpp b/storage/src/tests/distributor/removeoperationtest.cpp
index d88a0a574cc..d352d23bb8c 100644
--- a/storage/src/tests/distributor/removeoperationtest.cpp
+++ b/storage/src/tests/distributor/removeoperationtest.cpp
@@ -6,7 +6,8 @@
 #include <vespa/storage/distributor/distributor_stripe.h>
 #include <vespa/storage/distributor/operations/external/removeoperation.h>
 #include <vespa/storageapi/message/persistence.h>
-#include <vespa/vespalib/gtest/gtest.h>
+#include <gtest/gtest.h>
+#include <gmock/gmock.h>
 
 using documentapi::TestAndSetCondition;
 using document::test::makeDocumentBucket;
@@ -119,6 +120,7 @@ void ExtRemoveOperationTest::set_up_tas_remove_with_2_nodes(ReplicaState replica
 
     auto remove = createRemove(docId);
     remove->setCondition(TestAndSetCondition("test.foo"));
+    remove->getTrace().setLevel(9);
     sendRemove(std::move(remove));
     if (replica_state == ReplicaState::INCONSISTENT) {
         ASSERT_EQ("Get => 1,Get => 0", _sender.getCommands(true));
@@ -304,4 +306,41 @@ TEST_F(ExtRemoveOperationTest, failed_condition_probe_fails_op_with_returned_err
               _sender.getLastReply());
 }
 
+TEST_F(ExtRemoveOperationTest, trace_is_propagated_from_condition_probe_gets_ok_probe_case) {
+    ASSERT_NO_FATAL_FAILURE(set_up_tas_remove_with_2_nodes(ReplicaState::INCONSISTENT));
+
+    ASSERT_EQ(sent_get_command(0)->getTrace().getLevel(), 9);
+    auto get_reply = make_get_reply(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(1, 50, false, true));
+
+    ASSERT_EQ("Get => 1,Get => 0,Remove => 1,Remove => 0", _sender.getCommands(true));
+    reply_with(make_remove_reply(2, 50)); // remove from node 1
+    reply_with(make_remove_reply(3, 50)); // remove from node 0
+    ASSERT_EQ(_sender.replies().size(), 1);
+    auto remove_reply = sent_reply<api::RemoveReply>(0);
+
+    auto trace_str = remove_reply->getTrace().toString();
+    EXPECT_THAT(trace_str, HasSubstr("a foo walks into a bar"));
+}
+
+TEST_F(ExtRemoveOperationTest, trace_is_propagated_from_condition_probe_gets_failed_probe_case) {
+    ASSERT_NO_FATAL_FAILURE(set_up_tas_remove_with_2_nodes(ReplicaState::INCONSISTENT));
+
+    auto get_reply = make_get_reply(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(1, 50, false, false));
+
+    ASSERT_EQ("Get => 1,Get => 0", _sender.getCommands(true));
+    ASSERT_EQ(_sender.replies().size(), 1);
+    auto remove_reply = sent_reply<api::RemoveReply>(0);
+
+    auto trace_str = remove_reply->getTrace().toString();
+    EXPECT_THAT(trace_str, HasSubstr("a foo walks into a zoo"));
+}
+
 } // storage::distributor
diff --git a/storage/src/vespa/storage/distributor/operations/external/check_condition.h b/storage/src/vespa/storage/distributor/operations/external/check_condition.h
index 2a659c55081..062c9bb831d 100644
--- a/storage/src/vespa/storage/distributor/operations/external/check_condition.h
+++ b/storage/src/vespa/storage/distributor/operations/external/check_condition.h
@@ -136,7 +136,7 @@ public:
             const DistributorNodeContext& node_ctx,
             const DistributorStripeOperationContext& op_ctx,
             PersistenceOperationMetricSet& metric,
-            uint32_t trace_level = 0); // TODO remove default value
+            uint32_t trace_level);
 private:
     [[nodiscard]] bool replica_set_changed_after_get_operation() const;
 
diff --git a/storage/src/vespa/storage/distributor/operations/external/removeoperation.cpp b/storage/src/vespa/storage/distributor/operations/external/removeoperation.cpp
index d3001c37f7c..59ae4120fd6 100644
--- a/storage/src/vespa/storage/distributor/operations/external/removeoperation.cpp
+++ b/storage/src/vespa/storage/distributor/operations/external/removeoperation.cpp
@@ -48,13 +48,13 @@ void RemoveOperation::start_conditional_remove(DistributorStripeMessageSender& s
     document::Bucket bucket(_msg->getBucket().getBucketSpace(), _doc_id_bucket_id);
     _check_condition = CheckCondition::create_if_inconsistent_replicas(bucket, _bucket_space, _msg->getDocumentId(),
                                                                        _msg->getCondition(), _node_ctx, _op_ctx,
-                                                                       _temp_metric);
+                                                                       _temp_metric, _msg->getTrace().getLevel());
     if (!_check_condition) {
         start_direct_remove_dispatch(sender);
     } else {
         // Inconsistent replicas; need write repair
         _check_condition->start_and_send(sender);
-        const auto& outcome = _check_condition->maybe_outcome(); // Might be done immediately
+        auto& outcome = _check_condition->maybe_outcome(); // Might be done immediately
         if (outcome) {
             on_completed_check_condition(*outcome, sender);
         }
@@ -110,7 +110,7 @@ RemoveOperation::onReceive(DistributorStripeMessageSender& sender, const std::sh
 {
     if (_check_condition) {
         _check_condition->handle_reply(sender, msg);
-        const auto& outcome = _check_condition->maybe_outcome();
+        auto& outcome = _check_condition->maybe_outcome();
         if (!outcome) {
             return; // Condition check not done yet
         }
@@ -131,9 +131,12 @@ RemoveOperation::onReceive(DistributorStripeMessageSender& sender, const std::sh
     _tracker.receiveReply(sender, reply);
 }
 
-void RemoveOperation::on_completed_check_condition(const CheckCondition::Outcome& outcome,
+void RemoveOperation::on_completed_check_condition(CheckCondition::Outcome& outcome,
                                                    DistributorStripeMessageSender& sender)
 {
+    if (!outcome.trace().isEmpty()) {
+        _tracker.add_trace_tree_to_reply(outcome.steal_trace());
+    }
     if (outcome.matched_condition()) {
         _msg->clear_condition(); // Transform to unconditional Remove
         start_direct_remove_dispatch(sender);
diff --git a/storage/src/vespa/storage/distributor/operations/external/removeoperation.h b/storage/src/vespa/storage/distributor/operations/external/removeoperation.h
index ba6d42c5108..349a6182937 100644
--- a/storage/src/vespa/storage/distributor/operations/external/removeoperation.h
+++ b/storage/src/vespa/storage/distributor/operations/external/removeoperation.h
@@ -42,7 +42,7 @@ private:
 
     void start_direct_remove_dispatch(DistributorStripeMessageSender& sender);
     void start_conditional_remove(DistributorStripeMessageSender& sender);
-    void on_completed_check_condition(const CheckCondition::Outcome& outcome,
+    void on_completed_check_condition(CheckCondition::Outcome& outcome,
                                       DistributorStripeMessageSender& sender);
     [[nodiscard]] bool has_condition() const noexcept;
 };