From 5386739354fa11df2733ff399654aba3cd6bc324 Mon Sep 17 00:00:00 2001
From: Piotr Sarna <sarna@scylladb.com>
Date: Wed, 17 Mar 2021 11:14:28 +0100
Subject: [PATCH] tracing: allow providing a custom session record param

The mechanism of session record params is currently only used
to store query strings and a couple more params like consistency level,
but since we now have more frontends than just CQL and Thrift,
it would be nice to also allow the users to put custom parameters in
there.
An immediate first user of this mechanism would be alternator,
which is going to put the operation type under the "alternator_op" key.
The operation type is not part of the query string due to how DynamoDB's
protocol works - the op type is stored separately in the HTTP header.
While it's possible to extract the operation type from the session_id,
it might not be the case once #2572 is implemented.
---
 tracing/trace_state.cc |  4 ++++
 tracing/trace_state.hh | 17 +++++++++++++++++
 2 files changed, 21 insertions(+)

diff --git a/tracing/trace_state.cc b/tracing/trace_state.cc
index 8d7894010e..09d8b3a092 100644
--- a/tracing/trace_state.cc
+++ b/tracing/trace_state.cc
@@ -135,6 +135,10 @@ void trace_state::add_query(sstring_view val) {
     _params_ptr->queries.emplace_back(std::move(val));
 }
 
+void trace_state::add_session_param(sstring_view key, sstring_view val) {
+    _records->session_rec.parameters.emplace(std::move(key), std::move(val));
+}
+
 void trace_state::set_user_timestamp(api::timestamp_type val) {
     _params_ptr->user_timestamp.emplace(val);
 }
diff --git a/tracing/trace_state.hh b/tracing/trace_state.hh
index c6af17414a..8ac9cf589f 100644
--- a/tracing/trace_state.hh
+++ b/tracing/trace_state.hh
@@ -357,6 +357,16 @@ private:
      */
     void add_query(sstring_view val);
 
+    /**
+     * Store a custom session parameter.
+     * 
+     * Thus value will be stored in the params<string, string> map of a tracing session
+     * 
+     * @param key the parameter key
+     * @param val the parameter value
+     */
+    void add_session_param(sstring_view key, sstring_view val);
+
     /**
      * Store a user provided timestamp.
      *
@@ -475,6 +485,7 @@ private:
     friend void set_consistency_level(const trace_state_ptr& p, db::consistency_level val);
     friend void set_optional_serial_consistency_level(const trace_state_ptr& p, const std::optional<db::consistency_level>&val);
     friend void add_query(const trace_state_ptr& p, sstring_view val);
+    friend void add_session_param(const trace_state_ptr& p, sstring_view key, sstring_view val);
     friend void set_user_timestamp(const trace_state_ptr& p, api::timestamp_type val);
     friend void add_prepared_statement(const trace_state_ptr& p, prepared_checked_weak_ptr& prepared);
     friend void set_username(const trace_state_ptr& p, const std::optional<auth::authenticated_user>& user);
@@ -616,6 +627,12 @@ inline void add_query(const trace_state_ptr& p, sstring_view val) {
     }
 }
 
+inline void add_session_param(const trace_state_ptr& p, sstring_view key, sstring_view val) {
+    if (p) {
+        p->add_session_param(std::move(key), std::move(val));
+    }
+}
+
 inline void set_user_timestamp(const trace_state_ptr& p, api::timestamp_type val) {
     if (p) {
         p->set_user_timestamp(val);