InstrumentedRoundTripper handles cancelable transports

This change allows instrumenting a request made with an http.Client that
has a Timeout set.

Previously, an error would occur when using a Timeout with an
instrumented round tripper as the instrumented round tripper didn't
implement CancelRequest even when the underlying transport did. This
causes a runtime error when a request is made with the client.

Signed-off-by: Atul Kshirsagar <atul.kshirsagar@gmail.com>

dropsonde_test.go

@@ -23,7 +23,7 @@ var _ = Describe("Autowire", func() {
 	Describe("Initialize", func() {
 		It("resets the HTTP default transport to be instrumented", func() {
 			dropsonde.InitializeWithEmitter(&dropsonde.NullEventEmitter{})
-			Expect(reflect.TypeOf(http.DefaultTransport).Elem().Name()).To(Equal("instrumentedRoundTripper"))
+			Expect(reflect.TypeOf(http.DefaultTransport).Elem().Name()).To(Equal("instrumentedCancelableRoundTripper"))
 		})
 	})
 

instrumented_round_tripper/instrumented_round_tripper.go

@@ -1,24 +1,45 @@
 package instrumented_round_tripper
 
 import (
+	"log"
+	"net/http"
+
 	"github.com/cloudfoundry/dropsonde/emitter"
 	"github.com/cloudfoundry/dropsonde/events"
 	"github.com/cloudfoundry/dropsonde/factories"
 	uuid "github.com/nu7hatch/gouuid"
-	"log"
-	"net/http"
 )
 
 type instrumentedRoundTripper struct {
 	roundTripper http.RoundTripper
 	emitter      emitter.EventEmitter
 }
 
+type instrumentedCancelableRoundTripper struct {
+	instrumentedRoundTripper *instrumentedRoundTripper
+	cancelableRoundTripper   cancelableRoundTripper
+}
+
+type cancelableRoundTripper interface {
+	CancelRequest(*http.Request)
+	RoundTrip(*http.Request) (*http.Response, error)
+}
+
 /*
 Helper for creating an InstrumentedRoundTripper which will delegate to the given RoundTripper
 */
 func InstrumentedRoundTripper(roundTripper http.RoundTripper, emitter emitter.EventEmitter) http.RoundTripper {
-	return &instrumentedRoundTripper{roundTripper, emitter}
+	irt := &instrumentedRoundTripper{roundTripper, emitter}
+
+	tr, ok := roundTripper.(cancelableRoundTripper)
+	if ok {
+		return &instrumentedCancelableRoundTripper{
+			instrumentedRoundTripper: irt,
+			cancelableRoundTripper:   tr,
+		}
+	}
+
+	return irt
 }
 
 /*
@@ -65,4 +86,12 @@ func (irt *instrumentedRoundTripper) RoundTrip(req *http.Request) (*http.Respons
 	return resp, roundTripErr
 }
 
+func (icrt *instrumentedCancelableRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) {
+	return icrt.instrumentedRoundTripper.RoundTrip(req)
+}
+
+func (icrt *instrumentedCancelableRoundTripper) CancelRequest(req *http.Request) {
+	icrt.cancelableRoundTripper.CancelRequest(req)
+}
+
 var GenerateUuid = uuid.NewV4

instrumented_round_tripper/instrumented_round_tripper_test.go

@@ -2,14 +2,16 @@ package instrumented_round_tripper_test
 
 import (
 	"errors"
+	"net/http"
+	"reflect"
+
 	"github.com/cloudfoundry/dropsonde/emitter/fake"
 	"github.com/cloudfoundry/dropsonde/events"
 	"github.com/cloudfoundry/dropsonde/factories"
 	"github.com/cloudfoundry/dropsonde/instrumented_round_tripper"
 	uuid "github.com/nu7hatch/gouuid"
 	. "github.com/onsi/ginkgo"
 	. "github.com/onsi/gomega"
-	"net/http"
 )
 
 type FakeRoundTripper struct {
@@ -20,6 +22,16 @@ func (frt *FakeRoundTripper) RoundTrip(req *http.Request) (*http.Response, error
 	return &http.Response{StatusCode: 123, ContentLength: 1234}, frt.FakeError
 }
 
+type FakeCancelableRoundTripper struct {
+	FakeError error
+}
+
+func (frt *FakeCancelableRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) {
+	return &http.Response{StatusCode: 123, ContentLength: 1234}, frt.FakeError
+}
+
+func (frt *FakeCancelableRoundTripper) CancelRequest(req *http.Request) {}
+
 var _ = Describe("InstrumentedRoundTripper", func() {
 	var fakeRoundTripper *FakeRoundTripper
 	var rt http.RoundTripper
@@ -39,7 +51,27 @@ var _ = Describe("InstrumentedRoundTripper", func() {
 		Expect(err).ToNot(HaveOccurred())
 		req.RemoteAddr = "127.0.0.1"
 		req.Header.Set("User-Agent", "our-testing-client")
+	})
+
+	Context("when the round tripper is a cancelable round tripper", func() {
+		BeforeEach(func() {
+			rt = instrumented_round_tripper.InstrumentedRoundTripper(new(FakeCancelableRoundTripper), fakeEmitter)
+		})
 
+		It("returns an instrumentedCancelableRoundTripper", func() {
+			Expect(reflect.TypeOf(rt).Elem().Name()).To(Equal("instrumentedCancelableRoundTripper"))
+		})
+	})
+
+	Context("when the round tripper is not a cancelable round tripper", func() {
+		BeforeEach(func() {
+			fakeRoundTripper = new(FakeRoundTripper)
+			rt = instrumented_round_tripper.InstrumentedRoundTripper(fakeRoundTripper, fakeEmitter)
+		})
+
+		It("returns an instrumentedRoundTripper", func() {
+			Expect(reflect.TypeOf(rt).Elem().Name()).To(Equal("instrumentedRoundTripper"))
+		})
 	})
 
 	Describe("request ID", func() {