From e36f271312c8113514f2f8f5a93cd43809b3d633 Mon Sep 17 00:00:00 2001
From: Andrew Wilkins <axw@elastic.co>
Date: Wed, 16 Dec 2020 14:28:16 +0800
Subject: [PATCH] module/apmgrpc: set span destination context

Update the apmgrpc client interceptor to record
span destination address/port and context for
service maps.
---
 module/apmgrpc/client.go      | 15 +++++++++++++++
 module/apmgrpc/client_test.go | 11 +++++++++++
 2 files changed, 26 insertions(+)

diff --git a/module/apmgrpc/client.go b/module/apmgrpc/client.go
index c084a5e16..b104bdbab 100644
--- a/module/apmgrpc/client.go
+++ b/module/apmgrpc/client.go
@@ -20,10 +20,13 @@
 package apmgrpc // import "go.elastic.co/apm/module/apmgrpc"
 
 import (
+	"net"
+
 	"golang.org/x/net/context"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/metadata"
+	"google.golang.org/grpc/peer"
 
 	"go.elastic.co/apm"
 	"go.elastic.co/apm/module/apmhttp"
@@ -48,13 +51,25 @@ func NewUnaryClientInterceptor(o ...ClientOption) grpc.UnaryClientInterceptor {
 		invoker grpc.UnaryInvoker,
 		opts ...grpc.CallOption,
 	) error {
+		var peer peer.Peer // maybe set after call if span != nil
 		span, ctx := startSpan(ctx, method)
 		if span != nil {
 			defer span.End()
+			opts = append(opts, grpc.Peer(&peer))
 		}
 		err := invoker(ctx, method, req, resp, cc, opts...)
 		if span != nil {
 			setSpanOutcome(span, err)
+			if peer.Addr != nil {
+				if tcpAddr, ok := peer.Addr.(*net.TCPAddr); ok {
+					span.Context.SetDestinationAddress(tcpAddr.IP.String(), tcpAddr.Port)
+				}
+				addrString := peer.Addr.String()
+				span.Context.SetDestinationService(apm.DestinationServiceSpanContext{
+					Name:     addrString,
+					Resource: addrString,
+				})
+			}
 		}
 		return err
 	}
diff --git a/module/apmgrpc/client_test.go b/module/apmgrpc/client_test.go
index cd1069281..8f6ba83d8 100644
--- a/module/apmgrpc/client_test.go
+++ b/module/apmgrpc/client_test.go
@@ -20,6 +20,7 @@
 package apmgrpc_test
 
 import (
+	"net"
 	"os"
 	"testing"
 
@@ -56,6 +57,7 @@ func testClientSpan(t *testing.T, traceparentHeaders ...string) {
 	defer serverTracer.Close()
 	s, _, addr := newServer(t, serverTracer)
 	defer s.GracefulStop()
+	tcpAddr := addr.(*net.TCPAddr)
 
 	conn, client := newClient(t, addr)
 	defer conn.Close()
@@ -84,6 +86,15 @@ func testClientSpan(t *testing.T, traceparentHeaders ...string) {
 	assert.Equal(t, "/helloworld.Greeter/SayHello", clientSpans[0].Name)
 	assert.Equal(t, "external", clientSpans[0].Type)
 	assert.Equal(t, "grpc", clientSpans[0].Subtype)
+	assert.Equal(t, &model.DestinationSpanContext{
+		Address: tcpAddr.IP.String(),
+		Port:    tcpAddr.Port,
+		Service: &model.DestinationServiceSpanContext{
+			Type:     "external",
+			Name:     tcpAddr.String(),
+			Resource: tcpAddr.String(),
+		},
+	}, clientSpans[0].Context.Destination)
 
 	serverTracer.Flush(nil)
 	serverTransactions := serverTransport.Payloads().Transactions