module/apmhttp: end client span on body closure

Don't end HTTP client spans as soon as the response
headers have been received, but rather when the
response body has been closed or EOF received.

CHANGELOG.md

@@ -27,6 +27,7 @@
  - Report all HTTP request/response headers (#280)
  - Drop Context.SetCustom (#284)
  - Reuse memory for tags (#286)
+ - End HTTP client spans on response body closure (#289)
 
 ## [v0.5.2](https://github.com/elastic/apm-agent-go/releases/tag/v0.5.2)
 

module/apmhttp/client.go

@@ -1,7 +1,10 @@
 package apmhttp
 
 import (
+	"io"
 	"net/http"
+	"sync/atomic"
+	"unsafe"
 
 	"go.elastic.co/apm"
 )
@@ -80,15 +83,53 @@ func (r *roundTripper) RoundTrip(req *http.Request) (*http.Response, error) {
 	spanType := "ext.http"
 	span := tx.StartSpan(name, spanType, apm.SpanFromContext(ctx))
 	span.Context.SetHTTPRequest(req)
-	defer span.End()
 	if !span.Dropped() {
 		traceContext = span.TraceContext()
+		ctx = apm.ContextWithSpan(ctx, span)
+		req = RequestWithContext(ctx, req)
+	} else {
+		span.End()
+		span = nil
 	}
 
 	req.Header.Set(TraceparentHeader, FormatTraceparentHeader(traceContext))
-	ctx = apm.ContextWithSpan(ctx, span)
-	req = RequestWithContext(ctx, req)
-	return r.r.RoundTrip(req)
+	resp, err := r.r.RoundTrip(req)
+	if span != nil {
+		if err != nil {
+			span.End()
+		} else {
+			resp.Body = &responseBody{span: span, body: resp.Body}
+		}
+	}
+	return resp, err
+}
+
+type responseBody struct {
+	span *apm.Span
+	body io.ReadCloser
+}
+
+// Close closes the response body, and ends the span if it hasn't already been ended.
+func (b *responseBody) Close() error {
+	b.endSpan()
+	return b.body.Close()
+}
+
+// Read reads from the response body, and ends the span when io.EOF is returend if
+// the span hasn't already been ended.
+func (b *responseBody) Read(p []byte) (n int, err error) {
+	n, err = b.body.Read(p)
+	if err == io.EOF {
+		b.endSpan()
+	}
+	return n, err
+}
+
+func (b *responseBody) endSpan() {
+	addr := (*unsafe.Pointer)(unsafe.Pointer(&b.span))
+	if old := atomic.SwapPointer(addr, nil); old != nil {
+		(*apm.Span)(old).End()
+	}
 }
 
 // ClientOption sets options for tracing client requests.

module/apmhttp/client_test.go

@@ -2,17 +2,20 @@ package apmhttp_test
 
 import (
 	"context"
+	"io"
 	"io/ioutil"
 	"net/http"
 	"net/http/httptest"
 	"net/url"
 	"testing"
+	"time"
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	"golang.org/x/net/context/ctxhttp"
 
 	"go.elastic.co/apm"
+	"go.elastic.co/apm/apmtest"
 	"go.elastic.co/apm/model"
 	"go.elastic.co/apm/module/apmhttp"
 	"go.elastic.co/apm/transport/transporttest"
@@ -71,3 +74,30 @@ func TestClient(t *testing.T) {
 	assert.Equal(t, span.ID, model.SpanID(clientTraceContext.Span))
 	assert.Equal(t, transaction.ID, span.ParentID)
 }
+
+func TestClientDuration(t *testing.T) {
+	const delay = 500 * time.Millisecond
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
+		w.Write([]byte("hello"))
+		w.(http.Flusher).Flush()
+		time.Sleep(delay)
+		w.Write([]byte("world"))
+	}))
+	defer server.Close()
+
+	_, spans, _ := apmtest.WithTransaction(func(ctx context.Context) {
+		client := apmhttp.WrapClient(http.DefaultClient)
+
+		resp, err := ctxhttp.Get(ctx, client, server.URL)
+		assert.NoError(t, err)
+		defer resp.Body.Close()
+		io.Copy(ioutil.Discard, resp.Body)
+	})
+
+	require.Len(t, spans, 1)
+	span := spans[0]
+
+	assert.Equal(t, "GET "+server.Listener.Addr().String(), span.Name)
+	assert.Equal(t, "ext.http", span.Type)
+	assert.InDelta(t, delay/time.Millisecond, span.Duration, 100)
+}