From 828544cb2f9158b30759b2e7d0180c1b2923265b Mon Sep 17 00:00:00 2001
From: Josh Rickmar <jrick@conformal.com>
Date: Thu, 21 Nov 2013 13:45:14 -0500
Subject: [PATCH] Fix some races found by -race.

---
 disksync.go | 15 +++++++++---
 sockets.go  | 68 +++++++++++++++++++++++++----------------------------
 2 files changed, 44 insertions(+), 39 deletions(-)

diff --git a/disksync.go b/disksync.go
index fe23021..91aef36 100644
--- a/disksync.go
+++ b/disksync.go
@@ -84,7 +84,10 @@ func (w *Account) writeDirtyToDisk() error {
 	// with block N-1.
 
 	// UTXOs
-	if w.UtxoStore.dirty {
+	w.UtxoStore.RLock()
+	dirty := w.TxStore.dirty
+	w.UtxoStore.RUnlock()
+	if dirty {
 		w.UtxoStore.Lock()
 		defer w.UtxoStore.Unlock()
 		tmpfilepath := utxofilepath + "-" + timeStr
@@ -107,7 +110,10 @@ func (w *Account) writeDirtyToDisk() error {
 	}
 
 	// Transactions
-	if w.TxStore.dirty {
+	w.TxStore.RLock()
+	dirty = w.TxStore.dirty
+	w.TxStore.RUnlock()
+	if dirty {
 		w.TxStore.Lock()
 		defer w.TxStore.Unlock()
 		tmpfilepath := txfilepath + "-" + timeStr
@@ -130,7 +136,10 @@ func (w *Account) writeDirtyToDisk() error {
 	}
 
 	// Wallet
-	if w.dirty {
+	w.mtx.RLock()
+	dirty = w.dirty
+	w.mtx.RUnlock()
+	if dirty {
 		w.mtx.Lock()
 		defer w.mtx.Unlock()
 		tmpfilepath := wfilepath + "-" + timeStr
diff --git a/sockets.go b/sockets.go
index 9cdc366..377ebb8 100644
--- a/sockets.go
+++ b/sockets.go
@@ -272,25 +272,20 @@ func frontendSendRecv(ws *websocket.Conn) {
 // BtcdHandler listens for replies and notifications from btcd over a
 // websocket and sends messages that btcwallet does not understand to
 // btcd.  Unlike FrontendHandler, exactly one BtcdHandler goroutine runs.
-func BtcdHandler(ws *websocket.Conn) {
-	// Notification channel to return from listener goroutine when
-	// btcd disconnects.
-	disconnected := make(chan int)
-	defer func() {
-		close(disconnected)
-	}()
-
+func BtcdHandler(ws *websocket.Conn, done chan struct{}) {
 	// Listen for replies/notifications from btcd, and decide how to handle them.
 	replies := make(chan []byte)
 	go func() {
 		defer close(replies)
 		for {
 			select {
-			case <-disconnected:
+			case <-done:
 				return
+
 			default:
 				var m []byte
 				if err := websocket.Message.Receive(ws, &m); err != nil {
+					close(done)
 					return
 				}
 				replies <- m
@@ -298,23 +293,31 @@ func BtcdHandler(ws *websocket.Conn) {
 		}
 	}()
 
-	for {
-		select {
-		case rply, ok := <-replies:
-			if !ok {
-				// btcd disconnected
-				return
-			}
-			// Handle message here.
-			go ProcessBtcdNotificationReply(rply)
-		case r := <-btcdMsgs:
-			if err := websocket.Message.Send(ws, r); err != nil {
-				// btcd disconnected.
-				log.Errorf("Unable to send message to btcd: %v", err)
+	go func() {
+		for {
+			select {
+			case <-done:
 				return
+
+			case rply, ok := <-replies:
+				if !ok {
+					// btcd disconnected
+					close(done)
+					return
+				}
+				// Handle message here.
+				go ProcessBtcdNotificationReply(rply)
+
+			case r := <-btcdMsgs:
+				if err := websocket.Message.Send(ws, r); err != nil {
+					// btcd disconnected.
+					log.Errorf("Unable to send message to btcd: %v", err)
+					close(done)
+					return
+				}
 			}
 		}
-	}
+	}()
 }
 
 type notificationHandler func(btcws.Notification)
@@ -462,11 +465,6 @@ func NtfnBlockConnected(n btcws.Notification) {
 	// TODO(jrick): send frontend tx notifications once that's
 	// implemented.
 	for _, a := range accounts.m {
-		// Mark wallet as being synced with the new blockstamp.
-		a.mtx.Lock()
-		a.Wallet.SetSyncedWith(bs)
-		a.mtx.Unlock()
-
 		// The UTXO store will be dirty if it was modified
 		// from a tx notification.
 		if a.UtxoStore.dirty {
@@ -490,7 +488,10 @@ func NtfnBlockConnected(n btcws.Notification) {
 		//
 		// Instead, the wallet is queued to be written to disk at the
 		// next scheduled disk sync.
+		a.mtx.Lock()
+		a.Wallet.SetSyncedWith(bs)
 		a.dirty = true
+		a.mtx.Unlock()
 		dirtyAccounts.Lock()
 		dirtyAccounts.m[a] = true
 		dirtyAccounts.Unlock()
@@ -636,15 +637,10 @@ func BtcdConnect(certificates []byte, reply chan error) {
 	}
 	replyHandlers.Unlock()
 
-	handlerClosed := make(chan int)
-	go func() {
-		BtcdHandler(btcdws)
-		close(handlerClosed)
-	}()
-
+	done := make(chan struct{})
+	BtcdHandler(btcdws, done)
 	BtcdHandshake(btcdws)
-
-	<-handlerClosed
+	<-done
 	reply <- ErrConnLost
 }