More robust solution to close deadlock

diff --git a/device.go b/device.go
index 8b0d2a53b5f6787de3c68e036f96a48f1c6198da..c714b213776f25ff36ad370696b456ca0efa2e9a 100644 (file)
--- a/device.go
+++ b/device.go
@@ -13,6 +13,10 @@ import (
        "time"
 )
 
+const (
+       DeviceRoutineNumberPerCPU = 3
+)
+
 type Device struct {
        isUp     AtomicBool // device is (going) up
        isClosed AtomicBool // device is closed? (acting as guard)
@@ -21,6 +25,7 @@ type Device struct {
        // synchronized resources (locks acquired in order)
 
        state struct {
+               stopping sync.WaitGroup
                mutex    sync.Mutex
                changing AtomicBool
                current  bool
@@ -306,7 +311,9 @@ func NewDevice(tun TUNDevice, logger *Logger) *Device {
 
        // start workers
 
-       for i := 0; i < runtime.NumCPU(); i += 1 {
+       cpus := runtime.NumCPU()
+       device.state.stopping.Add(DeviceRoutineNumberPerCPU * cpus)
+       for i := 0; i < cpus; i += 1 {
                go device.RoutineEncryption()
                go device.RoutineDecryption()
                go device.RoutineHandshake()
@@ -360,6 +367,25 @@ func (device *Device) RemoveAllPeers() {
        device.peers.keyMap = make(map[NoisePublicKey]*Peer)
 }
 
+func (device *Device) FlushPacketQueues() {
+       for {
+               select {
+               case elem, ok := <-device.queue.decryption:
+                       if ok {
+                               elem.Drop()
+                       }
+               case elem, ok := <-device.queue.encryption:
+                       if ok {
+                               elem.Drop()
+                       }
+               case <-device.queue.handshake:
+               default:
+                       return
+               }
+       }
+
+}
+
 func (device *Device) Close() {
        if device.isClosed.Swap(true) {
                return
@@ -376,6 +402,9 @@ func (device *Device) Close() {
 
        device.signal.stop.Broadcast()
 
+       device.state.stopping.Wait()
+       device.FlushPacketQueues()
+
        device.RemoveAllPeers()
        device.rate.limiter.Close()
 

diff --git a/receive.go b/receive.go
index dc6488f9833da416eea593163d91625e58206c51..e323c6d0c6d972f976bdbfc4eadafe37be304e2f 100644 (file)
--- a/receive.go
+++ b/receive.go
@@ -238,17 +238,7 @@ func (device *Device) RoutineDecryption() {
 
        logDebug := device.log.Debug
        defer func() {
-               for {
-                       select {
-                       case elem, ok := <-device.queue.decryption:
-                               if ok {
-                                       elem.Drop()
-                               }
-                       default:
-                               goto out
-                       }
-               }
-               out:
+               device.state.stopping.Done()
                logDebug.Println("Routine: decryption worker - stopped")
        }()
        logDebug.Println("Routine: decryption worker - started")
@@ -314,14 +304,7 @@ func (device *Device) RoutineHandshake() {
        logDebug := device.log.Debug
 
        defer func() {
-               for {
-                       select {
-                       case <-device.queue.handshake:
-                       default:
-                               goto out
-                       }
-               }
-               out:
+               device.state.stopping.Done()
                logDebug.Println("Routine: handshake worker - stopped")
        }()
 

diff --git a/send.go b/send.go
index 7abe2117d1fade3388dfe74ae8699086540c5ee3..e41be838bc5362b0c9b094cb88c9a92785f9d6da 100644 (file)
--- a/send.go
+++ b/send.go
@@ -274,17 +274,7 @@ func (device *Device) RoutineEncryption() {
        logDebug := device.log.Debug
 
        defer func() {
-               for {
-                       select {
-                       case elem, ok := <-device.queue.encryption:
-                               if ok {
-                                       elem.Drop()
-                               }
-                       default:
-                               goto out
-                       }
-               }
-               out:
+               device.state.stopping.Done()
                logDebug.Println("Routine: encryption worker - stopped")
        }()