| // Copyright 2017 Google Inc. All Rights Reserved. |
| // |
| // Licensed under the Apache License, Version 2.0 (the "License"); |
| // you may not use this file except in compliance with the License. |
| // You may obtain a copy of the License at |
| // |
| // http://www.apache.org/licenses/LICENSE-2.0 |
| // |
| // Unless required by applicable law or agreed to in writing, software |
| // distributed under the License is distributed on an "AS IS" BASIS, |
| // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| // See the License for the specific language governing permissions and |
| // limitations under the License. |
| |
| // +build go1.7 |
| |
| package trace |
| |
| import "net/http" |
| |
| type tracerTransport struct { |
| base http.RoundTripper |
| } |
| |
| func (tt *tracerTransport) RoundTrip(req *http.Request) (*http.Response, error) { |
| span := FromContext(req.Context()).NewRemoteChild(req) |
| resp, err := tt.base.RoundTrip(req) |
| |
| // TODO(jbd): Is it possible to defer the span.Finish? |
| // In cases where RoundTrip panics, we still can finish the span. |
| span.Finish(WithResponse(resp)) |
| return resp, err |
| } |
| |
| // HTTPClient is an HTTP client that enhances http.Client |
| // with automatic tracing support. |
| type HTTPClient struct { |
| http.Client |
| traceClient *Client |
| } |
| |
| // Do behaves like (*http.Client).Do but automatically traces |
| // outgoing requests if tracing is enabled for the current request. |
| // |
| // If req.Context() contains a traced *Span, the outgoing request |
| // is traced with the existing span. If not, the request is not traced. |
| func (c *HTTPClient) Do(req *http.Request) (*http.Response, error) { |
| return c.Client.Do(req) |
| } |
| |
| // NewHTTPClient creates a new HTTPClient that will trace the outgoing |
| // requests using tc. The attributes of this client are inherited from the |
| // given http.Client. If orig is nil, http.DefaultClient is used. |
| func (c *Client) NewHTTPClient(orig *http.Client) *HTTPClient { |
| if orig == nil { |
| orig = http.DefaultClient |
| } |
| rt := orig.Transport |
| if rt == nil { |
| rt = http.DefaultTransport |
| } |
| client := http.Client{ |
| Transport: &tracerTransport{base: rt}, |
| CheckRedirect: orig.CheckRedirect, |
| Jar: orig.Jar, |
| Timeout: orig.Timeout, |
| } |
| return &HTTPClient{ |
| Client: client, |
| traceClient: c, |
| } |
| } |
| |
| // HTTPHandler returns a http.Handler from the given handler |
| // that is aware of the incoming request's span. |
| // The span can be extracted from the incoming request in handler |
| // functions from incoming request's context: |
| // |
| // span := trace.FromContext(r.Context()) |
| // |
| // The span will be auto finished by the handler. |
| func (c *Client) HTTPHandler(h http.Handler) http.Handler { |
| return &handler{traceClient: c, handler: h} |
| } |
| |
| type handler struct { |
| traceClient *Client |
| handler http.Handler |
| } |
| |
| func (h *handler) ServeHTTP(w http.ResponseWriter, r *http.Request) { |
| 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) |
| |
| } |