Skip to content

Commit 1ac4cea

Browse files
committed
Use Carrier.LookupKey if available for single key propagation.
1 parent bdf26c3 commit 1ac4cea

File tree

2 files changed

+115
-49
lines changed

2 files changed

+115
-49
lines changed

src/propagation.cpp

Lines changed: 56 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -43,6 +43,39 @@ static uint64_t HexToUint64(opentracing::string_view s) {
4343
return x;
4444
}
4545

46+
//------------------------------------------------------------------------------
47+
// LookupKey
48+
//------------------------------------------------------------------------------
49+
template <class KeyCompare>
50+
static opentracing::expected<opentracing::string_view> LookupKey(
51+
const opentracing::TextMapReader& carrier, opentracing::string_view key,
52+
KeyCompare key_compare) {
53+
// First try carrier.LookupKey since that can potentially be the fastest
54+
// approach.
55+
auto result = carrier.LookupKey(key);
56+
if (result || result.error() != opentracing::lookup_key_not_supported_error) {
57+
return result;
58+
}
59+
60+
// Fall back to iterating through all of the keys.
61+
result = opentracing::make_unexpected(opentracing::key_not_found_error);
62+
auto was_successful = carrier.ForeachKey(
63+
[&](opentracing::string_view key,
64+
opentracing::string_view value) -> opentracing::expected<void> {
65+
if (!key_compare(key, key)) {
66+
return {};
67+
}
68+
result = value;
69+
70+
// Found key, so bail out of the loop with a success error code.
71+
return opentracing::make_unexpected(std::error_code{});
72+
});
73+
if (!was_successful && was_successful.error() != std::error_code{}) {
74+
return opentracing::make_unexpected(was_successful.error());
75+
}
76+
return result;
77+
}
78+
4679
//------------------------------------------------------------------------------
4780
// InjectSpanContextMultiKey
4881
//------------------------------------------------------------------------------
@@ -246,41 +279,33 @@ static opentracing::expected<bool> ExtractSpanContextSingleKey(
246279
const opentracing::TextMapReader& carrier, uint64_t& trace_id,
247280
uint64_t& span_id, std::unordered_map<std::string, std::string>& baggage,
248281
KeyCompare key_compare) {
282+
auto value_maybe = LookupKey(carrier, PropagationSingleKey, key_compare);
283+
if (!value_maybe) {
284+
if (value_maybe.error() == opentracing::key_not_found_error) {
285+
return false;
286+
} else {
287+
return opentracing::make_unexpected(value_maybe.error());
288+
}
289+
}
290+
auto value = *value_maybe;
249291
std::stringstream stream;
250-
auto result = carrier.ForeachKey(
251-
[&](opentracing::string_view key,
252-
opentracing::string_view value) -> opentracing::expected<void> {
253-
if (!key_compare(key, PropagationSingleKey)) {
254-
return {};
255-
}
256-
257-
// Decode `value` into `stream`.
258-
try {
259-
in_memory_stream istream{value.data(), value.size()};
260-
base64::decoder decoder;
261-
decoder.decode(istream, stream);
262-
263-
// Flush so that when we call stream.good(), we'll get an accurate
264-
// view of the error state.
265-
stream.flush();
266-
if (!stream.good()) {
267-
return opentracing::make_unexpected(
268-
std::make_error_code(std::errc::io_error));
269-
}
270-
271-
} catch (const std::bad_alloc&) {
272-
return opentracing::make_unexpected(
273-
std::make_error_code(std::errc::not_enough_memory));
274-
}
292+
try {
293+
in_memory_stream istream{value.data(), value.size()};
294+
base64::decoder decoder;
295+
decoder.decode(istream, stream);
275296

276-
// Break out of the loop with success code.
277-
return opentracing::make_unexpected(std::error_code{});
278-
});
297+
// Flush so that when we call stream.good(), we'll get an accurate
298+
// view of the error state.
299+
stream.flush();
300+
if (!stream.good()) {
301+
return opentracing::make_unexpected(
302+
std::make_error_code(std::errc::io_error));
303+
}
279304

280-
if (!result && result.error() != std::error_code{}) {
281-
return opentracing::make_unexpected(result.error());
305+
} catch (const std::bad_alloc&) {
306+
return opentracing::make_unexpected(
307+
std::make_error_code(std::errc::not_enough_memory));
282308
}
283-
284309
return ExtractSpanContext(PropagationOptions{}, stream, trace_id, span_id,
285310
baggage);
286311
}

test/propagation_test.cpp

Lines changed: 59 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -15,48 +15,71 @@
1515
#include <lightstep/catch/catch.hpp>
1616

1717
using namespace lightstep;
18-
using namespace opentracing;
1918

2019
//------------------------------------------------------------------------------
2120
// TextMapCarrier
2221
//------------------------------------------------------------------------------
23-
struct TextMapCarrier : TextMapReader, TextMapWriter {
22+
struct TextMapCarrier : opentracing::TextMapReader, opentracing::TextMapWriter {
2423
TextMapCarrier(std::unordered_map<std::string, std::string>& text_map_)
2524
: text_map(text_map_) {}
2625

27-
expected<void> Set(string_view key, string_view value) const override {
26+
opentracing::expected<void> Set(
27+
opentracing::string_view key,
28+
opentracing::string_view value) const override {
2829
text_map[key] = value;
2930
return {};
3031
}
3132

32-
expected<void> ForeachKey(
33-
std::function<expected<void>(string_view key, string_view value)> f)
34-
const override {
33+
opentracing::expected<opentracing::string_view> LookupKey(
34+
opentracing::string_view key) const override {
35+
if (!supports_lookup) {
36+
return opentracing::make_unexpected(
37+
opentracing::lookup_key_not_supported_error);
38+
}
39+
auto iter = text_map.find(key);
40+
if (iter != text_map.end()) {
41+
return opentracing::string_view{iter->second};
42+
} else {
43+
return opentracing::make_unexpected(opentracing::key_not_found_error);
44+
}
45+
}
46+
47+
opentracing::expected<void> ForeachKey(
48+
std::function<opentracing::expected<void>(opentracing::string_view key,
49+
opentracing::string_view value)>
50+
f) const override {
51+
++foreach_key_call_count;
3552
for (const auto& key_value : text_map) {
3653
auto result = f(key_value.first, key_value.second);
3754
if (!result) return result;
3855
}
3956
return {};
4057
}
4158

59+
bool supports_lookup = false;
60+
mutable int foreach_key_call_count = 0;
4261
std::unordered_map<std::string, std::string>& text_map;
4362
};
4463

4564
//------------------------------------------------------------------------------
4665
// HTTPHeadersCarrier
4766
//------------------------------------------------------------------------------
48-
struct HTTPHeadersCarrier : HTTPHeadersReader, HTTPHeadersWriter {
67+
struct HTTPHeadersCarrier : opentracing::HTTPHeadersReader,
68+
opentracing::HTTPHeadersWriter {
4969
HTTPHeadersCarrier(std::unordered_map<std::string, std::string>& text_map_)
5070
: text_map(text_map_) {}
5171

52-
expected<void> Set(string_view key, string_view value) const override {
72+
opentracing::expected<void> Set(
73+
opentracing::string_view key,
74+
opentracing::string_view value) const override {
5375
text_map[key] = value;
5476
return {};
5577
}
5678

57-
expected<void> ForeachKey(
58-
std::function<expected<void>(string_view key, string_view value)> f)
59-
const override {
79+
opentracing::expected<void> ForeachKey(
80+
std::function<opentracing::expected<void>(opentracing::string_view key,
81+
opentracing::string_view value)>
82+
f) const override {
6083
for (const auto& key_value : text_map) {
6184
auto result = f(key_value.first, key_value.second);
6285
if (!result) return result;
@@ -132,13 +155,13 @@ TEST_CASE("propagation") {
132155

133156
SECTION(
134157
"Injecting a non-LightStep span returns invalid_span_context_error.") {
135-
auto noop_tracer = MakeNoopTracer();
158+
auto noop_tracer = opentracing::MakeNoopTracer();
136159
CHECK(noop_tracer);
137160
auto span = noop_tracer->StartSpan("a");
138161
CHECK(span);
139162
auto was_successful = tracer->Inject(span->context(), text_map_carrier);
140163
CHECK(!was_successful);
141-
CHECK(was_successful.error() == invalid_span_context_error);
164+
CHECK(was_successful.error() == opentracing::invalid_span_context_error);
142165
}
143166

144167
SECTION(
@@ -152,7 +175,8 @@ TEST_CASE("propagation") {
152175
text_map.erase(std::begin(text_map));
153176
auto span_context_maybe = tracer->Extract(text_map_carrier);
154177
CHECK(!span_context_maybe);
155-
CHECK(span_context_maybe.error() == span_context_corrupted_error);
178+
CHECK(span_context_maybe.error() ==
179+
opentracing::span_context_corrupted_error);
156180
}
157181

158182
SECTION("Extract is insensitive to changes in case for http header fields") {
@@ -192,7 +216,8 @@ TEST_CASE("propagation") {
192216
std::istringstream iss{invalid_context, std::ios::binary};
193217
auto span_context_maybe = tracer->Extract(iss);
194218
CHECK(!span_context_maybe);
195-
CHECK(span_context_maybe.error() == span_context_corrupted_error);
219+
CHECK(span_context_maybe.error() ==
220+
opentracing::span_context_corrupted_error);
196221
}
197222

198223
SECTION("Calling Extract on an empty stream yields a nullptr.") {
@@ -218,8 +243,6 @@ TEST_CASE("propagation - single key") {
218243
CHECK(tracer->Inject(span->context(), text_map_carrier));
219244
CHECK(text_map.size() == 1);
220245

221-
for (auto& kv : text_map) std::cout << kv.first << ": " << kv.second << "\n";
222-
223246
SECTION("Inject, extract, inject yields the same text_map.") {
224247
auto injection_map1 = text_map;
225248
auto span_context_maybe = tracer->Extract(text_map_carrier);
@@ -229,8 +252,26 @@ TEST_CASE("propagation - single key") {
229252
CHECK(injection_map1 == text_map);
230253
}
231254

255+
SECTION("If a carrier supports LookupKey, then ForeachKey won't be called") {
256+
text_map_carrier.supports_lookup = true;
257+
auto span_context_maybe = tracer->Extract(text_map_carrier);
258+
CHECK((span_context_maybe && span_context_maybe->get()));
259+
CHECK(text_map_carrier.foreach_key_call_count == 0);
260+
}
261+
262+
SECTION(
263+
"When LookupKey is used, a nullptr is returned if there is no "
264+
"span_context") {
265+
text_map.clear();
266+
text_map_carrier.supports_lookup = true;
267+
auto span_context_maybe = tracer->Extract(text_map_carrier);
268+
CHECK((span_context_maybe && span_context_maybe->get() == nullptr));
269+
CHECK(text_map_carrier.foreach_key_call_count == 0);
270+
}
271+
232272
SECTION("Verify only valid base64 characters are used.") {
233-
// Follows the guidelines given in RFC-4648. See
273+
// Follows the guidelines given in RFC-4648 on what characters are
274+
// permissible. See
234275
// http://www.rfc-editor.org/rfc/rfc4648.txt
235276
auto iter = text_map.begin();
236277
CHECK(iter != text_map.end());

0 commit comments

Comments
 (0)