1
#include "source/common/tracing/trace_context_impl.h"
2

            
3
#include "source/common/http/header_map_impl.h"
4

            
5
namespace Envoy {
6
namespace Tracing {
7

            
8
312
TraceContextHandler::TraceContextHandler(absl::string_view key) : key_(key) {
9
  // This will force the header map to be finalized in unit tests and do nothing in prod (
10
  // where the header map is already finalized when the server is initializing).
11
312
  Http::TypedHeaderMapImpl<Http::RequestHeaderMap>::inlineHeadersSize();
12

            
13
312
  handle_ = Http::CustomInlineHeaderRegistry::getInlineHeader<
14
312
      Http::CustomInlineHeaderRegistry::Type::RequestHeaders>(key_);
15
312
}
16

            
17
4
void TraceContextHandler::set(TraceContext& trace_context, absl::string_view value) const {
18
  // Will dynamic_cast be better?
19
4
  auto header_map = trace_context.requestHeaders();
20
4
  if (!header_map.has_value()) {
21
2
    trace_context.set(key_, value);
22
2
    return;
23
2
  }
24

            
25
2
  if (handle_.has_value()) {
26
1
    header_map->setInline(handle_.value(), value);
27
1
  } else {
28
1
    header_map->setCopy(key_, value);
29
1
  }
30
2
}
31

            
32
122
void TraceContextHandler::setRefKey(TraceContext& trace_context, absl::string_view value) const {
33
122
  auto header_map = trace_context.requestHeaders();
34
122
  if (!header_map.has_value()) {
35
46
    trace_context.set(key_, value);
36
46
    return;
37
46
  }
38

            
39
76
  if (handle_.has_value()) {
40
1
    header_map->setInline(handle_.value(), value);
41
75
  } else {
42
75
    header_map->setReferenceKey(key_, value);
43
75
  }
44
76
}
45

            
46
4
void TraceContextHandler::setRef(TraceContext& trace_context, absl::string_view value) const {
47
4
  auto header_map = trace_context.requestHeaders();
48
4
  if (!header_map.has_value()) {
49
2
    trace_context.set(key_, value);
50
2
    return;
51
2
  }
52

            
53
2
  if (handle_.has_value()) {
54
1
    header_map->setReferenceInline(handle_.value(), value);
55
1
  } else {
56
1
    header_map->setReference(key_, value);
57
1
  }
58
2
}
59

            
60
absl::optional<absl::string_view>
61
607
TraceContextHandler::get(const TraceContext& trace_context) const {
62
607
  auto header_map = trace_context.requestHeaders();
63
607
  if (!header_map.has_value()) {
64
555
    return trace_context.get(key_);
65
555
  }
66

            
67
52
  if (handle_.has_value()) {
68
3
    auto* entry = header_map->getInline(handle_.value());
69
3
    if (entry == nullptr) {
70
1
      return absl::nullopt;
71
1
    }
72
2
    return entry->value().getStringView();
73
49
  } else {
74
49
    auto results = header_map->get(key_);
75
49
    if (results.empty()) {
76
23
      return absl::nullopt;
77
23
    }
78
26
    return results[0]->value().getStringView();
79
49
  }
80
52
}
81

            
82
TraceContextHandler::GetAllResult
83
35
TraceContextHandler::getAll(const TraceContext& trace_context) const {
84
35
  auto header_map = trace_context.requestHeaders();
85
35
  if (!header_map.has_value()) {
86
18
    if (const auto value = trace_context.get(key_); value.has_value()) {
87
7
      return {value.value()};
88
7
    }
89
11
    return {};
90
18
  }
91

            
92
17
  if (handle_.has_value()) {
93
2
    auto* entry = header_map->getInline(handle_.value());
94
2
    if (entry == nullptr) {
95
1
      return {};
96
1
    }
97
1
    return {entry->value().getStringView()};
98
15
  } else {
99
15
    auto results = header_map->get(key_);
100
15
    if (results.empty()) {
101
6
      return {};
102
6
    }
103
9
    GetAllResult all_values;
104
9
    all_values.reserve(results.size());
105
21
    for (size_t i = 0; i < results.size(); ++i) {
106
12
      all_values.push_back(results[i]->value().getStringView());
107
12
    }
108
9
    return all_values;
109
15
  }
110
17
}
111

            
112
4
void TraceContextHandler::remove(TraceContext& trace_context) const {
113
4
  auto header_map = trace_context.requestHeaders();
114
4
  if (!header_map.has_value()) {
115
2
    trace_context.remove(key_);
116
2
    return;
117
2
  }
118

            
119
2
  if (handle_.has_value()) {
120
1
    header_map->removeInline(handle_.value());
121
1
  } else {
122
1
    header_map->remove(key_);
123
1
  }
124
2
}
125

            
126
} // namespace Tracing
127
} // namespace Envoy