From 583aef6bbb1939859d4d45ac597a2d138db1cc82 Mon Sep 17 00:00:00 2001
From: Ben Johnson <benbjohnson@yahoo.com>
Date: Wed, 29 Jun 2016 10:18:24 -0600
Subject: [PATCH] check for early close during AAE

This commit adds a `Closing` channel to the `IndexSyncer` and the
`FragmentSyncer` that is periodically checked during execution.
If the channel is closed then an in-progress sync is immediately
stopped.
---
 fragment.go | 32 ++++++++++++++++++++++++++++++++
 index.go    | 12 ++++++++++++
 server.go   |  1 +
 3 files changed, 45 insertions(+)

diff --git a/fragment.go b/fragment.go
index 8d226fb24..4c9be2782 100644
--- a/fragment.go
+++ b/fragment.go
@@ -1170,6 +1170,18 @@ type FragmentSyncer struct {
 
 	Host    string
 	Cluster *Cluster
+
+	Closing <-chan struct{}
+}
+
+// isClosing returns true if the closing channel is closed.
+func (s *FragmentSyncer) isClosing() bool {
+	select {
+	case <-s.Closing:
+		return true
+	default:
+		return false
+	}
 }
 
 // SyncFragment compares checksums for the local and remote fragments and
@@ -1197,6 +1209,11 @@ func (s *FragmentSyncer) SyncFragment() error {
 			return err
 		}
 		blockSets = append(blockSets, blocks)
+
+		// Verify sync is not prematurely closing.
+		if s.isClosing() {
+			return nil
+		}
 	}
 
 	// Iterate over all blocks and find differences.
@@ -1257,6 +1274,11 @@ func (s *FragmentSyncer) syncBlock(id int) error {
 			continue
 		}
 
+		// Verify sync is not prematurely closing.
+		if s.isClosing() {
+			return nil
+		}
+
 		client, err := NewClient(node.Host)
 		if err != nil {
 			return err
@@ -1274,6 +1296,11 @@ func (s *FragmentSyncer) syncBlock(id int) error {
 		})
 	}
 
+	// Verify sync is not prematurely closing.
+	if s.isClosing() {
+		return nil
+	}
+
 	// Merge blocks together.
 	sets, clears, err := f.MergeBlock(id, pairSets)
 	if err != nil {
@@ -1298,6 +1325,11 @@ func (s *FragmentSyncer) syncBlock(id int) error {
 			fmt.Fprintf(&buf, "ClearBit(frame=%q, id=%d, profileID=%d)\n", f.Frame(), clear.BitmapIDs[j], (f.Slice()*SliceWidth)+clear.ProfileIDs[j])
 		}
 
+		// Verify sync is not prematurely closing.
+		if s.isClosing() {
+			return nil
+		}
+
 		// Execute query.
 		_, err := clients[i].ExecuteQuery(f.DB(), buf.String(), false)
 		if err != nil {
diff --git a/index.go b/index.go
index 4ac614c58..0fe535768 100644
--- a/index.go
+++ b/index.go
@@ -203,6 +203,9 @@ type IndexSyncer struct {
 
 	Host    string
 	Cluster *Cluster
+
+	// Signals that the sync should stop.
+	Closing <-chan struct{}
 }
 
 // SyncIndex compares the index on host with the local index and resolves differences.
@@ -218,6 +221,13 @@ func (s *IndexSyncer) SyncIndex() error {
 					continue
 				}
 
+				// Verify syncer has not closed.
+				select {
+				case <-s.Closing:
+					return nil
+				default:
+				}
+
 				// Sync fragment if own it.
 				if err := s.syncFragment(di.Name, fi.Name, slice); err != nil {
 					return fmt.Errorf("sync error: db=%s, frame=%s, slice=%d, err=%s", di.Name, fi.Name, slice, err)
@@ -242,9 +252,11 @@ func (s *IndexSyncer) syncFragment(db, frame string, slice uint64) error {
 		Fragment: f,
 		Host:     s.Host,
 		Cluster:  s.Cluster,
+		Closing:  s.Closing,
 	}
 	if err := fs.SyncFragment(); err != nil {
 		return err
 	}
+
 	return nil
 }
diff --git a/server.go b/server.go
index 92047d35c..5fceb9602 100644
--- a/server.go
+++ b/server.go
@@ -167,6 +167,7 @@ func (s *Server) monitorAntiEntropy() {
 		syncer.Index = s.Index
 		syncer.Host = s.Host
 		syncer.Cluster = s.Cluster
+		syncer.Closing = s.closing
 
 		// Sync indexes.
 		if err := syncer.SyncIndex(); err != nil {