From 1efabca51591d8824e3c39824f4cbca49c44cc82 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Botond=20D=C3=A9nes?= <bdenes@scylladb.com>
Date: Fri, 29 Sep 2023 06:43:35 -0400
Subject: [PATCH] tools/scylla-nodetool: implement the traceprobability
 commands

gettraceprobability and settraceprobability
---
 test/nodetool/test_traceprobability.py | 48 ++++++++++++++++++++++++++
 tools/scylla-nodetool.cc               | 46 ++++++++++++++++++++++++
 2 files changed, 94 insertions(+)
 create mode 100644 test/nodetool/test_traceprobability.py

diff --git a/test/nodetool/test_traceprobability.py b/test/nodetool/test_traceprobability.py
new file mode 100644
index 0000000000..baddcfaa80
--- /dev/null
+++ b/test/nodetool/test_traceprobability.py
@@ -0,0 +1,48 @@
+#
+# Copyright 2023-present ScyllaDB
+#
+# SPDX-License-Identifier: AGPL-3.0-or-later
+#
+
+from rest_api_mock import expected_request
+import utils
+
+
+def test_gettraceprobability(nodetool):
+    out = nodetool("gettraceprobability", expected_requests=[
+        expected_request("GET", "/storage_service/trace_probability", response=0.2)])
+
+    assert out == "Current trace probability: 0.2\n"
+
+
+def test_settraceprobability(nodetool):
+    nodetool("settraceprobability", "0.2", expected_requests=[
+        expected_request("POST", "/storage_service/trace_probability", params={"probability": "0.2"})])
+
+
+def test_settraceprobability_missing_param(nodetool):
+    utils.check_nodetool_fails_with(
+            nodetool,
+            ("settraceprobability",),
+            {},
+            ["nodetool: Required parameters are missing: trace_probability",
+             "error processing arguments: required parameters are missing: trace_probability"])
+
+
+def test_settraceprobability_invalid_type(nodetool):
+    utils.check_nodetool_fails_with(
+            nodetool,
+            ("settraceprobability", "adadad"),
+            {},
+            ["nodetool: trace_probability: can not convert \"adadad\" to a Double",
+             "error: the argument ('adadad') for option '--trace_probability' is invalid"])
+
+
+def test_settraceprobability_out_of_bounds(nodetool):
+    for value in ("-0.1", "1.1", "9000"):
+        utils.check_nodetool_fails_with(
+                nodetool,
+                ("settraceprobability", "--", value),
+                {},
+                ["nodetool: Trace probability must be between 0 and 1",
+                 "error processing arguments: trace probability must be between 0 and 1"])
diff --git a/tools/scylla-nodetool.cc b/tools/scylla-nodetool.cc
index 39b2fb74b2..9c96a5c659 100644
--- a/tools/scylla-nodetool.cc
+++ b/tools/scylla-nodetool.cc
@@ -141,6 +141,22 @@ void enablegossip_operation(scylla_rest_client& client, const bpo::variables_map
     client.post("/storage_service/gossiping");
 }
 
+void gettraceprobability_operation(scylla_rest_client& client, const bpo::variables_map& vm) {
+    auto res = client.get("/storage_service/trace_probability");
+    fmt::print(std::cout, "Current trace probability: {}\n", res.GetDouble());
+}
+
+void settraceprobability_operation(scylla_rest_client& client, const bpo::variables_map& vm) {
+    if (!vm.count("trace_probability")) {
+        throw std::invalid_argument("required parameters are missing: trace_probability");
+    }
+    const auto value = vm["trace_probability"].as<double>();
+    if (value < 0.0 or value > 1.0) {
+        throw std::invalid_argument("trace probability must be between 0 and 1");
+    }
+    client.post("/storage_service/trace_probability", {{"probability", fmt::to_string(value)}});
+}
+
 void statusbackup_operation(scylla_rest_client& client, const bpo::variables_map& vm) {
     auto status = client.get("/storage_service/incremental_backups");
     fmt::print(std::cout, "{}\n", status.GetBool() ? "running" : "not running");
@@ -274,6 +290,36 @@ Fore more information, see: https://opensource.docs.scylladb.com/stable/operatin
             },
             enablegossip_operation
         },
+        {
+            {
+                "gettraceprobability",
+                "Displays the current trace probability value",
+R"(
+This value is the probability for tracing a request. To change this value see settraceprobability.
+
+Fore more information, see: https://opensource.docs.scylladb.com/stable/operating-scylla/nodetool-commands/gettraceprobability.html
+)",
+            },
+            gettraceprobability_operation
+        },
+        {
+            {
+                "settraceprobability",
+                "Sets the probability for tracing a request",
+R"(
+Value is trace probability between 0 and 1. 0 the trace will never happen and 1
+the trace will always happen. Anything in between is a percentage of the time,
+converted into a decimal. For example, 60% would be 0.6.
+
+Fore more information, see: https://opensource.docs.scylladb.com/stable/operating-scylla/nodetool-commands/settraceprobability.html
+)",
+                { },
+                {
+                    typed_option<double>("trace_probability", "trace probability value, must between 0 and 1, e.g. 0.2", 1),
+                },
+            },
+            settraceprobability_operation
+        },
         {
             {
                 "statusbackup",