Merge pull request #488 from magodo/wait_accept_custom_wg
`Wait()` option now accept *sync.WaitGroup
This commit is contained in:
commit
58a70562d8
@ -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)),
|
||||
|
@ -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) {
|
||||
|
@ -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)
|
||||
}
|
||||
}
|
||||
|
||||
|
@ -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()
|
||||
}
|
||||
|
||||
|
Loading…
Reference in New Issue
Block a user