feat(query): add trace response headers

http/query_handler.go

@@ -9,8 +9,11 @@ import (
 	"io/ioutil"
 	"net/http"
 	"net/url"
+	"strconv"
 	"time"
 
+	"github.com/influxdata/influxdb/logger"
+
 	"github.com/NYTimes/gziphandler"
 	"github.com/influxdata/flux"
 	"github.com/influxdata/flux/ast"
@@ -31,7 +34,9 @@ import (
 )
 
 const (
-	fluxPath = "/api/v2/query"
+	fluxPath           = "/api/v2/query"
+	traceIDHeader      = "Trace-Id"
+	traceSampledHeader = "Trace-Sampled"
 )
 
 // FluxBackend is all services and associated parameters required to construct
@@ -104,6 +109,11 @@ func (h *FluxHandler) handleQuery(w http.ResponseWriter, r *http.Request) {
 	defer span.Finish()
 
 	ctx := r.Context()
+	if id, sampled, found := logger.TraceInfo(ctx); found {
+		println(id, sampled)
+		w.Header().Set(traceIDHeader, id)
+		w.Header().Set(traceSampledHeader, strconv.FormatBool(sampled))
+	}
 
 	// TODO(desa): I really don't like how we're recording the usage metrics here
 	// Ideally this will be moved when we solve https://github.com/influxdata/influxdb/issues/13403

http/query_handler_test.go

@@ -15,6 +15,8 @@ import (
 	"testing"
 	"time"
 
+	"github.com/influxdata/influxdb/kit/tracing"
+
 	"github.com/google/go-cmp/cmp"
 	"github.com/influxdata/flux"
 	"github.com/influxdata/flux/csv"
@@ -321,6 +323,8 @@ var _ metric.EventRecorder = noopEventRecorder{}
 
 // Certain error cases must be encoded as influxdb.Error so they can be properly decoded clientside.
 func TestFluxHandler_PostQuery_Errors(t *testing.T) {
+	defer tracing.JaegerTestSetupAndTeardown(t.Name())()
+
 	i := inmem.NewService()
 	b := &FluxBackend{
 		HTTPErrorHandler:    ErrorHandler(0),
@@ -349,6 +353,14 @@ func TestFluxHandler_PostQuery_Errors(t *testing.T) {
 
 		defer resp.Body.Close()
 
+		if actual := resp.Header.Get("Trace-Id"); actual == "" {
+			t.Error("expected trace ID header")
+		}
+
+		if actual := resp.Header.Get("Trace-Sampled"); actual == "" {
+			t.Error("expected trace sampled header")
+		}
+
 		if resp.StatusCode != http.StatusUnauthorized {
 			t.Errorf("expected unauthorized status, got %d", resp.StatusCode)
 		}
@@ -380,6 +392,14 @@ func TestFluxHandler_PostQuery_Errors(t *testing.T) {
 
 		h.handleQuery(w, req)
 
+		if actual := w.Header().Get("Trace-Id"); actual == "" {
+			t.Error("expected trace ID header")
+		}
+
+		if actual := w.Header().Get("Trace-Sampled"); actual == "" {
+			t.Error("expected trace sampled header")
+		}
+
 		if w.Code != http.StatusBadRequest {
 			t.Errorf("expected bad request status, got %d", w.Code)
 		}
@@ -413,6 +433,14 @@ func TestFluxHandler_PostQuery_Errors(t *testing.T) {
 		w := httptest.NewRecorder()
 		h.handleQuery(w, req)
 
+		if actual := w.Header().Get("Trace-Id"); actual == "" {
+			t.Error("expected trace ID header")
+		}
+
+		if actual := w.Header().Get("Trace-Sampled"); actual == "" {
+			t.Error("expected trace sampled header")
+		}
+
 		if w.Code != http.StatusBadRequest {
 			t.Errorf("expected bad request status, got %d", w.Code)
 		}

http/swagger.yml

@@ -3208,6 +3208,16 @@ paths:
                     enum:
                     - gzip
                     - identity
+              Trace-Id:
+                description: The Trace-Id header reports the request's trace ID, if one was generated.
+                schema:
+                  type: string
+                  description: Specifies the request's trace ID.
+              Trace-Sampled:
+                description: The Trace-Sampled header reports whether the trace was sampled.
+                schema:
+                  type: boolean
+                  description: Specifies whether the trace was sampled.
             content:
               text/csv:
                 schema:

kit/tracing/tracing.go

@@ -7,6 +7,8 @@ import (
 	"runtime"
 	"strings"
 
+	"github.com/uber/jaeger-client-go"
+
 	"github.com/influxdata/httprouter"
 	"github.com/opentracing/opentracing-go"
 	"github.com/opentracing/opentracing-go/ext"
@@ -165,3 +167,18 @@ func StartSpanFromContextWithOperationName(ctx context.Context, operationName st
 
 	return span, ctx
 }
+
+// JaegerTestSetupAndTeardown sets the global tracer to an in memory Jaeger instance for testing.
+// The returned function should be deferred by the caller to tear down this setup after testing is complete.
+func JaegerTestSetupAndTeardown(name string) func() {
+	old := opentracing.GlobalTracer()
+	tracer, closer := jaeger.NewTracer(name,
+		jaeger.NewConstSampler(true),
+		jaeger.NewInMemoryReporter(),
+	)
+	opentracing.SetGlobalTracer(tracer)
+	return func() {
+		_ = closer.Close()
+		opentracing.SetGlobalTracer(old)
+	}
+}

logger/fields.go

@@ -94,6 +94,7 @@ func TraceInfo(ctx context.Context) (traceID string, sampled bool, found bool) {
 		if spanContext, ok := span.Context().(jaeger.SpanContext); ok {
 			traceID = spanContext.TraceID().String()
 			sampled = spanContext.IsSampled()
+			println(traceID, sampled)
 			return traceID, sampled, true
 		}
 	}