From 92716ddb4bea36a4c24dd18a9f618e8f5c790dd6 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?P=C3=A9ter=20Szil=C3=A1gyi?= <peterke@gmail.com>
Date: Thu, 1 Feb 2018 20:37:03 +0200
Subject: [PATCH] eth/downloader: termiante state sync cleanly

---
 eth/downloader/downloader.go | 19 ++++++++++++++-----
 eth/downloader/statesync.go  |  4 ++++
 2 files changed, 18 insertions(+), 5 deletions(-)

diff --git a/eth/downloader/downloader.go b/eth/downloader/downloader.go
index 84eaa99709..82887dea8a 100644
--- a/eth/downloader/downloader.go
+++ b/eth/downloader/downloader.go
@@ -1349,8 +1349,8 @@ func (d *Downloader) importBlockResults(results []*fetchResult) error {
 // processFastSyncContent takes fetch results from the queue and writes them to the
 // database. It also controls the synchronisation of state nodes of the pivot block.
 func (d *Downloader) processFastSyncContent(latest *types.Header) error {
-	// Start syncing state of the reported head block.
-	// This should get us most of the state of the pivot block.
+	// Start syncing state of the reported head block. This should get us most of
+	// the state of the pivot block.
 	stateSync := d.syncState(latest.Root)
 	defer stateSync.Cancel()
 	go func() {
@@ -1372,10 +1372,19 @@ func (d *Downloader) processFastSyncContent(latest *types.Header) error {
 	)
 	for {
 		// Wait for the next batch of downloaded data to be available, and if the pivot
-		// block became stale, move the goalpost.
+		// block became stale, move the goalpost
 		results := d.queue.Results(oldPivot == nil) // Block if we're not monitoring pivot staleness
-		if len(results) == 0 && oldPivot == nil {
-			return stateSync.Cancel()
+		if len(results) == 0 {
+			// If pivot sync is done, stop
+			if oldPivot == nil {
+				return stateSync.Cancel()
+			}
+			// If sync failed, stop
+			select {
+			case <-d.cancelCh:
+				return stateSync.Cancel()
+			default:
+			}
 		}
 		if d.chainInsertHook != nil {
 			d.chainInsertHook(results)
diff --git a/eth/downloader/statesync.go b/eth/downloader/statesync.go
index 7a1a874912..9cc65a208c 100644
--- a/eth/downloader/statesync.go
+++ b/eth/downloader/statesync.go
@@ -293,6 +293,9 @@ func (s *stateSync) loop() error {
 		case <-s.cancel:
 			return errCancelStateFetch
 
+		case <-s.d.cancelCh:
+			return errCancelStateFetch
+
 		case req := <-s.deliver:
 			// Response, disconnect or timeout triggered, drop the peer if stalling
 			log.Trace("Received node data response", "peer", req.peer.id, "count", len(req.response), "dropped", req.dropped, "timeout", !req.dropped && req.timedOut())
@@ -347,6 +350,7 @@ func (s *stateSync) assignTasks() {
 			case s.d.trackStateReq <- req:
 				req.peer.FetchNodeData(req.items)
 			case <-s.cancel:
+			case <-s.d.cancelCh:
 			}
 		}
 	}