77
88import datadog .trace .api .Config ;
99import datadog .trace .api .DDSpanId ;
10+ import datadog .trace .api .DDTraceId ;
1011import datadog .trace .api .naming .SpanNaming ;
1112import datadog .trace .bootstrap .ContextStore ;
1213import datadog .trace .bootstrap .instrumentation .api .AgentScope ;
2021import datadog .trace .bootstrap .instrumentation .jdbc .DBQueryInfo ;
2122import datadog .trace .bootstrap .instrumentation .jdbc .JDBCConnectionUrlParser ;
2223import java .io .ByteArrayOutputStream ;
24+ import java .io .DataOutputStream ;
2325import java .sql .Connection ;
2426import java .sql .DatabaseMetaData ;
2527import java .sql .PreparedStatement ;
@@ -244,16 +246,6 @@ public boolean isSqlServer(final DBInfo dbInfo) {
244246 return "sqlserver" .equals (dbInfo .getType ());
245247 }
246248
247- public static byte [] hexStringToByteArray (String s ) {
248- int len = s .length ();
249- byte [] data = new byte [len / 2 ];
250- for (int i = 0 ; i < len ; i += 2 ) {
251- data [i / 2 ] =
252- (byte ) ((Character .digit (s .charAt (i ), 16 ) << 4 ) + Character .digit (s .charAt (i + 1 ), 16 ));
253- }
254- return data ;
255- }
256-
257249 /**
258250 * Executes a `SET CONTEXT_INFO` statement on the DB with the active trace ID and the given span
259251 * ID. This context will be "attached" to future queries on the same connection. See <a
@@ -269,21 +261,28 @@ public static byte[] hexStringToByteArray(String s) {
269261 * @return spanID pre-created spanID
270262 */
271263 public long setContextInfo (Connection connection , DBInfo dbInfo ) {
272- final String VERSION = "0" ;
264+ final byte VERSION = 0 ;
273265 final long spanID = Config .get ().getIdGenerationStrategy ().generateSpanId ();
274266 AgentSpan instrumentationSpan =
275267 AgentTracer .get ().buildSpan ("set context_info" ).withTag ("dd.instrumentation" , true ).start ();
276268 DECORATE .afterStart (instrumentationSpan );
277269 DECORATE .onConnection (instrumentationSpan , dbInfo );
278270 PreparedStatement instrumentationStatement = null ;
279271 try (AgentScope scope = activateSpan (instrumentationSpan )) {
280- String samplingDecision = instrumentationSpan .forceSamplingDecision () > 0 ? "1" : "0" ;
281-
282- ByteArrayOutputStream outputStream = new ByteArrayOutputStream ();
283- outputStream .write (hexStringToByteArray (VERSION + samplingDecision ));
284- outputStream .write (hexStringToByteArray (DDSpanId .toHexStringPadded (spanID )));
285- outputStream .write (hexStringToByteArray (instrumentationSpan .getTraceId ().toHexString ()));
286- final byte [] contextInfo = outputStream .toByteArray ();
272+ final byte samplingDecision =
273+ (byte ) (instrumentationSpan .forceSamplingDecision () > 0 ? 1 : 0 );
274+ final byte versionAndSamplingDecision =
275+ (byte ) ((VERSION << 4 ) & 0b11110000 | samplingDecision & 0b00000001);
276+
277+ ByteArrayOutputStream byteArrayOutputStream = new ByteArrayOutputStream ();
278+ DataOutputStream dataOutputStream = new DataOutputStream (byteArrayOutputStream );
279+ dataOutputStream .writeByte (versionAndSamplingDecision );
280+ dataOutputStream .writeLong (spanID );
281+ final DDTraceId traceId = instrumentationSpan .getTraceId ();
282+ dataOutputStream .writeLong (traceId .toHighOrderLong ());
283+ dataOutputStream .writeLong (traceId .toLong ());
284+ dataOutputStream .flush ();
285+ final byte [] contextInfo = byteArrayOutputStream .toByteArray ();
287286
288287 String instrumentationSql = "set context_info ?" ;
289288 instrumentationStatement = connection .prepareStatement (instrumentationSql );
0 commit comments