diff --git a/broker.go b/broker.go index 3c793c1..43e1253 100644 --- a/broker.go +++ b/broker.go @@ -66,6 +66,10 @@ func (br *broker) run() { for { select { case c := <-br.register: + if c.Closed() { + break + } + devid := c.DeviceID() if c.IsDevice() { diff --git a/client/client.go b/client/client.go index c4bfe46..4809ce4 100644 --- a/client/client.go +++ b/client/client.go @@ -11,4 +11,6 @@ type Client interface { IsDevice() bool Close() + + Closed() bool } diff --git a/device.go b/device.go index 7c9e6de..0065247 100644 --- a/device.go +++ b/device.go @@ -131,10 +131,15 @@ func (dev *device) WriteMsg(typ int, data []byte) { dev.send <- append(b, data...) } +func (dev *device) Closed() bool { + return atomic.LoadUint32(&dev.closed) == 1 +} + func (dev *device) Close() { - if atomic.LoadUint32(&dev.closed) == 1 { + if dev.Closed() { return } + atomic.StoreUint32(&dev.closed, 1) log.Debug().Msgf("Device '%s' disconnected", dev.conn.RemoteAddr()) diff --git a/user.go b/user.go index ab38969..d7ff913 100644 --- a/user.go +++ b/user.go @@ -55,10 +55,15 @@ func (u *user) WriteMsg(typ int, data []byte) { } } +func (u *user) Closed() bool { + return atomic.LoadUint32(&u.closed) == 1 +} + func (u *user) Close() { - if atomic.LoadUint32(&u.closed) == 1 { + if u.Closed() { return } + atomic.StoreUint32(&u.closed, 1) u.conn.Close()