diff --git a/pkg/dmsg/client.go b/pkg/dmsg/client.go index b70a967415..660ce15b28 100644 --- a/pkg/dmsg/client.go +++ b/pkg/dmsg/client.go @@ -159,15 +159,7 @@ func (c *ClientConn) Serve(ctx context.Context, accept chan<- *Transport) (err e log := c.log.WithField("remoteServer", c.remoteSrv) log.WithField("connCount", incrementServeCount()).Infoln("ServingConn") defer func() { - log.WithError(err).WithField("connCount", decrementServeCount()).Infoln("ClosingConn") - c.mx.Lock() - for _, tp := range c.tps { - if tp != nil { - go tp.Close() //nolint:errcheck - } - } - _ = c.Conn.Close() //nolint:errcheck - c.mx.Unlock() + log.WithError(err).WithField("connCount", decrementServeCount()).Infoln("ConnectionClosed") c.wg.Done() }() @@ -255,8 +247,17 @@ func (c *ClientConn) Close() error { closed = true c.log.WithField("remoteServer", c.remoteSrv).Infoln("ClosingConnection") close(c.done) + c.mx.Lock() + for _, tp := range c.tps { + if tp != nil { + go tp.Close() //nolint:errcheck + } + } + _ = c.Conn.Close() //nolint:errcheck + c.mx.Unlock() + c.wg.Wait() }) - c.wg.Wait() + if !closed { return ErrClientClosed }