implement a clean shutdown of the probe method

server.go

@@ -1,7 +1,6 @@
 package zeroconf
 
 import (
-	"errors"
 	"fmt"
 	"log"
 	"math/rand"
@@ -75,8 +74,7 @@ func Register(instance, service, domain string, port int, text []string, ifaces
 	}
 
 	s.service = entry
-	go s.mainloop()
-	go s.probe()
+	s.start()
 
 	return s, nil
 }
@@ -132,8 +130,7 @@ func RegisterProxy(instance, service, domain string, port int, host string, ips
 	}
 
 	s.service = entry
-	go s.mainloop()
-	go s.probe()
+	s.start()
 
 	return s, nil
 }
@@ -151,7 +148,7 @@ type Server struct {
 
 	shouldShutdown chan struct{}
 	shutdownLock   sync.Mutex
-	shutdownEnd    sync.WaitGroup
+	refCount       sync.WaitGroup
 	isShutdown     bool
 	ttl            uint32
 }
@@ -182,19 +179,17 @@ func newServer(ifaces []net.Interface) (*Server, error) {
 	return s, nil
 }
 
-// Start listeners and waits for the shutdown signal from exit channel
-func (s *Server) mainloop() {
+func (s *Server) start() {
 	if s.ipv4conn != nil {
+		s.refCount.Add(1)
 		go s.recv4(s.ipv4conn)
 	}
 	if s.ipv6conn != nil {
+		s.refCount.Add(1)
 		go s.recv6(s.ipv6conn)
 	}
-}
-
-// Shutdown closes all udp connections and unregisters the service
-func (s *Server) Shutdown() {
-	s.shutdown()
+	s.refCount.Add(1)
+	go s.probe()
 }
 
 // SetText updates and announces the TXT records
@@ -208,15 +203,17 @@ func (s *Server) TTL(ttl uint32) {
 	s.ttl = ttl
 }
 
-// Shutdown server will close currently open connections & channel
-func (s *Server) shutdown() error {
+// Shutdown closes all udp connections and unregisters the service
+func (s *Server) Shutdown() {
 	s.shutdownLock.Lock()
 	defer s.shutdownLock.Unlock()
 	if s.isShutdown {
-		return errors.New("server is already shutdown")
+		return
 	}
 
-	err := s.unregister()
+	if err := s.unregister(); err != nil {
+		log.Printf("failed to unregister: %s", err)
+	}
 
 	close(s.shouldShutdown)
 
@@ -228,20 +225,17 @@ func (s *Server) shutdown() error {
 	}
 
 	// Wait for connection and routines to be closed
-	s.shutdownEnd.Wait()
+	s.refCount.Wait()
 	s.isShutdown = true
-
-	return err
 }
 
-// recv is a long running routine to receive packets from an interface
+// recv4 is a long running routine to receive packets from an interface
 func (s *Server) recv4(c *ipv4.PacketConn) {
+	defer s.refCount.Done()
 	if c == nil {
 		return
 	}
 	buf := make([]byte, 65536)
-	s.shutdownEnd.Add(1)
-	defer s.shutdownEnd.Done()
 	for {
 		select {
 		case <-s.shouldShutdown:
@@ -260,14 +254,13 @@ func (s *Server) recv4(c *ipv4.PacketConn) {
 	}
 }
 
-// recv is a long running routine to receive packets from an interface
+// recv6 is a long running routine to receive packets from an interface
 func (s *Server) recv6(c *ipv6.PacketConn) {
+	defer s.refCount.Done()
 	if c == nil {
 		return
 	}
 	buf := make([]byte, 65536)
-	s.shutdownEnd.Add(1)
-	defer s.shutdownEnd.Done()
 	for {
 		select {
 		case <-s.shouldShutdown:
@@ -528,6 +521,8 @@ func (s *Server) serviceTypeName(resp *dns.Msg, ttl uint32) {
 // Perform probing & announcement
 //TODO: implement a proper probing & conflict resolution
 func (s *Server) probe() {
+	defer s.refCount.Done()
+
 	q := new(dns.Msg)
 	q.SetQuestion(s.service.ServiceInstanceName(), dns.TypePTR)
 	q.RecursionDesired = false
@@ -555,16 +550,25 @@ func (s *Server) probe() {
 	}
 	q.Ns = []dns.RR{srv, txt}
 
-	randomizer := rand.New(rand.NewSource(time.Now().UnixNano()))
-
 	// Wait for a random duration uniformly distributed between 0 and 250 ms
 	// before sending the first probe packet.
-	time.Sleep(time.Duration(randomizer.Intn(250)) * time.Millisecond)
+	timer := time.NewTimer(time.Duration(rand.Intn(250)) * time.Millisecond)
+	defer timer.Stop()
+	select {
+	case <-timer.C:
+	case <-s.shouldShutdown:
+		return
+	}
 	for i := 0; i < 3; i++ {
 		if err := s.multicastResponse(q, 0); err != nil {
 			log.Println("[ERR] zeroconf: failed to send probe:", err.Error())
 		}
-		time.Sleep(250 * time.Millisecond)
+		timer.Reset(250 * time.Millisecond)
+		select {
+		case <-timer.C:
+		case <-s.shouldShutdown:
+			return
+		}
 	}
 
 	// From RFC6762
@@ -573,7 +577,7 @@ func (s *Server) probe() {
 	//    packet loss, a responder MAY send up to eight unsolicited responses,
 	//    provided that the interval between unsolicited responses increases by
 	//    at least a factor of two with every response sent.
-	timeout := 1 * time.Second
+	timeout := time.Second
 	for i := 0; i < multicastRepetitions; i++ {
 		for _, intf := range s.ifaces {
 			resp := new(dns.Msg)
@@ -587,7 +591,12 @@ func (s *Server) probe() {
 				log.Println("[ERR] zeroconf: failed to send announcement:", err.Error())
 			}
 		}
-		time.Sleep(timeout)
+		timer.Reset(timeout)
+		select {
+		case <-timer.C:
+		case <-s.shouldShutdown:
+			return
+		}
 		timeout *= 2
 	}
 }
@@ -719,7 +728,7 @@ func (s *Server) unicastResponse(resp *dns.Msg, ifIndex int, from net.Addr) erro
 	}
 }
 
-// multicastResponse us used to send a multicast response packet
+// multicastResponse is used to send a multicast response packet
 func (s *Server) multicastResponse(msg *dns.Msg, ifIndex int) error {
 	buf, err := msg.Pack()
 	if err != nil {

service_test.go

@@ -25,6 +25,24 @@ func startMDNS(t *testing.T, port int, name, service, domain string) {
 	log.Printf("Published service: %s, type: %s, domain: %s", name, service, domain)
 }
 
+func TestQuickShutdown(t *testing.T) {
+	server, err := Register(mdnsName, mdnsService, mdnsDomain, mdnsPort, []string{"txtv=0", "lo=1", "la=2"}, nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	done := make(chan struct{})
+	go func() {
+		defer close(done)
+		server.Shutdown()
+	}()
+	select {
+	case <-done:
+	case <-time.After(500 * time.Millisecond):
+		t.Fatal("shutdown took longer than 500ms")
+	}
+}
+
 func TestBasic(t *testing.T) {
 	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
 	defer cancel()