Tracing for gRPC Server (#619)

* Implements tracing in the native gRPC server

* Undo some unnecessary changes

* Properly revert trace/setup.go this time

* Some very very basic tests

* Remove warning for non-Trace middleware

6 files changed, 72 insertions, 11 deletions

diff --git a/core/dnsserver/server-grpc.go b/core/dnsserver/server-grpc.go
index 5eed5fc54..a91255f28 100644
--- a/core/dnsserver/server-grpc.go
+++ b/core/dnsserver/server-grpc.go
@@ -6,7 +6,9 @@ import (
 	"fmt"
 	"net"
 
+	"github.com/grpc-ecosystem/grpc-opentracing/go/otgrpc"
 	"github.com/miekg/dns"
+	opentracing "github.com/opentracing/opentracing-go"
 	"golang.org/x/net/context"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/peer"
@@ -30,10 +32,6 @@ func NewServergRPC(addr string, group []*Config) (*servergRPC, error) {
 		return nil, err
 	}
 	gs := &servergRPC{Server: s}
-	gs.grpcServer = grpc.NewServer()
-	// trace foo... TODO(miek)
-	pb.RegisterDnsServiceServer(gs.grpcServer, gs)
-
 	return gs, nil
 }
 
@@ -43,6 +41,18 @@ func (s *servergRPC) Serve(l net.Listener) error {
 	s.listenAddr = l.Addr()
 	s.m.Unlock()
 
+	if s.Tracer() != nil {
+		onlyIfParent := func(parentSpanCtx opentracing.SpanContext, method string, req, resp interface{}) bool {
+			return parentSpanCtx != nil
+		}
+		intercept := otgrpc.OpenTracingServerInterceptor(s.Tracer(), otgrpc.IncludingSpans(onlyIfParent))
+		s.grpcServer = grpc.NewServer(grpc.UnaryInterceptor(intercept))
+	} else {
+		s.grpcServer = grpc.NewServer()
+	}
+
+	pb.RegisterDnsServiceServer(s.grpcServer, s)
+
 	return s.grpcServer.Serve(l)
 }
 
diff --git a/core/dnsserver/server.go b/core/dnsserver/server.go
index 89060bc78..7366c69d1 100644
--- a/core/dnsserver/server.go
+++ b/core/dnsserver/server.go
@@ -13,9 +13,11 @@ import (
 	"github.com/coredns/coredns/middleware/metrics/vars"
 	"github.com/coredns/coredns/middleware/pkg/edns"
 	"github.com/coredns/coredns/middleware/pkg/rcode"
+	"github.com/coredns/coredns/middleware/pkg/trace"
 	"github.com/coredns/coredns/request"
 
 	"github.com/miekg/dns"
+	ot "github.com/opentracing/opentracing-go"
 	"golang.org/x/net/context"
 )
 
@@ -33,6 +35,7 @@ type Server struct {
 	zones       map[string]*Config // zones keyed by their address
 	dnsWg       sync.WaitGroup     // used to wait on outstanding connections
 	connTimeout time.Duration      // the maximum duration of a graceful shutdown
+	trace       trace.Trace        // the trace middleware for the server
 }
 
 // NewServer returns a new CoreDNS server and compiles all middleware in to it.
@@ -59,6 +62,13 @@ func NewServer(addr string, group []*Config) (*Server, error) {
 		var stack middleware.Handler
 		for i := len(site.Middleware) - 1; i >= 0; i-- {
 			stack = site.Middleware[i](stack)
+			if s.trace == nil && stack.Name() == "trace" {
+				// we have to stash away the middleware, not the
+				// Tracer object, because the Tracer won't be initialized yet
+				if t, ok := stack.(trace.Trace); ok {
+					s.trace = t
+				}
+			}
 		}
 		site.middlewareChain = stack
 	}
@@ -242,6 +252,14 @@ func (s *Server) OnStartupComplete() {
 	}
 }
 
+func (s *Server) Tracer() ot.Tracer {
+	if s.trace == nil {
+		return nil
+	}
+
+	return s.trace.Tracer()
+}
+
 // DefaultErrorFunc responds to an DNS request with an error.
 func DefaultErrorFunc(w dns.ResponseWriter, r *dns.Msg, rc int) {
 	state := request.Request{W: w, Req: r}
diff --git a/core/dnsserver/server_test.go b/core/dnsserver/server_test.go
new file mode 100644
index 000000000..0c2deeb43
--- /dev/null
+++ b/core/dnsserver/server_test.go
@@ -0,0 +1,26 @@
+package dnsserver
+
+import (
+	"testing"
+)
+
+func makeConfig(transport string) *Config {
+	return &Config{Zone: "example.com", Transport: transport, ListenHost: "127.0.0.1", Port: "53"}
+}
+
+func TestNewServer(t *testing.T) {
+	_, err := NewServer("127.0.0.1:53", []*Config{makeConfig("dns")})
+	if err != nil {
+		t.Errorf("Expected no error for NewServer, got %s.", err)
+	}
+
+	_, err = NewServergRPC("127.0.0.1:53", []*Config{makeConfig("grpc")})
+	if err != nil {
+		t.Errorf("Expected no error for NewServergRPC, got %s.", err)
+	}
+
+	_, err = NewServerTLS("127.0.0.1:53", []*Config{makeConfig("tls")})
+	if err != nil {
+		t.Errorf("Expected no error for NewServerTLS, got %s.", err)
+	}
+}
diff --git a/middleware/pkg/trace/trace.go b/middleware/pkg/trace/trace.go
new file mode 100644
index 000000000..e4858942b
--- /dev/null
+++ b/middleware/pkg/trace/trace.go
@@ -0,0 +1,12 @@
+package trace
+
+import (
+	"github.com/coredns/coredns/middleware"
+	ot "github.com/opentracing/opentracing-go"
+)
+
+// Trace holds the tracer and endpoint info
+type Trace interface {
+	middleware.Handler
+	Tracer() ot.Tracer
+}
diff --git a/middleware/proxy/grpc.go b/middleware/proxy/grpc.go
index d3ed14ccd..031869c60 100644
--- a/middleware/proxy/grpc.go
+++ b/middleware/proxy/grpc.go
@@ -5,7 +5,7 @@ import (
 	"crypto/tls"
 	"log"
 
-	"github.com/coredns/coredns/middleware/trace"
+	"github.com/coredns/coredns/middleware/pkg/trace"
 	"github.com/coredns/coredns/pb"
 	"github.com/coredns/coredns/request"
 
diff --git a/middleware/trace/trace.go b/middleware/trace/trace.go
index 3413fa681..78089b0e1 100644
--- a/middleware/trace/trace.go
+++ b/middleware/trace/trace.go
@@ -7,6 +7,7 @@ import (
 	"sync/atomic"
 
 	"github.com/coredns/coredns/middleware"
+	_ "github.com/coredns/coredns/middleware/pkg/trace"
 	"github.com/miekg/dns"
 	ot "github.com/opentracing/opentracing-go"
 	zipkin "github.com/openzipkin/zipkin-go-opentracing"
@@ -14,12 +15,6 @@ import (
 	"golang.org/x/net/context"
 )
 
-// Trace holds the tracer and endpoint info
-type Trace interface {
-	middleware.Handler
-	Tracer() ot.Tracer
-}
-
 type trace struct {
 	Next            middleware.Handler
 	ServiceEndpoint string