From 8371c87fd00fb816628cd8c7052a1de508136a7f Mon Sep 17 00:00:00 2001
From: James Phillips <james@hashicorp.com>
Date: Mon, 19 Oct 2015 14:07:57 -0700
Subject: [PATCH] Converts KVS snapshot over to iterator.

---
 consul/fsm.go                    | 10 +++++-----
 consul/state/state_store.go      | 13 ++++---------
 consul/state/state_store_test.go |  6 +++++-
 3 files changed, 14 insertions(+), 15 deletions(-)

diff --git a/consul/fsm.go b/consul/fsm.go
index 2dd38c7971..0982fd5c63 100644
--- a/consul/fsm.go
+++ b/consul/fsm.go
@@ -378,7 +378,7 @@ func (s *consulSnapshot) Persist(sink raft.SnapshotSink) error {
 		return err
 	}
 
-	if err := s.persistKV(sink, encoder); err != nil {
+	if err := s.persistKVs(sink, encoder); err != nil {
 		sink.Cancel()
 		return err
 	}
@@ -475,16 +475,16 @@ func (s *consulSnapshot) persistACLs(sink raft.SnapshotSink,
 	return nil
 }
 
-func (s *consulSnapshot) persistKV(sink raft.SnapshotSink,
+func (s *consulSnapshot) persistKVs(sink raft.SnapshotSink,
 	encoder *codec.Encoder) error {
-	entries, err := s.state.KVSDump()
+	iter, err := s.state.KVs()
 	if err != nil {
 		return err
 	}
 
-	for _, e := range entries {
+	for ki := iter.Next(); ki != nil; ki = iter.Next() {
 		sink.Write([]byte{byte(structs.KVSRequestType)})
-		if err := encoder.Encode(e); err != nil {
+		if err := encoder.Encode(ki.(*structs.DirEntry)); err != nil {
 			return err
 		}
 	}
diff --git a/consul/state/state_store.go b/consul/state/state_store.go
index 969e50003d..4176ab556d 100644
--- a/consul/state/state_store.go
+++ b/consul/state/state_store.go
@@ -157,17 +157,12 @@ func (s *StateSnapshot) Checks(node string) (memdb.ResultIterator, error) {
 }
 
 // KVSDump is used to pull the full list of KVS entries for use during snapshots.
-func (s *StateSnapshot) KVSDump() (structs.DirEntries, error) {
-	entries, err := s.tx.Get("kvs", "id_prefix")
+func (s *StateSnapshot) KVs() (memdb.ResultIterator, error) {
+	iter, err := s.tx.Get("kvs", "id_prefix")
 	if err != nil {
-		return nil, fmt.Errorf("failed kvs lookup: %s", err)
-	}
-
-	var dump structs.DirEntries
-	for entry := entries.Next(); entry != nil; entry = entries.Next() {
-		dump = append(dump, entry.(*structs.DirEntry))
+		return nil, err
 	}
-	return dump, nil
+	return iter, nil
 }
 
 // TombstoneDump is used to pull all the tombstones from the graveyard.
diff --git a/consul/state/state_store_test.go b/consul/state/state_store_test.go
index 9c13101dec..3137023b54 100644
--- a/consul/state/state_store_test.go
+++ b/consul/state/state_store_test.go
@@ -3065,10 +3065,14 @@ func TestStateStore_KVS_Snapshot_Restore(t *testing.T) {
 	if idx := snap.LastIndex(); idx != 7 {
 		t.Fatalf("bad index: %d", idx)
 	}
-	dump, err := snap.KVSDump()
+	iter, err := snap.KVs()
 	if err != nil {
 		t.Fatalf("err: %s", err)
 	}
+	var dump structs.DirEntries
+	for ki := iter.Next(); ki != nil; ki = iter.Next() {
+		dump = append(dump, ki.(*structs.DirEntry))
+	}
 	if !reflect.DeepEqual(dump, entries) {
 		t.Fatalf("bad: %#v", dump)
 	}