Wait() option now accept *sync.WaitGroup

The original signature accept a boolean, and it feel like a little
verbose, since when people pass in this option, he/she always want to
pass a `true`.

Now if input `wg` is nil, it has same effect as passing `true` in
original code. Furthermore, if user want's finer grained control during
shutdown, one can pass in a predefined `wg`, so that server will wait
against it during shutdown.

function.go

@@ -54,7 +54,7 @@ func newFunction(opts ...Option) Function {
 
 	service.Server().Init(
 		// ensure the service waits for requests to finish
-		server.Wait(true),
+		server.Wait(nil),
 		// wrap handlers and subscribers to finish execution
 		server.WrapHandler(fnHandlerWrapper(fn)),
 		server.WrapSubscriber(fnSubWrapper(fn)),

server/context.go

@@ -2,19 +2,20 @@ package server
 
 import (
 	"context"
+	"sync"
 )
 
 type serverKey struct{}
 
-func wait(ctx context.Context) bool {
+func wait(ctx context.Context) *sync.WaitGroup {
 	if ctx == nil {
-		return false
+		return nil
 	}
-	wait, ok := ctx.Value("wait").(bool)
+	wg, ok := ctx.Value("wait").(*sync.WaitGroup)
 	if !ok {
-		return false
+		return nil
 	}
-	return wait
+	return wg
 }
 
 func FromContext(ctx context.Context) (Server, bool) {

server/options.go

@@ -2,6 +2,7 @@ package server
 
 import (
 	"context"
+	"sync"
 	"time"
 
 	"github.com/micro/go-micro/broker"
@@ -198,12 +199,18 @@ func WithRouter(r Router) Option {
 }
 
 // Wait tells the server to wait for requests to finish before exiting
-func Wait(b bool) Option {
+// If `wg` is nil, server only wait for completion of rpc handler.
+// For user need finer grained control, pass a concrete `wg` here, server will
+// wait against it on stop.
+func Wait(wg *sync.WaitGroup) Option {
 	return func(o *Options) {
 		if o.Context == nil {
 			o.Context = context.Background()
 		}
-		o.Context = context.WithValue(o.Context, "wait", b)
+		if wg == nil {
+			wg = new(sync.WaitGroup)
+		}
+		o.Context = context.WithValue(o.Context, "wait", wg)
 	}
 }
 

server/rpc_server.go

@@ -31,7 +31,7 @@ type rpcServer struct {
 	// used for first registration
 	registered bool
 	// graceful exit
-	wg sync.WaitGroup
+	wg *sync.WaitGroup
 }
 
 func newRpcServer(opts ...Option) Server {
@@ -42,6 +42,7 @@ func newRpcServer(opts ...Option) Server {
 		handlers:    make(map[string]Handler),
 		subscribers: make(map[*subscriber][]broker.Subscriber),
 		exit:        make(chan chan error),
+		wg:          wait(options.Context),
 	}
 }
 
@@ -63,8 +64,10 @@ func (s *rpcServer) ServeConn(sock transport.Socket) {
 			return
 		}
 
-		// add to wait group
+		// add to wait group if "wait" is opt-in
+		if s.wg != nil {
 			s.wg.Add(1)
+		}
 
 		// we use this Timeout header to set a server deadline
 		to := msg.Header["Timeout"]
@@ -111,7 +114,9 @@ func (s *rpcServer) ServeConn(sock transport.Socket) {
 					},
 					Body: []byte(err.Error()),
 				})
+				if s.wg != nil {
 					s.wg.Done()
+				}
 				return
 			}
 		}
@@ -167,13 +172,17 @@ func (s *rpcServer) ServeConn(sock transport.Socket) {
 			if err != nil {
 				log.Logf("rpc: unable to write error response: %v", err)
 			}
+			if s.wg != nil {
 				s.wg.Done()
+			}
 			return
 		}
 
 		// done
+		if s.wg != nil {
 			s.wg.Done()
 		}
+	}
 }
 
 func (s *rpcServer) newCodec(contentType string) (codec.NewCodec, error) {
@@ -555,7 +564,7 @@ func (s *rpcServer) Start() error {
 		}
 
 		// wait for requests to finish
-		if wait(s.opts.Context) {
+		if s.wg != nil {
 			s.wg.Wait()
 		}