Bugfix: fix closing of UDP routine (#118)

utils/utils.go

@@ -6,7 +6,7 @@ import (
 	"io"
 	"net"
 	"strconv"
-	"sync/atomic"
+	"sync"
 	"time"
 
 	reuseport "github.com/libp2p/go-reuseport"
@@ -158,37 +158,46 @@ func UDPStoppableRoutine(stopCh <-chan struct{}, name string, decodeFunc decoder
 		payload []byte
 	}
 
-	stopped := atomic.Value{}
-	stopped.Store(false)
 	udpDataCh := make(chan udpData)
+	defer close(udpDataCh)
+
+	wg := &sync.WaitGroup{}
+	wg.Add(1)
 	go func() {
+		defer wg.Done()
 		for {
 			u := udpData{}
 			u.size, u.pktAddr, _ = udpconn.ReadFromUDP(payload)
-			if stopped.Load() == false {
 			if u.size == 0 { // Ignore 0 byte packets.
 				continue
 			}
 			u.payload = make([]byte, u.size)
 			copy(u.payload, payload[0:u.size])
-				udpDataCh <- u
-			} else {
+			select {
+			case <-stopCh:
 				return
+			default:
+				udpDataCh <- u
 			}
 		}
 	}()
+	func() {
 		for {
 			select {
 			case u := <-udpDataCh:
 				process(u.size, u.payload, u.pktAddr, processor, localIP, addrUDP, name)
 			case <-stopCh:
-			stopped.Store(true)
-			udpconn.Close()
-			close(udpDataCh)
+				return
+			}
+		}
+	}()
+
+	for _ = range udpDataCh {
+		// drain
+	}
+	wg.Wait()
 	return nil
 }
-	}
-}
 
 func process(size int, payload []byte, pktAddr *net.UDPAddr, processor decoder.Processor, localIP string, addrUDP net.UDPAddr, name string) {
 	baseMessage := BaseMessage{

utils/utils_test.go

@@ -53,6 +53,7 @@ func TestCancelUDPRoutine(t *testing.T) {
 	require.Contains(t, []string{"message 1", "message 2", "message 3"}, readMessage())
 
 	dp.Shutdown()
+	time.Sleep(100 * time.Millisecond)
 
 	_ = sendMessage("no more messages should be processed")