propagateTraceContext extraction

Author: jaybcee

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/request/TraceContextPropagator.java

@@ -0,0 +1,33 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+package org.elasticsearch.xpack.inference.external.request;
+
+import org.apache.http.client.methods.HttpPost;
+import org.elasticsearch.tasks.Task;
+import org.elasticsearch.xpack.inference.telemetry.TraceContext;
+
+public final class TraceContextPropagator {
+    private TraceContextPropagator() {} // Utility class
+
+    public static void propagateTraceContext(HttpPost httpPost, TraceContext traceContext) {
+        if (traceContext == null) {
+            return;
+        }
+
+        var traceParent = traceContext.traceParent();
+        var traceState = traceContext.traceState();
+
+        if (traceParent != null) {
+            httpPost.setHeader(Task.TRACE_PARENT_HTTP_HEADER, traceParent);
+        }
+
+        if (traceState != null) {
+            httpPost.setHeader(Task.TRACE_STATE, traceState);
+        }
+    }
+}

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/request/elastic/ElasticInferenceServiceSparseEmbeddingsRequest.java

@@ -12,13 +12,13 @@
 import org.apache.http.entity.ByteArrayEntity;
 import org.apache.http.message.BasicHeader;
 import org.elasticsearch.common.Strings;
-import org.elasticsearch.tasks.Task;
 import org.elasticsearch.xcontent.XContentType;
 import org.elasticsearch.xpack.inference.common.Truncator;
 import org.elasticsearch.xpack.inference.external.request.HttpRequest;
 import org.elasticsearch.xpack.inference.external.request.Request;
 import org.elasticsearch.xpack.inference.services.elastic.ElasticInferenceServiceSparseEmbeddingsModel;
 import org.elasticsearch.xpack.inference.telemetry.TraceContext;
+import org.elasticsearch.xpack.inference.external.request.TraceContextPropagator;
 
 import java.net.URI;
 import java.nio.charset.StandardCharsets;
@@ -56,9 +56,7 @@ public HttpRequest createHttpRequest() {
         ByteArrayEntity byteEntity = new ByteArrayEntity(requestEntity.getBytes(StandardCharsets.UTF_8));
         httpPost.setEntity(byteEntity);
 
-        if (traceContext != null) {
-            propagateTraceContext(httpPost);
-        }
+        TraceContextPropagator.propagateTraceContext(httpPost, traceContext);
 
         httpPost.setHeader(new BasicHeader(HttpHeaders.CONTENT_TYPE, XContentType.JSON.mediaType()));
 
@@ -90,17 +88,4 @@ public Request truncate() {
     public boolean[] getTruncationInfo() {
         return truncationResult.truncated().clone();
     }
-
-    private void propagateTraceContext(HttpPost httpPost) {
-        var traceParent = traceContext.traceParent();
-        var traceState = traceContext.traceState();
-
-        if (traceParent != null) {
-            httpPost.setHeader(Task.TRACE_PARENT_HTTP_HEADER, traceParent);
-        }
-
-        if (traceState != null) {
-            httpPost.setHeader(Task.TRACE_STATE, traceState);
-        }
-    }
 }