adds set_kind func for oc_trace

census-instrumentation · Aug 14, 2020 · b522ef0 · b522ef0
1 parent 33de39d
commit b522ef0
Show file tree

Hide file tree

Showing 3 changed files with 18 additions and 10 deletions.
diff --git a/src/oc_propagation_http_b3.erl b/src/oc_propagation_http_b3.erl
@@ -39,21 +39,14 @@ to_headers(#span_ctx{trace_id=TraceId,
                      trace_options=TraceOptions}) ->
     Options = case TraceOptions band 1 of 1 -> "1"; _ -> "0" end,
     %% iolist_to_binary only needed for versions before otp-21
-    EncodedTraceId = encode_trace_id(TraceId),
+    EncodedTraceId = iolist_to_binary(io_lib:format("~32.16.0b", [TraceId])),
     EncodedSpanId = iolist_to_binary(io_lib:format("~16.16.0b", [SpanId])),
     [{?B3_TRACE_ID, EncodedTraceId},
      {?B3_SPAN_ID, EncodedSpanId},
      {?B3_SAMPLED, Options}];
 to_headers(undefined) ->
     [].
 
-encode_trace_id(TraceId) ->
-    encode_trace_id(TraceId, binary:encode_unsigned(TraceId)).
-encode_trace_id(TraceId, TraceIdBin) when bit_size(TraceIdBin) == 64 ->
-    iolist_to_binary(io_lib:format("~16.16.0b", [TraceId]));
-encode_trace_id(TraceId, _TraceIdBin) ->
-    iolist_to_binary(io_lib:format("~32.16.0b", [TraceId])).
-
 -spec from_headers(list() | map()) -> maybe(opencensus:span_ctx()).
 from_headers(Headers) when is_map(Headers) ->
     from_headers(maps:to_list(Headers));

diff --git a/src/oc_span.erl b/src/oc_span.erl
@@ -32,6 +32,8 @@
 
          message_event/4,
 
+         set_kind/2,
+
          set_status/3]).
 
 -include("opencensus.hrl").
@@ -138,6 +140,16 @@ set_status(Code, Message, Span=#span{}) ->
 set_status(_, _, undefined) ->
     undefined.
 
+%%--------------------------------------------------------------------
+%% @doc
+%% Set Kind.
+%% @end
+%%--------------------------------------------------------------------
+-spec set_kind(Kind, Span) -> Span when Kind :: opencensus:span_kind(), Span :: maybe(opencensus:span()).
+set_kind(Kind, Span=#span{}) ->
+    Span#span{kind=Kind};
+set_kind(_, undefined) ->
+    undefined.
 
 %%--------------------------------------------------------------------
 %% @doc

diff --git a/test/oc_span_SUITE.erl b/test/oc_span_SUITE.erl
@@ -53,10 +53,13 @@ modifications(_Config) ->
     Span6 = oc_span:add_link(Link, Span5),
     ?assertEqual(undefined, oc_span:add_link(Link, undefined)),
 
-    ?assertEqual({error, no_report_buffer}, oc_span:finish_span(#span_ctx{}, Span6)),
+    Span7 = oc_span:set_kind(?SPAN_KIND_SERVER, Span6),
+    ?assertEqual(undefined, oc_span:set_kind(?SPAN_KIND_SERVER, undefined)),
+
+    ?assertEqual({error, no_report_buffer}, oc_span:finish_span(#span_ctx{}, Span7)),
 
     {ok, _} = application:ensure_all_started(opencensus),
-    ?assertEqual(true, oc_span:finish_span(#span_ctx{}, Span6)),
+    ?assertEqual(true, oc_span:finish_span(#span_ctx{}, Span7)),
     ?assertEqual(true, oc_span:finish_span(#span_ctx{}, undefined)),
 
     application:stop(opencensus).