ochttp plugin: makes Body a transparent wrapper (#1069)

plugin/ochttp/client_stats.go

@@ -68,7 +68,7 @@ func (t statsTransport) RoundTrip(req *http.Request) (*http.Response, error) {
 			track.end()
 		} else {
 			track.body = resp.Body
-			resp.Body = track
+			resp.Body = wrappedBody(track, resp.Body)
 		}
 	}
 	return resp, err

plugin/ochttp/trace.go

@@ -93,7 +93,8 @@ func (t *traceTransport) RoundTrip(req *http.Request) (*http.Response, error) {
 	// span.End() will be invoked after
 	// a read from resp.Body returns io.EOF or when
 	// resp.Body.Close() is invoked.
-	resp.Body = &bodyTracker{rc: resp.Body, span: span}
+	bt := &bodyTracker{rc: resp.Body, span: span}
+	resp.Body = wrappedBody(bt, resp.Body)
 	return resp, err
 }
 

plugin/ochttp/wrapped_body.go

@@ -0,0 +1,44 @@
+// Copyright 2019, OpenCensus Authors
+//
+// 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.
+
+package ochttp
+
+import (
+	"io"
+)
+
+// wrappedBody returns a wrapped version of the original
+// Body and only implements the same combination of additional
+// interfaces as the original.
+func wrappedBody(wrapper io.ReadCloser, body io.ReadCloser) io.ReadCloser {
+	var (
+		wr, i0 = body.(io.Writer)
+	)
+	switch {
+	case !i0:
+		return struct {
+			io.ReadCloser
+		}{wrapper}
+
+	case i0:
+		return struct {
+			io.ReadCloser
+			io.Writer
+		}{wrapper, wr}
+	default:
+		return struct {
+			io.ReadCloser
+		}{wrapper}
+	}
+}