1 // Copyright 2023 gRPC authors.
2 //
3 // Licensed under the Apache License, Version 2.0 (the "License");
4 // you may not use this file except in compliance with the License.
5 // You may obtain a copy of the License at
6 //
7 //     http://www.apache.org/licenses/LICENSE-2.0
8 //
9 // Unless required by applicable law or agreed to in writing, software
10 // distributed under the License is distributed on an "AS IS" BASIS,
11 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 // See the License for the specific language governing permissions and
13 // limitations under the License.
14 
15 #include "python_census_context.h"
16 
17 #include <string.h>
18 
19 #include <iomanip>
20 #include <iostream>
21 #include <new>
22 
23 #include "absl/numeric/int128.h"
24 #include "absl/random/random.h"
25 #include "absl/strings/escaping.h"
26 #include "rpc_encoding.h"
27 
28 #include "src/core/lib/transport/transport.h"
29 
30 namespace grpc_observability {
31 
EnablePythonCensusStats(bool enable)32 void EnablePythonCensusStats(bool enable) {
33   g_python_census_stats_enabled = enable;
34 }
35 
EnablePythonCensusTracing(bool enable)36 void EnablePythonCensusTracing(bool enable) {
37   g_python_census_tracing_enabled = enable;
38 }
39 
PythonCensusStatsEnabled()40 bool PythonCensusStatsEnabled() {
41   return g_python_census_stats_enabled.load(std::memory_order_relaxed);
42 }
43 
PythonCensusTracingEnabled()44 bool PythonCensusTracingEnabled() {
45   return g_python_census_tracing_enabled.load(std::memory_order_relaxed);
46 }
47 
GenerateClientContext(absl::string_view method,absl::string_view trace_id,absl::string_view parent_span_id,PythonCensusContext * context)48 void GenerateClientContext(absl::string_view method, absl::string_view trace_id,
49                            absl::string_view parent_span_id,
50                            PythonCensusContext* context) {
51   // Destruct the current CensusContext to free the Span memory before
52   // overwriting it below.
53   context->~PythonCensusContext();
54   if (method.empty()) {
55     new (context) PythonCensusContext();
56     return;
57   }
58   if (!parent_span_id.empty()) {
59     // Note that parent_span_id exist also means it was marked as sampled at
60     // Python OC, we'll respect that decision.
61     SpanContext parent_context =
62         SpanContext(std::string(trace_id), std::string(parent_span_id), true);
63     new (context) PythonCensusContext(method, parent_context);
64     return;
65   }
66   // Create span without parent.
67   new (context) PythonCensusContext(method, trace_id);
68 }
69 
GenerateServerContext(absl::string_view header,absl::string_view method,PythonCensusContext * context)70 void GenerateServerContext(absl::string_view header, absl::string_view method,
71                            PythonCensusContext* context) {
72   // Destruct the current CensusContext to free the Span memory before
73   // overwriting it below.
74   context->~PythonCensusContext();
75   if (method.empty()) {
76     new (context) PythonCensusContext();
77     return;
78   }
79   SpanContext parent_ctx = FromGrpcTraceBinHeader(header);
80   if (parent_ctx.IsValid()) {
81     new (context) PythonCensusContext(method, parent_ctx);
82   } else {
83     new (context) PythonCensusContext(method);
84   }
85 }
86 
ToGrpcTraceBinHeader(const PythonCensusContext & ctx,uint8_t * out)87 void ToGrpcTraceBinHeader(const PythonCensusContext& ctx, uint8_t* out) {
88   out[kVersionOfs] = kVersionId;
89   out[kTraceIdOfs] = kTraceIdField;
90   uint8_t trace_options_rep_[kSizeTraceOptions];
91 
92   std::string trace_id =
93       absl::HexStringToBytes(absl::string_view(ctx.GetSpanContext().TraceId()));
94   std::string span_id =
95       absl::HexStringToBytes(absl::string_view(ctx.GetSpanContext().SpanId()));
96   trace_options_rep_[0] = ctx.GetSpanContext().IsSampled() ? 1 : 0;
97 
98   memcpy(reinterpret_cast<uint8_t*>(&out[kTraceIdOfs + 1]), trace_id.c_str(),
99          kSizeTraceID);
100 
101   out[kSpanIdOfs] = kSpanIdField;
102   memcpy(reinterpret_cast<uint8_t*>(&out[kSpanIdOfs + 1]), span_id.c_str(),
103          kSizeSpanID);
104 
105   out[kTraceOptionsOfs] = kTraceOptionsField;
106   memcpy(reinterpret_cast<uint8_t*>(&out[kTraceOptionsOfs + 1]),
107          trace_options_rep_, kSizeTraceOptions);
108 }
109 
FromGrpcTraceBinHeader(absl::string_view header)110 SpanContext FromGrpcTraceBinHeader(absl::string_view header) {
111   if (header.size() < kGrpcTraceBinHeaderLen ||
112       header[kVersionOfs] != kVersionId ||
113       header[kTraceIdOfs] != kTraceIdField ||
114       header[kSpanIdOfs] != kSpanIdField ||
115       header[kTraceOptionsOfs] != kTraceOptionsField) {
116     return SpanContext();  // Invalid.
117   }
118 
119   uint8_t options = header[kTraceOptionsOfs + 1] & 1;
120   constexpr uint8_t kIsSampled = 1;
121 
122   uint8_t trace_id_rep_[kTraceIdSize];
123   memcpy(trace_id_rep_,
124          reinterpret_cast<const uint8_t*>(&header[kTraceIdOfs + 1]),
125          kTraceIdSize);
126 
127   uint8_t span_id_rep_[kSpanIdSize];
128   memcpy(span_id_rep_,
129          reinterpret_cast<const uint8_t*>(&header[kSpanIdOfs + 1]),
130          kSpanIdSize);
131 
132   uint8_t trace_option_rep_[kTraceOptionsLen];
133   memcpy(trace_option_rep_, &options, kTraceOptionsLen);
134 
135   SpanContext context(
136       absl::BytesToHexString(absl::string_view(
137           reinterpret_cast<const char*>(trace_id_rep_), kTraceIdSize)),
138       absl::BytesToHexString(absl::string_view(
139           reinterpret_cast<const char*>(span_id_rep_), kSpanIdSize)),
140       trace_option_rep_[0] & kIsSampled);
141   return context;
142 }
143 
TraceContextSerialize(const PythonCensusContext & context,char * tracing_buf,size_t tracing_buf_size)144 size_t TraceContextSerialize(const PythonCensusContext& context,
145                              char* tracing_buf, size_t tracing_buf_size) {
146   if (tracing_buf_size < kGrpcTraceBinHeaderLen) {
147     return 0;
148   }
149   ToGrpcTraceBinHeader(context, reinterpret_cast<uint8_t*>(tracing_buf));
150   return kGrpcTraceBinHeaderLen;
151 }
152 
StatsContextSerialize(size_t,grpc_slice *)153 size_t StatsContextSerialize(size_t /*max_tags_len*/, grpc_slice* /*tags*/) {
154   return 0;
155 }
156 
ServerStatsDeserialize(const char * buf,size_t buf_size,uint64_t * server_elapsed_time)157 size_t ServerStatsDeserialize(const char* buf, size_t buf_size,
158                               uint64_t* server_elapsed_time) {
159   return RpcServerStatsEncoding::Decode(absl::string_view(buf, buf_size),
160                                         server_elapsed_time);
161 }
162 
ServerStatsSerialize(uint64_t server_elapsed_time,char * buf,size_t buf_size)163 size_t ServerStatsSerialize(uint64_t server_elapsed_time, char* buf,
164                             size_t buf_size) {
165   return RpcServerStatsEncoding::Encode(server_elapsed_time, buf, buf_size);
166 }
167 
GetIncomingDataSize(const grpc_call_final_info * final_info)168 uint64_t GetIncomingDataSize(const grpc_call_final_info* final_info) {
169   return final_info->stats.transport_stream_stats.incoming.data_bytes;
170 }
171 
GetOutgoingDataSize(const grpc_call_final_info * final_info)172 uint64_t GetOutgoingDataSize(const grpc_call_final_info* final_info) {
173   return final_info->stats.transport_stream_stats.outgoing.data_bytes;
174 }
175 
176 namespace {
177 // span_id is a 16-character hexadecimal encoded string.
GenerateSpanId()178 std::string GenerateSpanId() {
179   uint64_t span_id = absl::Uniform<uint64_t>(absl::BitGen());
180   std::stringstream hex_string;
181   hex_string << std::setfill('0') << std::setw(16) << std::hex << span_id;
182   return std::string(hex_string.str());
183 }
184 
185 // trace_id is a 32-character hexadecimal encoded string
GenerateTraceId()186 std::string GenerateTraceId() {
187   absl::uint128 trace_id = absl::Uniform<absl::uint128>(absl::BitGen());
188   std::stringstream hex_string;
189   hex_string << std::setfill('0') << std::setw(32) << std::hex << trace_id;
190   return std::string(hex_string.str());
191 }
192 
193 }  // namespace
194 
195 //
196 // Span
197 //
198 
StartSpan(absl::string_view name,const Span * parent)199 Span Span::StartSpan(absl::string_view name, const Span* parent) {
200   std::string span_id = GenerateSpanId();
201   std::string trace_id;
202   std::string parent_span_id;
203   bool should_sample;
204   auto start_time = absl::Now();
205 
206   if (parent != nullptr) {
207     parent_span_id = parent->Context().SpanId();
208     trace_id = parent->Context().TraceId();
209     should_sample = parent->Context().IsSampled();
210   } else {
211     trace_id = GenerateTraceId();
212     should_sample = ShouldSample(trace_id);
213   }
214 
215   SpanContext context = SpanContext(trace_id, span_id, should_sample);
216   return Span(std::string(name), parent_span_id, start_time, context);
217 }
218 
StartSpan(absl::string_view name,const SpanContext & parent_context)219 Span Span::StartSpan(absl::string_view name,
220                      const SpanContext& parent_context) {
221   std::string trace_id = parent_context.TraceId();
222   std::string parent_span_id = parent_context.SpanId();
223   std::string span_id = GenerateSpanId();
224   bool should_sample = parent_context.IsSampled();
225   if (!should_sample) {
226     // Resampling here so that it's possible to collect trace on server side
227     // if client tracing is not enabled.
228     should_sample = ShouldSample(std::string(trace_id));
229   }
230   auto start_time = absl::Now();
231   SpanContext context(trace_id, span_id, should_sample);
232   return Span(std::string(name), parent_span_id, start_time, context);
233 }
234 
StartSpan(absl::string_view name,absl::string_view trace_id)235 Span Span::StartSpan(absl::string_view name, absl::string_view trace_id) {
236   std::string span_id = GenerateSpanId();
237   auto start_time = absl::Now();
238   bool should_sample = ShouldSample(std::string(trace_id));
239   SpanContext context(std::string(trace_id), span_id, should_sample);
240   return Span(std::string(name), "", start_time, context);
241 }
242 
SetStatus(absl::string_view status)243 void Span::SetStatus(absl::string_view status) {
244   status_ = std::string(status);
245 }
246 
AddAttribute(absl::string_view key,absl::string_view value)247 void Span::AddAttribute(absl::string_view key, absl::string_view value) {
248   span_labels_.emplace_back(std::string(key), std::string(value));
249 }
250 
AddAnnotation(absl::string_view description)251 void Span::AddAnnotation(absl::string_view description) {
252   // Need a string format which can be converted to Python datetime.datetime
253   // class directly.
254   std::string time_stamp =
255       absl::FormatTime("%Y-%m-%d %H:%M:%E3S", absl::Now(), absl::UTCTimeZone());
256   span_annotations_.emplace_back(
257       Annotation{time_stamp, std::string(description)});
258 }
259 
ToCensusData() const260 SpanCensusData Span::ToCensusData() const {
261   SpanCensusData census_data;
262   absl::TimeZone utc = absl::UTCTimeZone();
263   census_data.name = name_;
264   // Need a string format which can be exported to StackDriver directly.
265   // See format details:
266   // https://cloud.google.com/trace/docs/reference/v2/rest/v2/projects.traces/batchWrite
267   census_data.start_time =
268       absl::FormatTime("%Y-%m-%dT%H:%M:%E6SZ", start_time_, utc);
269   census_data.end_time =
270       absl::FormatTime("%Y-%m-%dT%H:%M:%E6SZ", end_time_, utc);
271   census_data.trace_id = Context().TraceId();
272   census_data.span_id = Context().SpanId();
273   census_data.should_sample = Context().IsSampled();
274   census_data.parent_span_id = parent_span_id_;
275   census_data.status = status_;
276   census_data.span_labels = span_labels_;
277   census_data.span_annotations = span_annotations_;
278   census_data.child_span_count = child_span_count_;
279   return census_data;
280 }
281 
282 }  // namespace grpc_observability
283