github.com/ManabuSeki/goa-v1@v1.4.3/middleware/xray/transport.go (about)

     1  package xray
     2  
     3  import (
     4  	"net/http"
     5  )
     6  
     7  // WrapTransport wraps a http RoundTripper with a RoundTripper which creates subsegments of the
     8  // segment in each request's context. The subsegments created this way have their namespace set to
     9  // "remote". The request's ctx must be set and contain the current request segment as set by the
    10  // xray middleware.
    11  //
    12  // Example of how to wrap http.Client's transport:
    13  //   httpClient := &http.Client{
    14  //     Transport: WrapTransport(http.DefaultTransport),
    15  //   }
    16  func WrapTransport(rt http.RoundTripper) http.RoundTripper {
    17  	return &xrayTransport{rt}
    18  }
    19  
    20  // xrayTransport wraps an http RoundTripper to add a tracing subsegment of the request's context segment
    21  type xrayTransport struct {
    22  	wrapped http.RoundTripper
    23  }
    24  
    25  // RoundTrip wraps the original RoundTripper.RoundTrip to create xray tracing segments
    26  func (t *xrayTransport) RoundTrip(req *http.Request) (*http.Response, error) {
    27  	s := ContextSegment(req.Context())
    28  	if s == nil {
    29  		return t.wrapped.RoundTrip(req)
    30  	}
    31  
    32  	sub := s.NewSubsegment(req.URL.Host)
    33  	sub.RecordRequest(req, "remote")
    34  	sub.SubmitInProgress()
    35  	defer sub.Close()
    36  
    37  	resp, err := t.wrapped.RoundTrip(req)
    38  
    39  	if err != nil {
    40  		sub.RecordError(err)
    41  	} else {
    42  		sub.RecordResponse(resp)
    43  	}
    44  
    45  	return resp, err
    46  }