Merge pull request #66 from hashicorp/fix-data-races

Fix races found with go test -race

Author: mkeeler

client.go

@@ -5,7 +5,7 @@ import (
 	"log"
 	"net"
 	"strings"
-	"sync"
+	"sync/atomic"
 	"time"
 
 	"github.com/miekg/dns"
@@ -102,9 +102,8 @@ type client struct {
 	ipv4MulticastConn *net.UDPConn
 	ipv6MulticastConn *net.UDPConn
 
-	closed    bool
-	closedCh  chan struct{} // TODO(reddaly): This doesn't appear to be used.
-	closeLock sync.Mutex
+	closed   int32
+	closedCh chan struct{} // TODO(reddaly): This doesn't appear to be used.
 }
 
 // NewClient creates a new mdns Client that can be used to query
@@ -150,13 +149,10 @@ func newClient() (*client, error) {
 
 // Close is used to cleanup the client
 func (c *client) Close() error {
-	c.closeLock.Lock()
-	defer c.closeLock.Unlock()
-
-	if c.closed {
+	if !atomic.CompareAndSwapInt32(&c.closed, 0, 1) {
+		// something else already closed it
 		return nil
 	}
-	c.closed = true
 
 	log.Printf("[INFO] mdns: Closing client %v", *c)
 	close(c.closedCh)
@@ -326,8 +322,13 @@ func (c *client) recv(l *net.UDPConn, msgCh chan *dns.Msg) {
 		return
 	}
 	buf := make([]byte, 65536)
-	for !c.closed {
+	for atomic.LoadInt32(&c.closed) == 0 {
 		n, err := l.Read(buf)
+
+		if atomic.LoadInt32(&c.closed) == 1 {
+			return
+		}
+
 		if err != nil {
 			log.Printf("[ERR] mdns: Failed to read packet: %v", err)
 			continue

server.go

@@ -5,7 +5,7 @@ import (
 	"log"
 	"net"
 	"strings"
-	"sync"
+	"sync/atomic"
 
 	"github.com/miekg/dns"
 )
@@ -51,9 +51,8 @@ type Server struct {
 	ipv4List *net.UDPConn
 	ipv6List *net.UDPConn
 
-	shutdown     bool
-	shutdownCh   chan struct{}
-	shutdownLock sync.Mutex
+	shutdown   int32
+	shutdownCh chan struct{}
 }
 
 // NewServer is used to create a new mDNS server from a config
@@ -87,13 +86,11 @@ func NewServer(config *Config) (*Server, error) {
 
 // Shutdown is used to shutdown the listener
 func (s *Server) Shutdown() error {
-	s.shutdownLock.Lock()
-	defer s.shutdownLock.Unlock()
-
-	if s.shutdown {
+	if !atomic.CompareAndSwapInt32(&s.shutdown, 0, 1) {
+		// something else already closed us
 		return nil
 	}
-	s.shutdown = true
+
 	close(s.shutdownCh)
 
 	if s.ipv4List != nil {
@@ -111,8 +108,9 @@ func (s *Server) recv(c *net.UDPConn) {
 		return
 	}
 	buf := make([]byte, 65536)
-	for !s.shutdown {
+	for atomic.LoadInt32(&s.shutdown) == 0 {
 		n, from, err := c.ReadFrom(buf)
+
 		if err != nil {
 			continue
 		}

server_test.go

@@ -1,6 +1,7 @@
 package mdns
 
 import (
+	"sync/atomic"
 	"testing"
 	"time"
 )
@@ -22,7 +23,7 @@ func TestServer_Lookup(t *testing.T) {
 	defer serv.Shutdown()
 
 	entries := make(chan *ServiceEntry, 1)
-	found := false
+	var found int32 = 0
 	go func() {
 		select {
 		case e := <-entries:
@@ -35,7 +36,7 @@ func TestServer_Lookup(t *testing.T) {
 			if e.Info != "Local web server" {
 				t.Fatalf("bad: %v", e)
 			}
-			found = true
+			atomic.StoreInt32(&found, 1)
 
 		case <-time.After(80 * time.Millisecond):
 			t.Fatalf("timeout")
@@ -52,7 +53,7 @@ func TestServer_Lookup(t *testing.T) {
 	if err != nil {
 		t.Fatalf("err: %v", err)
 	}
-	if !found {
+	if atomic.LoadInt32(&found) == 0 {
 		t.Fatalf("record not found")
 	}
 }