Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions src/postgres/src/backend/utils/misc/pg_yb_utils.c
Original file line number Diff line number Diff line change
Expand Up @@ -1206,7 +1206,7 @@ YBInitPostgresBackend(const char *program_name, const YbcPgInitPostgresInfo *ini
hex_encode((const char *) YbGetLocalTServerUuid(), UUID_LEN, hex_uuid);
hex_uuid[2 * UUID_LEN] = '\0';

YBCInitDistTrace(MyProcPid, hex_uuid);
YBCInitDistTrace(hex_uuid);
}
}
}
Expand All @@ -1217,7 +1217,7 @@ YBOnPostgresBackendShutdown()
YBCDestroyPgGate();

if (YBCIsDistTraceEnabled())
YBCCleanupDistTrace();
YBCShutdownDistTrace();
}

void
Expand Down
89 changes: 77 additions & 12 deletions src/yb/rpc/outbound_call.cc
Original file line number Diff line number Diff line change
Expand Up @@ -151,13 +151,11 @@ bool FinishedState(RpcCallState state) {
return false;
}

void SetSpanStatus(opentelemetry::trace::Span& span, RpcCallState state) {
void SetSpanStatus(opentelemetry::trace::Span& span, RpcCallState state, const Status& status) {
switch (state) {
case TIMED_OUT:
span.SetStatus(opentelemetry::trace::StatusCode::kError, "Call TimedOut");
return;
case FINISHED_ERROR:
span.SetStatus(opentelemetry::trace::StatusCode::kError, "Call ErroredOut");
span.SetStatus(opentelemetry::trace::StatusCode::kError, status.ToUserMessage());
return;
case FINISHED_SUCCESS:
span.SetStatus(opentelemetry::trace::StatusCode::kOk);
Expand Down Expand Up @@ -279,9 +277,16 @@ OutboundCall::OutboundCall(const RemoteMethod& remote_method,
IncrementCounter(rpc_metrics_->outbound_calls_created);
IncrementGauge(rpc_metrics_->outbound_calls_alive);

if (dist_trace::HasActiveContext()) {
otel_span_ = dist_trace::StartSpan(
Format("rpc $0", remote_method_.ToString()), dist_trace::GetPendingRpcAttrPairs());
// Capture this call's parent before StartClientSpanWithScope makes otel_span_ current;
// InvokeCallbackSync restores it around the callback so its follow-on work nests as a sibling.
trace_parent_ = dist_trace::GetActiveSpanContext();

otel_span_ = dist_trace::StartClientSpanWithScope(Format("rpc $0", remote_method_.ToString()));
if (otel_span_) {
otel_span_->SetAttribute("rpc.system", "outbound_rpc");
otel_span_->SetAttribute("rpc.service", remote_method_.service_name());
otel_span_->SetAttribute("rpc.method", remote_method_.method_name());
otel_span_->SetAttribute("rpc.call_id", call_id_);
}
}

Expand Down Expand Up @@ -372,6 +377,33 @@ Status OutboundCall::SetRequestParam(
metadata_size += 1; // add tag size of RequestHeader::kMetadataFieldNumber
}

// Distributed-trace context: extract the outbound span's SpanContext (if any) and pre-compute the
// serialized size of the TraceContextPB submessage so it can be folded into the header length.
size_t trace_context_size = 0;
size_t trace_context_message_size = 0;
uint32_t version_and_flags = 0;
opentelemetry::trace::TraceId trace_id;
opentelemetry::trace::SpanId span_id;
if (otel_span_) {
auto span_context = otel_span_->GetContext();
if (span_context.IsValid()) {
trace_id = span_context.trace_id();
span_id = span_context.span_id();
auto trace_flags = span_context.trace_flags();
// TraceContextPB submessage layout:
// - trace_id_hi: 1 byte tag + 8 bytes fixed64
// - trace_id_lo: 1 byte tag + 8 bytes fixed64
// - span_id: 1 byte tag + 8 bytes fixed64
// - version_and_flags: 1 byte tag + varint (size depends on the combined value)
constexpr uint32_t kVersion = 0;
version_and_flags = (kVersion << 8) | trace_flags.flags();
size_t version_and_flags_varint_size = Output::VarintSize32(version_and_flags);
trace_context_message_size = 3 * 9 + 1 + version_and_flags_varint_size;
trace_context_size =
1 + Output::VarintSize64(trace_context_message_size) + trace_context_message_size;
}
}

auto use_crc = FLAGS_rpc_enable_crc;
size_t header_pb_len = 1 + call_id_size + // int32 call_id = 1
serialized_remote_method.size() + // RemoteMethodPB remote_method = 2
Expand All @@ -380,6 +412,7 @@ Status OutboundCall::SetRequestParam(
if (pool_tag) {
header_pb_len += 1 + Output::VarintSize64(pool_tag); // uint64 pool_tag = 7
}
header_pb_len += trace_context_size; // TraceContext trace_context = 8
if (use_crc) {
header_pb_len += 1 + sizeof(uint32_t); // fixed32 crc = 15
}
Expand Down Expand Up @@ -431,6 +464,31 @@ Status OutboundCall::SetRequestParam(
dst = Output::WriteVarint64ToArray(pool_tag, dst);
}

if (trace_context_size > 0) {
// Write the TraceContextPB submessage. Field numbers match yb.TraceContextPB in common.proto.
// The 16-byte trace id splits into two big-endian 64-bit halves; the span id is one big-endian
// 64-bit.
dst = Output::WriteTagToArray(
(RequestHeader::kTraceContextFieldNumber << 3) | WireFormatLite::WIRETYPE_LENGTH_DELIMITED,
dst);
dst = Output::WriteVarint32ToArray(narrow_cast<uint32_t>(trace_context_message_size), dst);

dst = Output::WriteTagToArray(
(TraceContextPB::kTraceIdHiFieldNumber << 3) | WireFormatLite::WIRETYPE_FIXED64, dst);
dst = Output::WriteLittleEndian64ToArray(BigEndian::Load64(trace_id.Id().data()), dst);

dst = Output::WriteTagToArray(
(TraceContextPB::kTraceIdLoFieldNumber << 3) | WireFormatLite::WIRETYPE_FIXED64, dst);
dst = Output::WriteLittleEndian64ToArray(BigEndian::Load64(trace_id.Id().data() + 8), dst);

dst = Output::WriteTagToArray(
(TraceContextPB::kSpanIdFieldNumber << 3) | WireFormatLite::WIRETYPE_FIXED64, dst);
dst = Output::WriteLittleEndian64ToArray(BigEndian::Load64(span_id.Id().data()), dst);

dst = Output::WriteTagToArray(TraceContextPB::kVersionAndFlagsFieldNumber << 3, dst);
dst = Output::WriteVarint32ToArray(version_and_flags, dst);
}

// CRC should be at the end of header, otherwise adjust CRC filling logic below.
if (use_crc) {
dst = Output::WriteTagToArray(
Expand Down Expand Up @@ -472,7 +530,7 @@ OutboundCall::State OutboundCall::state() const {
return state_.load(std::memory_order_acquire);
}

bool OutboundCall::SetState(State new_state) {
bool OutboundCall::SetState(State new_state, const Status& status) {
auto old_state = state();
// Sanity check state transitions.
DVLOG(3) << "OutboundCall " << this << " (" << ToString() << ") switching from "
Expand All @@ -487,7 +545,8 @@ bool OutboundCall::SetState(State new_state) {
}
if (state_.compare_exchange_weak(old_state, new_state, std::memory_order_acq_rel)) {
if (otel_span_ && FinishedState(new_state)) {
SetSpanStatus(*otel_span_, new_state);
DCHECK(otel_span_->span);
SetSpanStatus(*otel_span_->span, new_state, status);
otel_span_->End();
}
return true;
Expand Down Expand Up @@ -557,7 +616,13 @@ void OutboundCall::InvokeCallbackSync(std::optional<CoarseTimePoint> now_optiona
// TODO: consider removing the cycle-based mechanism of reporting slow callbacks below.

int64_t start_cycles = CycleClock::Now();
callback_();
// Re-activate the call's parent context so RPCs the callback issues nest as siblings, not
// parentless roots. No-op when trace_parent_ is invalid; parent_scope drops at block end so it
// can't leak.
{
auto parent_scope = dist_trace::ActivateParentScope(trace_parent_);
callback_();
}
// Clear the callback, since it may be holding onto reference counts
// via bound parameters. We do this inside the timer because it's possible
// the user has naughty destructors that block, and we want to account for that
Expand Down Expand Up @@ -692,7 +757,7 @@ void OutboundCall::SetFailed(const Status &status, std::unique_ptr<ErrorStatusPB
bool invoke_callback;
{
std::lock_guard l(mtx_);
invoke_callback = SetState(RpcCallState::FINISHED_ERROR);
invoke_callback = SetState(RpcCallState::FINISHED_ERROR, status);
if (invoke_callback) {
status_ = status;
if (status_.IsRemoteError()) {
Expand Down Expand Up @@ -731,7 +796,7 @@ void OutboundCall::SetTimedOut() {
call_id_);
std::lock_guard l(mtx_);
status_ = std::move(status);
invoke_callback = SetState(RpcCallState::TIMED_OUT);
invoke_callback = SetState(RpcCallState::TIMED_OUT, status_);
}
if (invoke_callback) {
InvokeCallback();
Expand Down
29 changes: 23 additions & 6 deletions src/yb/rpc/outbound_call.h
Original file line number Diff line number Diff line change
Expand Up @@ -72,12 +72,11 @@
#include "yb/util/memory/memory_usage.h"
#include "yb/util/monotime.h"
#include "yb/util/net/sockaddr.h"
#include "yb/util/dist_trace_fwd.h"
#include "yb/util/dist_trace.h"
#include "yb/util/object_pool.h"
#include "yb/util/ref_cnt_buffer.h"
#include "yb/util/shared_lock.h"
#include "yb/util/slice.h"
#include "yb/util/status_fwd.h"
#include "yb/util/trace.h"

using namespace std::literals;
Expand Down Expand Up @@ -349,6 +348,16 @@ class OutboundCall : public RpcCall {
void SetConnectionId(const ConnectionId& value, const std::string* hostname) {
conn_id_ = value;
hostname_ = hostname;
if (otel_span_) {
if (hostname) {
otel_span_->SetAttribute("network.peer.name", *hostname);
}
otel_span_->SetAttribute("network.peer.address", yb::ToString(value.remote()));
// Drop the OTEL span's scope here, on the calling thread, before the call is queued to the
// reactor. The span ends later, but the scope must be released on the thread that installed
// it.
otel_span_->DropScope();
}
}

void SetThreadPoolFailure(const Status& status) EXCLUDES(mtx_) {
Expand Down Expand Up @@ -453,6 +462,10 @@ class OutboundCall : public RpcCall {
virtual size_t GetSidecarsCount() const;
virtual size_t TransferSidecars(Sidecars* dest);

// Distributed-trace span for this call; LocalOutboundCall reads its context to parent the local
// inbound span.
const dist_trace::SpanWithScopePtr& otel_span() const { return otel_span_; }

// ----------------------------------------------------------------------------------------------
// Protected fields set in constructor or during initialization
// ----------------------------------------------------------------------------------------------
Expand Down Expand Up @@ -481,7 +494,7 @@ class OutboundCall : public RpcCall {

void NotifyTransferred(const Status& status, const ConnectionPtr& conn) override;

MUST_USE_RESULT bool SetState(State new_state);
MUST_USE_RESULT bool SetState(State new_state, const Status& status = Status::OK());
State state() const;

// return current status
Expand Down Expand Up @@ -595,9 +608,13 @@ class OutboundCall : public RpcCall {

std::unique_ptr<MetadataSerializer> metadata_serializer_;

// OpenTelemetry span for distributed tracing. Created when the call starts if there is an
// active trace context, ended when the call completes (success, failure, or timeout).
opentelemetry::nostd::shared_ptr<opentelemetry::trace::Span> otel_span_;
// OpenTelemetry span for this call, created at start (if a trace context is active) and ended at
// completion.
dist_trace::SpanWithScopePtr otel_span_;

// The trace context active when this call was constructed -- its PARENT, re-activated around the
// completion callback so follow-on RPCs nest as SIBLINGS of this call.
dist_trace::trace::SpanContext trace_parent_ = dist_trace::trace::SpanContext::GetInvalid();

// InvokeCallbackTask should be able to call InvokeCallbackSync and we don't want other that
// method to be public.
Expand Down
Loading
Loading