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.

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) {