stubserver: support xds-enabled grpc server

balancer/weightedroundrobin/balancer_test.go

@@ -138,7 +138,7 @@ func startServer(t *testing.T, r reportType) *testServer {
 			MinReportingInterval:  10 * time.Millisecond,
 		}
 		internal.ORCAAllowAnyMinReportingInterval.(func(so *orca.ServiceOptions))(&oso)
-		sopts = append(sopts, stubserver.RegisterServiceServerOption(func(s *grpc.Server) {
+		sopts = append(sopts, stubserver.RegisterServiceServerOption(func(s grpc.ServiceRegistrar) {
 			if err := orca.Register(s, oso); err != nil {
 				t.Fatalf("Failed to register orca service: %v", err)
 			}

internal/stubserver/stubserver.go

@@ -24,6 +24,7 @@
 	"context"
 	"fmt"
 	"net"
+	"net/http"
 	"testing"
 	"time"
 
@@ -39,6 +40,15 @@
 	testpb "google.golang.org/grpc/interop/grpc_testing"
 )
 
+// GRPCServer is an interface that groups methods implemented by a grpc.Server
+// or an xds.GRPCServer that are used by the StubServer.
+type GRPCServer interface {
+	grpc.ServiceRegistrar
+	Stop()
+	GracefulStop()
+	Serve(net.Listener) error
+}
+
 // StubServer is a server that is easy to customize within individual test
 // cases.
 type StubServer struct {
@@ -53,7 +63,12 @@
 	// A client connected to this service the test may use.  Created in Start().
 	Client testgrpc.TestServiceClient
 	CC     *grpc.ClientConn
-	S      *grpc.Server
+
+	// Server to serve this service from.
+	//
+	// If nil, a new grpc.Server is created, listening on the provided Network
+	// and Address fields, or listening using the provided Listener.
+	S GRPCServer
 
 	// Parameters for Listen and Dial. Defaults will be used if these are empty
 	// before Start.
@@ -100,14 +115,14 @@
 
 type registerServiceServerOption struct {
 	grpc.EmptyServerOption
-	f func(*grpc.Server)
+	f func(grpc.ServiceRegistrar)
 }
 
 // RegisterServiceServerOption returns a ServerOption that will run f() in
 // Start or StartServer with the grpc.Server created before serving.  This
 // allows other services to be registered on the test server (e.g. ORCA,
 // health, or reflection).
-func RegisterServiceServerOption(f func(*grpc.Server)) grpc.ServerOption {
+func RegisterServiceServerOption(f func(grpc.ServiceRegistrar)) grpc.ServerOption {
 	return &registerServiceServerOption{f: f}
 }
 
@@ -132,7 +147,9 @@
 	}
 	ss.Address = lis.Addr().String()
 
-	ss.S = grpc.NewServer(sopts...)
+	if ss.S == nil {
+		ss.S = grpc.NewServer(sopts...)
+	}
 	for _, so := range sopts {
 		switch x := so.(type) {
 		case *registerServiceServerOption:
@@ -154,9 +171,14 @@
 		return err
 	}
 
+	handler, ok := ss.S.(interface{ http.Handler })
+	if !ok {
+		panic(fmt.Sprintf("server of type %T does not implement http.Handler", ss.S))
+	}
+
 	go func() {
 		hs := &http2.Server{}
-		opts := &http2.ServeConnOpts{Handler: ss.S}
+		opts := &http2.ServeConnOpts{Handler: handler}
 		for {
 			conn, err := lis.Accept()
 			if err != nil {

orca/service.go

@@ -111,9 +111,7 @@ func NewService(opts ServiceOptions) (*Service, error) {
 
 // Register creates a new ORCA service implementation configured using the
 // provided options and registers the same on the provided grpc Server.
-func Register(s *grpc.Server, opts ServiceOptions) error {
-	// TODO(https://github.com/cncf/xds/issues/41): replace *grpc.Server with
-	// grpc.ServiceRegistrar when possible.
+func Register(s grpc.ServiceRegistrar, opts ServiceOptions) error {
 	service, err := NewService(opts)
 	if err != nil {
 		return err