use WG opts and add func gracefulStop

This commit is contained in:
Денис Евстигнеев 2024-03-06 13:00:49 +03:00
parent ec82a143b2
commit bbe2425174
2 changed files with 13 additions and 10 deletions

View File

@ -184,11 +184,11 @@ func validateSubscriber(sub server.Subscriber) error {
return nil
}
func (s *tcpServer) createSubHandler(sb *tcpSubscriber, opts server.Options) broker.Handler {
func (h *tcpServer) createSubHandler(sb *tcpSubscriber, opts server.Options) broker.Handler {
return func(p broker.Event) error {
msg := p.Message()
ct := msg.Header["Content-Type"]
cf, err := s.newCodec(ct)
cf, err := h.newCodec(ct)
if err != nil {
return err
}

19
tcp.go
View File

@ -404,6 +404,7 @@ func (h *tcpServer) Start() error {
}
}
h.gracefulStop()
ch <- ts.Close()
// deregister
@ -420,19 +421,21 @@ func (h *tcpServer) Start() error {
}
func (h *tcpServer) Stop() error {
ch := make(chan error)
h.exit <- ch
return <-ch
}
func (h *tcpServer) gracefulStop() {
ctx, cancel := context.WithTimeout(context.Background(), h.opts.GracefulTimeout)
go func() {
defer cancel()
h.wg.Wait()
h.opts.Wait.Wait()
cancel()
}()
// wait timeout or finish all connects
<-ctx.Done()
ch := make(chan error)
h.exit <- ch
return <-ch
}
func (h *tcpServer) String() string {
@ -483,9 +486,9 @@ func (h *tcpServer) serve(ln net.Listener, hd Handler) {
return
}
h.wg.Add(1)
h.opts.Wait.Add(1)
go func() {
defer h.wg.Done()
defer h.opts.Wait.Done()
hd.Serve(c)
}()
}