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.
This commit is contained in:
magodo
2019-05-27 21:17:57 +08:00
parent e035664a8c
commit ebc479ef2c
4 changed files with 32 additions and 15 deletions

View File

@@ -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)
}
}