use timers instead of sleeps when probing

server.go

@@ -189,14 +189,10 @@ func (s *Server) start() {
 		s.refCount.Add(1)
 		go s.recv6(s.ipv6conn)
 	}
+	s.refCount.Add(1)
 	go s.probe()
 }
 
-// Shutdown closes all udp connections and unregisters the service
-func (s *Server) Shutdown() {
-	s.shutdown()
-}
-
 // SetText updates and announces the TXT records
 func (s *Server) SetText(text []string) {
 	s.service.Text = text
@@ -208,8 +204,8 @@ 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() error {
 	s.shutdownLock.Lock()
 	defer s.shutdownLock.Unlock()
 	if s.isShutdown {
@@ -526,6 +522,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,12 +553,23 @@ func (s *Server) probe() {
 
 	// Wait for a random duration uniformly distributed between 0 and 250 ms
 	// before sending the first probe packet.
-	time.Sleep(time.Duration(rand.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
@@ -569,7 +578,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)
@@ -583,7 +592,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
 	}
 }

service_test.go

@@ -29,7 +29,24 @@ func startMDNS(ctx context.Context, port int, name, service, domain string) {
 	<-ctx.Done()
 
 	log.Printf("Shutting down.")
+}
 
+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) {