From 189df5c5356df6d54cfccb667d1292753459f808 Mon Sep 17 00:00:00 2001
From: Josh Rickmar <jrick@conformal.com>
Date: Fri, 28 Mar 2014 11:28:59 -0500
Subject: [PATCH] Mark partially synced addresses.

This change adds a notification handler for the new rescanprogress
notification and takes advantage of the recent rescan manager and
partial syncing support to mark addresses as partially synced.  If the
network connection to btcd is lost or wallet is restarted during a
rescan, a new rescan will start at the earliest block height for any
wallet address, taking partial syncs into consideration.
---
 acctmgr.go   | 25 ++++++++++++++++++++++---
 ntfns.go     | 16 ++++++++++++++++
 rpcserver.go |  2 +-
 3 files changed, 39 insertions(+), 4 deletions(-)

diff --git a/acctmgr.go b/acctmgr.go
index 6ef5eff..4de791e 100644
--- a/acctmgr.go
+++ b/acctmgr.go
@@ -159,7 +159,22 @@ func (am *AccountManager) rescanListener() {
 			log.Infof("Started rescan at height %d for %d %s", e.StartHeight, n, noun)
 
 		case *RescanProgressMsg:
-			// TODO: mark addresses as partially synced.
+			for acct, addrs := range e.Addresses {
+				for i := range addrs {
+					err := acct.SetSyncStatus(addrs[i], wallet.PartialSync(e.Height))
+					if err != nil {
+						log.Errorf("Error marking address partially synced: %v", err)
+						continue
+					}
+				}
+				am.ds.ScheduleWalletWrite(acct)
+				err := am.ds.FlushAccount(acct)
+				if err != nil {
+					log.Errorf("Could not write rescan progress: %v", err)
+				}
+			}
+
+			log.Infof("Rescanned through block height %d", e.Height)
 
 		case *RescanFinishedMsg:
 			if e.Error != nil {
@@ -169,15 +184,19 @@ func (am *AccountManager) rescanListener() {
 
 			n := 0
 			for acct, addrs := range e.Addresses {
+				n += len(addrs)
 				for i := range addrs {
-					n++
 					err := acct.SetSyncStatus(addrs[i], wallet.FullSync{})
 					if err != nil {
 						log.Errorf("Error marking address synced: %v", err)
 						continue
 					}
 				}
-				AcctMgr.ds.FlushAccount(acct)
+				am.ds.ScheduleWalletWrite(acct)
+				err := am.ds.FlushAccount(acct)
+				if err != nil {
+					log.Errorf("Could not write rescan progress: %v", err)
+				}
 			}
 
 			noun := pickNoun(n, "address", "addresses")
diff --git a/ntfns.go b/ntfns.go
index f4c7c3a..6b8d802 100644
--- a/ntfns.go
+++ b/ntfns.go
@@ -56,6 +56,7 @@ var notificationHandlers = map[string]notificationHandler{
 	btcws.BlockDisconnectedNtfnMethod: NtfnBlockDisconnected,
 	btcws.RecvTxNtfnMethod:            NtfnRecvTx,
 	btcws.RedeemingTxNtfnMethod:       NtfnRedeemingTx,
+	btcws.RescanProgressNtfnMethod:    NtfnRescanProgress,
 }
 
 // NtfnRecvTx handles the btcws.RecvTxNtfn notification.
@@ -264,3 +265,18 @@ func NtfnRedeemingTx(n btcjson.Cmd) error {
 
 	return nil
 }
+
+// NtfnRescanProgress handles btcd rescanprogress notifications resulting
+// from a partially completed rescan.
+func NtfnRescanProgress(n btcjson.Cmd) error {
+	cn, ok := n.(*btcws.RescanProgressNtfn)
+	if !ok {
+		return fmt.Errorf("%v handler: unexpected type", n.Method())
+	}
+
+	// Notify the rescan manager of the completed partial progress for
+	// the current rescan.
+	AcctMgr.rm.MarkProgress(cn.LastProcessed)
+
+	return nil
+}
diff --git a/rpcserver.go b/rpcserver.go
index 2c0e83d..3d6ad48 100644
--- a/rpcserver.go
+++ b/rpcserver.go
@@ -218,7 +218,7 @@ func WalletRequestProcessor() {
 				// notifications.  Restart the connection
 				// to reload accounts files from their last
 				// known good state.
-				log.Warn("Reconnecting to recover from "+
+				log.Warn("Reconnecting to recover from " +
 					"out-of-order btcd notification")
 				s := CurrentServerConn()
 				if btcd, ok := s.(*BtcdRPCConn); ok {