trace: respond with trace context to report the sampling options

For cases where incoming request doesn't contain sampling options,
client determines whether the request should be traced. Determined
sampling option needs to be returned back to the source of the
incoming request in the response back.

Change-Id: I3ccce25d1e0d653a2d865fd4434e5a34d2eef61c
Reviewed-on: https://code-review.googlesource.com/12211
Reviewed-by: kokoro <[email protected]>
Reviewed-by: Ross Light <[email protected]>

Author: rakyll

trace/http.go

@@ -89,9 +89,28 @@ type handler struct {
 }
 
 func (h *handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
-	span := h.traceClient.SpanFromRequest(r)
+	traceID, parentSpanID, options, optionsOk, ok := traceInfoFromHeader(r.Header.Get(httpHeader))
+	if !ok {
+		traceID = nextTraceID()
+	}
+	t := &trace{
+		traceID:       traceID,
+		client:        h.traceClient,
+		globalOptions: options,
+		localOptions:  options,
+	}
+	span := startNewChildWithRequest(r, t, parentSpanID)
+	span.span.Kind = spanKindServer
+	span.rootSpan = true
+	configureSpanFromPolicy(span, h.traceClient.policy, ok)
 	defer span.Finish()
 
 	r = r.WithContext(NewContext(r.Context(), span))
+	if ok && !optionsOk {
+		// Inject the trace context back to the response with the sampling options.
+		// TODO(jbd): Remove when there is a better way to report the client's sampling.
+		w.Header().Set(httpHeader, spanHeader(traceID, parentSpanID, span.trace.localOptions))
+	}
 	h.handler.ServeHTTP(w, r)
+
 }

trace/http_test.go

@@ -107,3 +107,52 @@ func TestHTTPHandlerNoTrace(t *testing.T) {
 		t.Fatal(err)
 	}
 }
+
+func TestHTTPHandler_response(t *testing.T) {
+	tc := newTestClient(&noopTransport{})
+	p, _ := NewLimitedSampler(1, 1<<32) // all
+	tc.SetSamplingPolicy(p)
+	handler := tc.HTTPHandler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}))
+	ts := httptest.NewServer(handler)
+	defer ts.Close()
+
+	tests := []struct {
+		name            string
+		traceHeader     string
+		wantTraceHeader string
+	}{
+		{
+			name:            "no global",
+			traceHeader:     "0123456789ABCDEF0123456789ABCDEF/123",
+			wantTraceHeader: "0123456789ABCDEF0123456789ABCDEF/123;o=1",
+		},
+		{
+			name:            "global=1",
+			traceHeader:     "0123456789ABCDEF0123456789ABCDEF/123;o=1",
+			wantTraceHeader: "",
+		},
+		{
+			name:            "global=0",
+			traceHeader:     "0123456789ABCDEF0123456789ABCDEF/123;o=0",
+			wantTraceHeader: "",
+		},
+		{
+			name:            "no trace context",
+			traceHeader:     "",
+			wantTraceHeader: "",
+		},
+	}
+
+	for _, tt := range tests {
+		req, _ := http.NewRequest("GET", ts.URL, nil)
+		req.Header.Set(httpHeader, tt.traceHeader)
+
+		res, err := http.DefaultClient.Do(req)
+		if err != nil {
+			t.Errorf("failed to request: %v", err)
+		}
+		if got, want := res.Header.Get(httpHeader), tt.wantTraceHeader; got != want {
+			t.Errorf("%v: response context header = %q; want %q", tt.name, got, want)
+		}
+	}
+}

trace/trace.go

@@ -335,7 +335,7 @@ func (c *Client) SpanFromHeader(name string, header string) *Span {
 	if c == nil {
 		return nil
 	}
-	traceID, parentSpanID, options, ok := traceInfoFromHeader(header)
+	traceID, parentSpanID, options, _, ok := traceInfoFromHeader(header)
 	if !ok {
 		traceID = nextTraceID()
 	}
@@ -373,7 +373,7 @@ func (c *Client) SpanFromRequest(r *http.Request) *Span {
 	if c == nil {
 		return nil
 	}
-	traceID, parentSpanID, options, ok := traceInfoFromHeader(r.Header.Get(httpHeader))
+	traceID, parentSpanID, options, _, ok := traceInfoFromHeader(r.Header.Get(httpHeader))
 	if !ok {
 		traceID = nextTraceID()
 	}
@@ -446,20 +446,22 @@ func FromContext(ctx context.Context) *Span {
 	return s
 }
 
-func traceInfoFromHeader(h string) (string, uint64, optionFlags, bool) {
+func traceInfoFromHeader(h string) (traceID string, spanID uint64, options optionFlags, optionsOk bool, ok bool) {
 	// See https://cloud.google.com/trace/docs/faq for the header format.
 	// Return if the header is empty or missing, or if the header is unreasonably
 	// large, to avoid making unnecessary copies of a large string.
 	if h == "" || len(h) > 200 {
-		return "", 0, 0, false
+		return "", 0, 0, false, false
+
 	}
 
 	// Parse the trace id field.
 	slash := strings.Index(h, `/`)
 	if slash == -1 {
-		return "", 0, 0, false
+		return "", 0, 0, false, false
+
 	}
-	traceID, h := h[:slash], h[slash+1:]
+	traceID, h = h[:slash], h[slash+1:]
 
 	// Parse the span id field.
 	spanstr := h
@@ -469,19 +471,22 @@ func traceInfoFromHeader(h string) (string, uint64, optionFlags, bool) {
 	}
 	spanID, err := strconv.ParseUint(spanstr, 10, 64)
 	if err != nil {
-		return "", 0, 0, false
+		return "", 0, 0, false, false
+
 	}
 
 	// Parse the options field, options field is optional.
 	if !strings.HasPrefix(h, "o=") {
-		return traceID, spanID, 0, true
+		return traceID, spanID, 0, false, true
+
 	}
 	o, err := strconv.ParseUint(h[2:], 10, 64)
 	if err != nil {
-		return "", 0, 0, false
+		return "", 0, 0, false, false
+
 	}
-	options := optionFlags(o)
-	return traceID, spanID, options, true
+	options = optionFlags(o)
+	return traceID, spanID, options, true, true
 }
 
 type optionFlags uint32

trace/trace_test.go

@@ -231,7 +231,7 @@ func TestHeader(t *testing.T) {
 		},
 	}
 	for _, tt := range tests {
-		traceID, parentSpanID, opts, ok := traceInfoFromHeader(tt.header)
+		traceID, parentSpanID, opts, _, ok := traceInfoFromHeader(tt.header)
 		if got, want := traceID, tt.wantTraceID; got != want {
 			t.Errorf("TraceID(%v) = %q; want %q", tt.header, got, want)
 		}