pkg/server: serve drpc listeners along with grpc

The fundamental problem is that both drpc and grpc servers
want to close the listener and they both want to ignore the
error from Accept after the listener is closed. There's no
way to do this in a race free way. Fortunately, the mux
hands out listeners that can be independently closed. That
means they can both do their own shutdown logic where they
ignore the error, and then after they're closed, the code
orchestrating the servers can close the listeners.

The final weird bit is that the server's Close method is
required to wait until the Run method has exited (or at
least enough for the listeners to definitely be closed)
because tests depend on that behavior, so we have to add
some channels/mutexes/onces to ensure that Run has exited
and that a new call can't start after Close is called.

Change-Id: I7c4ef293f7963f83138815f51824fd5b8d09ce15

go.mod

@@ -125,5 +125,5 @@ require (
 	gopkg.in/olivere/elastic.v5 v5.0.76 // indirect
 	gopkg.in/spacemonkeygo/monkit.v2 v2.0.0-20190612171030-cf5a9e6f8fd2
 	gopkg.in/yaml.v2 v2.2.2
-	storj.io/drpc v0.0.1
+	storj.io/drpc v0.0.2
 )

go.sum

@@ -527,5 +527,5 @@ gopkg.in/yaml.v2 v2.2.2 h1:ZCJp+EgiOT7lHqUV2J862kp8Qj64Jo6az82+3Td9dZw=
 gopkg.in/yaml.v2 v2.2.2/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI=
 honnef.co/go/tools v0.0.0-20190102054323-c2f93a96b099/go.mod h1:rf3lG4BRIbNafJWhAfAdb/ePZxsR/4RtNHQocxwk9r4=
 honnef.co/go/tools v0.0.0-20190523083050-ea95bdfd59fc/go.mod h1:rf3lG4BRIbNafJWhAfAdb/ePZxsR/4RtNHQocxwk9r4=
-storj.io/drpc v0.0.1 h1:pCun+YtrgpQV+ISFhs3yT8GDfEJi8YyCslwaOpEzFjk=
-storj.io/drpc v0.0.1/go.mod h1:/ascUDbzNAv0A3Jj7wUIKFBH2JdJ2uJIBO/b9+2yHgQ=
+storj.io/drpc v0.0.2 h1:cf3EMOpEtd9F2xLlx8Lp7+UrcQsRBC1f50JWfbg8co8=
+storj.io/drpc v0.0.2/go.mod h1:/ascUDbzNAv0A3Jj7wUIKFBH2JdJ2uJIBO/b9+2yHgQ=

pkg/server/server.go

@@ -6,12 +6,14 @@ package server
 import (
 	"context"
 	"net"
+	"sync"
 
 	"github.com/zeebo/errs"
 	"go.uber.org/zap"
 	"golang.org/x/sync/errgroup"
 	"google.golang.org/grpc"
 
+	"storj.io/drpc/drpcserver"
 	"storj.io/storj/pkg/identity"
 	"storj.io/storj/pkg/listenmux"
 	"storj.io/storj/pkg/peertls/tlsopts"
@@ -26,11 +28,13 @@ type Service interface {
 
 type public struct {
 	listener net.Listener
+	drpc     *drpcserver.Server
 	grpc     *grpc.Server
 }
 
 type private struct {
 	listener net.Listener
+	drpc     *drpcserver.Server
 	grpc     *grpc.Server
 }
 
@@ -42,6 +46,11 @@ type Server struct {
 	private  private
 	next     []Service
 	identity *identity.FullIdentity
+
+	mu   sync.Mutex
+	wg   sync.WaitGroup
+	once sync.Once
+	done chan struct{}
 }
 
 // New creates a Server out of an Identity, a net.Listener,
@@ -51,6 +60,7 @@ func New(log *zap.Logger, opts *tlsopts.Options, publicAddr, privateAddr string,
 		log:      log,
 		next:     services,
 		identity: opts.Ident,
+		done:     make(chan struct{}),
 	}
 
 	unaryInterceptor := server.logOnErrorUnaryInterceptor
@@ -64,6 +74,7 @@ func New(log *zap.Logger, opts *tlsopts.Options, publicAddr, privateAddr string,
 	}
 	server.public = public{
 		listener: publicListener,
+		drpc:     drpcserver.New(),
 		grpc: grpc.NewServer(
 			grpc.StreamInterceptor(server.logOnErrorStreamInterceptor),
 			grpc.UnaryInterceptor(unaryInterceptor),
@@ -77,6 +88,7 @@ func New(log *zap.Logger, opts *tlsopts.Options, publicAddr, privateAddr string,
 	}
 	server.private = private{
 		listener: privateListener,
+		drpc:     drpcserver.New(),
 		grpc:     grpc.NewServer(),
 	}
 
@@ -95,13 +107,30 @@ func (p *Server) PrivateAddr() net.Addr { return p.private.listener.Addr() }
 // GRPC returns the server's gRPC handle for registration purposes
 func (p *Server) GRPC() *grpc.Server { return p.public.grpc }
 
+// DRPC returns the server's dRPC handle for registration purposes
+func (p *Server) DRPC() *drpcserver.Server { return p.public.drpc }
+
 // PrivateGRPC returns the server's gRPC handle for registration purposes
 func (p *Server) PrivateGRPC() *grpc.Server { return p.private.grpc }
 
+// PrivateDRPC returns the server's dRPC handle for registration purposes
+func (p *Server) PrivateDRPC() *drpcserver.Server { return p.private.drpc }
+
 // Close shuts down the server
 func (p *Server) Close() error {
-	p.public.grpc.GracefulStop()
-	p.private.grpc.GracefulStop()
+	p.mu.Lock()
+	defer p.mu.Unlock()
+
+	// Close done and wait for any Runs to exit.
+	p.once.Do(func() { close(p.done) })
+	p.wg.Wait()
+
+	// Ensure the listeners are closed in case Run was never called.
+	// We ignore these errors because there's not really anything to do
+	// even if they happen, and they'll just be errors due to duplicate
+	// closes anyway.
+	_ = p.public.listener.Close()
+	_ = p.private.listener.Close()
 	return nil
 }
 
@@ -117,31 +146,87 @@ func (p *Server) Run(ctx context.Context) (err error) {
 		return next.Run(ctx, p)
 	}
 
+	// Make sure the server isn't already closed. If it is, register
+	// ourselves in the wait group so that Close can wait on it.
+	p.mu.Lock()
+	select {
+	case <-p.done:
+		p.mu.Unlock()
+		return errs.New("server closed")
+	default:
+		p.wg.Add(1)
+		defer p.wg.Done()
+	}
+	p.mu.Unlock()
+
+	// We want to launch the muxes in a different group so that they are
+	// only closed after we're sure that p.Close is called. The reason why
+	// is so that we don't get "listener closed" errors because the
+	// Run call exits and closes the listeners before the servers have had
+	// a chance to be notified that they're done running.
+	const drpcHeader = "DRPC!!!1"
+
+	publicMux := listenmux.New(p.public.listener, len(drpcHeader))
+	publicDRPCListener := publicMux.Route(drpcHeader)
+
+	privateMux := listenmux.New(p.private.listener, len(drpcHeader))
+	privateDRPCListener := privateMux.Route(drpcHeader)
+
+	// We need a new context chain because we require this context to be
+	// canceled only after all of the upcoming grpc/drpc servers have
+	// fully exited. The reason why is because Run closes listener for
+	// the mux when it exits, and we can only do that after all of the
+	// Servers are no longer accepting.
+	muxCtx, muxCancel := context.WithCancel(context.Background())
+	defer muxCancel()
+
+	var muxGroup errgroup.Group
+	muxGroup.Go(func() error {
+		return publicMux.Run(muxCtx)
+	})
+	muxGroup.Go(func() error {
+		return privateMux.Run(muxCtx)
+	})
+
+	// Now we launch all the stuff that uses the listeners.
 	ctx, cancel := context.WithCancel(ctx)
 	defer cancel()
 
-	publicMux := listenmux.New(p.public.listener, 8)
-	privateMux := listenmux.New(p.private.listener, 8)
-
 	var group errgroup.Group
 	group.Go(func() error {
-		return publicMux.Run(ctx)
-	})
-	group.Go(func() error {
-		return privateMux.Run(ctx)
-	})
-	group.Go(func() error {
-		<-ctx.Done()
-		return p.Close()
+		select {
+		case <-p.done:
+			cancel()
+		case <-ctx.Done():
+		}
+
+		p.public.grpc.GracefulStop()
+		p.private.grpc.GracefulStop()
+
+		return nil
 	})
+
 	group.Go(func() error {
 		defer cancel()
 		return p.public.grpc.Serve(publicMux.Default())
 	})
+	group.Go(func() error {
+		defer cancel()
+		return p.public.drpc.Serve(ctx, publicDRPCListener)
+	})
 	group.Go(func() error {
 		defer cancel()
 		return p.private.grpc.Serve(privateMux.Default())
 	})
+	group.Go(func() error {
+		defer cancel()
+		return p.private.drpc.Serve(ctx, privateDRPCListener)
+	})
 
-	return group.Wait()
+	// Now we wait for all the stuff using the listeners to exit.
+	err = group.Wait()
+
+	// Now we close down our listeners.
+	muxCancel()
+	return errs.Combine(err, muxGroup.Wait())
 }