xref: /aosp_15_r20/external/federated-compute/fcp/tracing/tracing_context_utils.h (revision 14675a029014e728ec732f129a32e299b2da0601)
1 // Copyright 2020 Google LLC
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 #ifndef FCP_TRACING_TRACING_CONTEXT_UTILS_H_
16 #define FCP_TRACING_TRACING_CONTEXT_UTILS_H_
17 
18 #include <string>
19 
20 #include "google/protobuf/descriptor.h"
21 #include "google/protobuf/message.h"
22 #include "fcp/base/monitoring.h"
23 #include "fcp/tracing/tracing_traits.h"
24 #include "flatbuffers/flatbuffers.h"
25 
26 namespace fcp::tracing_internal {
27 
28 constexpr char kContextFieldName[] = "tracing_context";
29 constexpr char kContextWrongTypeMessage[] =
30     "Type of tracing_context field should be bytes or string";
31 
32 // Given a proto message, checks to see if it has a tracing_context field and
33 // if so, serializes the provided `context` message and sets the field contents
34 // to the serialized context.
35 // If no tracing_context field is present on `message` or it is of a type other
36 // than string or bytes, this will be a no-op.
37 void SetTracingContextOnMessage(const google::protobuf::Message& context,
38                                 google::protobuf::Message& message);
39 
40 // Given a proto `message` which may have a field tracing_context which contains
41 // a serialized tracing context proto of type ContextT, uses reflection to
42 // extract and parse the serialized proto to return a ContextT.
43 // If the proto `message` does not have a tracing_context field, or it is empty,
44 // returns the default value of ContextT.
45 // If the tracing_context field is of a type other than string or bytes, this
46 // will fail.
47 template <class ContextT>
GetContextFromMessage(const google::protobuf::Message & message)48 ContextT GetContextFromMessage(const google::protobuf::Message& message) {
49   const google::protobuf::FieldDescriptor* field_descriptor =
50       message.GetDescriptor()->FindFieldByName(kContextFieldName);
51   ContextT context;
52   if (field_descriptor == nullptr) {
53     return context;
54   }
55   FCP_CHECK(field_descriptor->type() == google::protobuf::FieldDescriptor::TYPE_BYTES ||
56             field_descriptor->type() == google::protobuf::FieldDescriptor::TYPE_STRING)
57       << kContextWrongTypeMessage;
58   std::string serialized_context =
59       message.GetReflection()->GetString(message, field_descriptor);
60   if (serialized_context.empty()) {
61     return context;
62   }
63   FCP_CHECK(context.ParseFromString(serialized_context));
64 
65   return context;
66 }
67 
68 }  // namespace fcp::tracing_internal
69 
70 #endif  // FCP_TRACING_TRACING_CONTEXT_UTILS_H_
71