tracing: make opentelemetry state data sharded

Move gathered telemetry data to opentelemetry_state_data and make it sharded.

Author: margdoc

service/storage_proxy.cc

@@ -3559,8 +3559,8 @@ class abstract_read_executor : public enable_shared_from_this<abstract_read_exec
             return _proxy->query_result_local(_schema, _cmd, _partition_range, opts, _trace_state, timeout);
         } else {
             tracing::trace(_trace_state, "read_data: sending a message to /{}", ep);
-            return _proxy->_messaging.send_read_data(netw::messaging_service::msg_addr{ep, 0}, timeout, *_cmd, _partition_range, opts.digest_algo).then([this, ep](rpc::tuple<query::result, rpc::optional<cache_temperature>, tracing::trace_state_ptr::cache_counter_t> result_hit_rate) {
-                auto&& [result, hit_rate, cache_counter] = result_hit_rate;
+            return _proxy->_messaging.send_read_data(netw::messaging_service::msg_addr{ep, 0}, timeout, *_cmd, _partition_range, opts.digest_algo).then([this, ep](rpc::tuple<query::result, rpc::optional<cache_temperature>, tracing::trace_state_ptr::cache_counter_t> response) {
+                auto&& [result, hit_rate, cache_counter] = response;
                 tracing::trace(_trace_state, "read_data: got response from /{}", ep);
                 tracing::modify_cache_counter(_trace_state, cache_counter);
                 return make_ready_future<rpc::tuple<foreign_ptr<lw_shared_ptr<query::result>>, cache_temperature, tracing::trace_state_ptr::cache_counter_t>>(rpc::tuple(make_foreign(::make_lw_shared<query::result>(std::move(result))), hit_rate.value_or(cache_temperature::invalid()), cache_counter));
@@ -5109,22 +5109,25 @@ void storage_proxy::init_messaging_service(shared_ptr<migration_manager> mm) {
             auto& cfg = sp->local_db().get_config();
             cmd.max_result_size.emplace(cfg.max_memory_for_unlimited_query_soft_limit(), cfg.max_memory_for_unlimited_query_hard_limit());
         }
-        return do_with(std::move(pr), std::move(sp), std::move(trace_state_ptr), [&cinfo, cmd = make_lw_shared<query::read_command>(std::move(cmd)), src_addr = std::move(src_addr), da, t, mm] (::compat::wrapping_partition_range& pr, shared_ptr<storage_proxy>& p, tracing::trace_state_ptr& trace_state_ptr) mutable {
-            p->get_stats().replica_data_reads++;
-            auto src_ip = src_addr.addr;
-            return mm->get_schema_for_read(cmd->schema_version, std::move(src_addr), p->_messaging).then([cmd, da, &pr, &p, &trace_state_ptr, t] (schema_ptr s) {
-                auto pr2 = ::compat::unwrap(std::move(pr), *s);
-                if (pr2.second) {
-                    // this function assumes singular queries but doesn't validate
-                    throw std::runtime_error("READ_DATA called with wrapping range");
-                }
-                query::result_options opts;
-                opts.digest_algo = da;
-                opts.request = da == query::digest_algorithm::none ? query::result_request::only_result : query::result_request::result_and_digest;
-                auto timeout = t ? *t : db::no_timeout;
-                return p->query_result_local(std::move(s), cmd, std::move(pr2.first), opts, trace_state_ptr, timeout);
-            }).finally([&trace_state_ptr, src_ip] () mutable {
-                tracing::trace(trace_state_ptr, "read_data handling is done, sending a response to /{}", src_ip);
+        return tracing::start(trace_state_ptr).then([pr = std::move(pr), sp = std::move(sp), trace_state_ptr, &cinfo, cmd = make_lw_shared<query::read_command>(std::move(cmd)), src_addr = std::move(src_addr), da, t, mm] () {
+            return do_with(std::move(pr), std::move(sp), std::move(trace_state_ptr), [&cinfo, cmd = std::move(cmd), src_addr = std::move(src_addr), da, t, mm] (::compat::wrapping_partition_range& pr, shared_ptr<storage_proxy>& p, tracing::trace_state_ptr& trace_state_ptr) mutable {
+                p->get_stats().replica_data_reads++;
+                auto src_ip = src_addr.addr;
+                return mm->get_schema_for_read(cmd->schema_version, std::move(src_addr), p->_messaging).then([cmd, da, &pr, &p, &trace_state_ptr, t] (schema_ptr s) {
+                    auto pr2 = ::compat::unwrap(std::move(pr), *s);
+                    if (pr2.second) {
+                        // this function assumes singular queries but doesn't validate
+                        throw std::runtime_error("READ_DATA called with wrapping range");
+                    }
+                    query::result_options opts;
+                    opts.digest_algo = da;
+                    opts.request = da == query::digest_algorithm::none ? query::result_request::only_result : query::result_request::result_and_digest;
+                    auto timeout = t ? *t : db::no_timeout;
+                    return p->query_result_local(std::move(s), cmd, std::move(pr2.first), opts, trace_state_ptr, timeout);
+                }).finally([&trace_state_ptr, src_ip] () mutable {
+                    tracing::trace(trace_state_ptr, "read_data handling is done, sending a response to /{}", src_ip);
+                    return tracing::stop(trace_state_ptr);
+                });
             });
         });
     });

tracing/trace_state.cc

@@ -339,7 +339,7 @@ sstring trace_state::raw_value_to_sstring(const cql3::raw_value_view& v, const d
     }
 }
 
-void opentelemetry_state::serialize_replicas(bytes& serialized) const {
+void opentelemetry_state_data::serialize_replicas(bytes& serialized) const {
     const auto size = htonl(_replicas.size());
     const auto *size_ptr = reinterpret_cast<const int8_t*>(&size);
     serialized += bytes{size_ptr, sizeof(size)};
@@ -356,7 +356,7 @@ void opentelemetry_state::serialize_replicas(bytes& serialized) const {
     }
 }
 
-void opentelemetry_state::serialize_cache_counter(bytes& serialized) const {
+void opentelemetry_state_data::serialize_cache_counter(bytes& serialized) const {
     const auto counter = htonl(_cache_counter);
     const auto *counter_ptr = reinterpret_cast<const int8_t*>(&counter);
     serialized += bytes{counter_ptr, sizeof(counter)};

tracing/trace_state.hh

@@ -507,37 +507,90 @@ private:
 };
 
 
-class opentelemetry_state final {
+class opentelemetry_state_data final {
 public:
     using cache_counter_t = int32_t;
 
-private:
-    lw_shared_ptr<trace_state> _state_ptr;
-    bool const _opentelemetry_tracing{false};
     inet_address_vector_replica_set _replicas;
     // Number of read partitions that were found in cache.
     cache_counter_t _cache_counter{0};
 
     void serialize_replicas(bytes& serialized) const;
     void serialize_cache_counter(bytes& serialized) const;
+};
+
+
+class opentelemetry_state final {
+
+private:
+    lw_shared_ptr<trace_state> _state_ptr;
+    std::shared_ptr<sharded<opentelemetry_state_data>> _data;
+
+    class reducer {
+    private:
+        opentelemetry_state_data _data;
+
+    public:
+        void operator()(const opentelemetry_state_data& to_reduce) {
+            _data._replicas.insert(_data._replicas.end(),
+                                   to_reduce._replicas.begin(),
+                                   to_reduce._replicas.end());
+            _data._cache_counter += to_reduce._cache_counter;
+        }
+
+        opentelemetry_state_data get() const {
+            return _data;
+        }
+    };
 
 public:
     opentelemetry_state() = default;
-    opentelemetry_state(lw_shared_ptr<trace_state> state_ptr, bool opentelemetry_tracing = false)
-            : _state_ptr(std::move(state_ptr)), _opentelemetry_tracing(opentelemetry_tracing)
+    opentelemetry_state(lw_shared_ptr<trace_state> state_ptr, std::shared_ptr<sharded<opentelemetry_state_data>> data)
+            : _state_ptr(std::move(state_ptr)), _data(std::move(data))
+    {}
+    opentelemetry_state(std::nullptr_t , std::shared_ptr<sharded<opentelemetry_state_data>> data)
+            : _state_ptr(nullptr), _data(std::move(data))
     {}
-    opentelemetry_state(std::nullptr_t, bool opentelemetry_tracing = false)
-            : _state_ptr(nullptr), _opentelemetry_tracing(opentelemetry_tracing)
+    template<typename TraceState>
+    opentelemetry_state(TraceState state_ptr, bool opentelemetry_tracing = false)
+            : opentelemetry_state(std::move(state_ptr),
+                                  opentelemetry_tracing
+                                      ? std::make_shared<sharded<opentelemetry_state_data>>()
+                                      : nullptr)
     {}
 
+    /**
+     * Helper function to initialize sharded data in opentelemetry_state.
+     */
+    future<> start() {
+        return _data->start();
+    }
+
+    /**
+     * Helper function to properly stop sharded data in opentelemetry_state.
+     */
+    future<> stop() {
+        return _data->stop();
+    }
+
+    /**
+     * Helper function to collect opentelemetry data from all shards.
+     * It should be called before serializing opentelemetry_state.
+     */
+    future<> collect_data() {
+        return _data->map_reduce(reducer{}, std::identity{}).then([&local_data = _data->local()] (auto data) {
+            local_data = data;
+        });
+    }
+
     /**
      * @return serialized opentelemetry state.
      */
     bytes serialize() const noexcept {
         bytes serialized{};
 
-        serialize_replicas(serialized);
-        serialize_cache_counter(serialized);
+        _data->local().serialize_replicas(serialized);
+        _data->local().serialize_cache_counter(serialized);
 
         return serialized;
     }
@@ -548,30 +601,30 @@ public:
      * @param replicas list of contacted replicas
      */
     void set_replicas(const inet_address_vector_replica_set& replicas) {
-        _replicas = replicas;
+        _data->local()._replicas = replicas;
     }
 
     /**
      * Increment counter of partitions read from cache.
      *
      * @param count number of partitions
      */
-    void modify_cache_counter(cache_counter_t count) {
-        _cache_counter += count;
+    void modify_cache_counter(opentelemetry_state_data::cache_counter_t count) {
+        _data->local()._cache_counter += count;
     }
 
     /**
      * @return number of partitions that were found in cache
      */
-    cache_counter_t get_cache_counter() const {
-        return _cache_counter;
+    opentelemetry_state_data::cache_counter_t get_cache_counter() const {
+        return _data->local()._cache_counter;
     }
 
     /**
      * @return True if OpenTelemetry trace state is stored.
      */
     bool has_opentelemetry() const noexcept {
-        return _opentelemetry_tracing;
+        return __builtin_expect(bool(_data), false);
     };
 
     /**
@@ -594,6 +647,13 @@ public:
     trace_state& get_tracing() const noexcept {
         return *_state_ptr;
     }
+
+    /**
+     * @return A pointer to sharded data.
+     */
+    std::shared_ptr<sharded<opentelemetry_state_data>> get_data() const {
+        return _data;
+    }
 };
 
 
@@ -620,7 +680,7 @@ public:
         : _state_ptr(nullptr)
     {}
 
-    using cache_counter_t = opentelemetry_state::cache_counter_t;
+    using cache_counter_t = opentelemetry_state_data::cache_counter_t;
 
     /**
      * @return True if classic trace state is stored.
@@ -872,7 +932,7 @@ inline std::optional<trace_info> make_trace_info(const trace_state_ptr& state) {
     if (state.has_opentelemetry()) {
         trace_state_props_set props{};
         props.set(trace_state_props::opentelemetry);
-        return trace_info{props};
+        return trace_info{props, state.get_opentelemetry().get_data()};
     }
 
     return std::nullopt;
@@ -910,6 +970,30 @@ inline trace_state_ptr::cache_counter_t get_cache_counter(const trace_state_ptr&
     return 0;
 }
 
+inline future<> start(const trace_state_ptr& p) {
+    if (p.has_opentelemetry()) {
+        return p.get_opentelemetry_ptr()->start();
+    }
+
+    return make_ready_future<>();
+}
+
+inline future<> stop(const trace_state_ptr& p) {
+    if (p.has_opentelemetry()) {
+        return p.get_opentelemetry_ptr()->stop();
+    }
+
+    return make_ready_future<>();
+}
+
+inline future<> collect_data(const trace_state_ptr& p) {
+    if (p.has_opentelemetry()) {
+        return p.get_opentelemetry_ptr()->collect_data();
+    }
+
+    return make_ready_future<>();
+}
+
 // global_trace_state_ptr is a helper class that may be used for creating spans
 // of an existing tracing session on other shards. When a tracing span on a
 // different shard is needed global_trace_state_ptr would create a secondary

tracing/tracing.cc

@@ -141,19 +141,19 @@ trace_state_ptr tracing::create_session(const trace_info& secondary_session_info
     try {
         bool opentelemetry_tracing = secondary_session_info.state_props.contains(trace_state_props::opentelemetry);
         if (!started()) {
-            return opentelemetry_tracing ? make_lw_shared<opentelemetry_state>(nullptr, opentelemetry_tracing) : nullptr;
+            return opentelemetry_tracing ? make_lw_shared<opentelemetry_state>(nullptr, secondary_session_info.otel_data) : nullptr;
         }
 
         // Don't create a session if its records are likely to be dropped
         if (!may_create_new_session(secondary_session_info.session_id)) {
-            return opentelemetry_tracing ? make_lw_shared<opentelemetry_state>(nullptr, opentelemetry_tracing) : nullptr;
+            return opentelemetry_tracing ? make_lw_shared<opentelemetry_state>(nullptr, secondary_session_info.otel_data) : nullptr;
         }
 
         ++_active_sessions;
         if (secondary_session_info.state_props.contains(trace_state_props::classic)) {
-            return make_lw_shared<opentelemetry_state>(make_lw_shared<trace_state>(secondary_session_info), opentelemetry_tracing);
+            return make_lw_shared<opentelemetry_state>(make_lw_shared<trace_state>(secondary_session_info), secondary_session_info.otel_data);
         }
-        return make_lw_shared<opentelemetry_state>(nullptr, opentelemetry_tracing);
+        return make_lw_shared<opentelemetry_state>(nullptr, secondary_session_info.otel_data);
     } catch (...) {
         // return an uninitialized state in case of any error (OOM?)
         return nullptr;

tracing/tracing.hh

@@ -141,6 +141,8 @@ using trace_state_props_set = enum_set<super_enum<trace_state_props,
     trace_state_props::opentelemetry,
     trace_state_props::classic>>;
 
+class opentelemetry_state_data;
+
 class trace_info {
 public:
     utils::UUID session_id;
@@ -150,6 +152,7 @@ public:
     uint32_t slow_query_threshold_us; // in microseconds
     uint32_t slow_query_ttl_sec; // in seconds
     span_id parent_id;
+    std::shared_ptr<seastar::sharded<opentelemetry_state_data>> otel_data;
 
 public:
     trace_info(utils::UUID sid, trace_type t, bool w_o_c, trace_state_props_set s_p, uint32_t slow_query_threshold, uint32_t slow_query_ttl, span_id p_id)
@@ -164,7 +167,9 @@ public:
         state_props.set_if<trace_state_props::write_on_close>(write_on_close);
     }
 
-    trace_info(trace_state_props_set s_p) : state_props(s_p) {}
+    trace_info(trace_state_props_set s_p, std::shared_ptr<seastar::sharded<opentelemetry_state_data>> otel_data)
+        : state_props(s_p), otel_data(otel_data)
+    {}
 };
 
 struct one_session_records;