Serve() should not return error on Stop() or GracefulStop() (#1485)

* Serve() should not return error on Stop() or GracefulStop()
* Block Serve() until Stop() or GracefulStop() returns

server.go

@@ -96,6 +96,11 @@ type Server struct {
 	cv     *sync.Cond
 	m      map[string]*service // service name -> service info
 	events trace.EventLog
+
+	quit     chan struct{}
+	done     chan struct{}
+	quitOnce sync.Once
+	doneOnce sync.Once
 }
 
 type options struct {
@@ -307,6 +312,8 @@ func NewServer(opt ...ServerOption) *Server {
 		opts:  opts,
 		conns: make(map[io.Closer]bool),
 		m:     make(map[string]*service),
+		quit:  make(chan struct{}),
+		done:  make(chan struct{}),
 	}
 	s.cv = sync.NewCond(&s.mu)
 	s.ctx, s.cancel = context.WithCancel(context.Background())
@@ -485,6 +492,14 @@ func (s *Server) Serve(lis net.Listener) error {
 			s.mu.Lock()
 			s.printf("done serving; Accept = %v", err)
 			s.mu.Unlock()
+
+			// If Stop or GracefulStop is called, block until they are done and return nil
+			select {
+			case <-s.quit:
+				<-s.done
+				return nil
+			default:
+			}
 			return err
 		}
 		tempDelay = 0
@@ -1052,6 +1067,16 @@ func (s *Server) handleStream(t transport.ServerTransport, stream *transport.Str
 // pending RPCs on the client side will get notified by connection
 // errors.
 func (s *Server) Stop() {
+	s.quitOnce.Do(func() {
+		close(s.quit)
+	})
+
+	defer func() {
+		s.doneOnce.Do(func() {
+			close(s.done)
+		})
+	}()
+
 	s.mu.Lock()
 	listeners := s.lis
 	s.lis = nil
@@ -1081,6 +1106,16 @@ func (s *Server) Stop() {
 // accepting new connections and RPCs and blocks until all the pending RPCs are
 // finished.
 func (s *Server) GracefulStop() {
+	s.quitOnce.Do(func() {
+		close(s.quit)
+	})
+
+	defer func() {
+		s.doneOnce.Do(func() {
+			close(s.done)
+		})
+	}()
+
 	s.mu.Lock()
 	defer s.mu.Unlock()
 	if s.conns == nil {

server_test.go

@@ -23,6 +23,7 @@ import (
 	"reflect"
 	"strings"
 	"testing"
+	"time"
 
 	"google.golang.org/grpc/test/leakcheck"
 )
@@ -53,6 +54,27 @@ func TestStopBeforeServe(t *testing.T) {
 	}
 }
 
+func TestGracefulStop(t *testing.T) {
+	defer leakcheck.Check(t)
+
+	lis, err := net.Listen("tcp", "localhost:0")
+	if err != nil {
+		t.Fatalf("failed to create listener: %v", err)
+	}
+
+	server := NewServer()
+	go func() {
+		// make sure Serve() is called
+		time.Sleep(time.Millisecond * 500)
+		server.GracefulStop()
+	}()
+
+	err = server.Serve(lis)
+	if err != nil {
+		t.Fatalf("Serve() returned non-nil error on GracefulStop: %v", err)
+	}
+}
+
 func TestGetServiceInfo(t *testing.T) {
 	defer leakcheck.Check(t)
 	testSd := ServiceDesc{