From e827a2ca27b4b20f3141ed9d601a2beeeff03c1f Mon Sep 17 00:00:00 2001
From: Armon Dadgar <armon.dadgar@gmail.com>
Date: Wed, 18 Dec 2013 15:09:38 -0800
Subject: [PATCH] Adding snapshot support

---
 consul/fsm.go         |  2 +-
 consul/state_store.go | 55 +++++++++++++++++++++++++++++++++++++++++--
 2 files changed, 54 insertions(+), 3 deletions(-)

diff --git a/consul/fsm.go b/consul/fsm.go
index f93397894c..bc98726f34 100644
--- a/consul/fsm.go
+++ b/consul/fsm.go
@@ -21,7 +21,7 @@ type consulFSM struct {
 // state in a way that can be accessed concurrently with operations
 // that may modify the live state.
 type consulSnapshot struct {
-	state *StateStore
+	state *StateSnapshot
 }
 
 // NewFSM is used to construct a new FSM with a blank state
diff --git a/consul/state_store.go b/consul/state_store.go
index 4f4ba69d7c..113ee759b7 100644
--- a/consul/state_store.go
+++ b/consul/state_store.go
@@ -27,6 +27,19 @@ type StateStore struct {
 	env  *mdb.Env
 }
 
+// StateSnapshot is used to provide a point-in-time snapshot
+// It works by starting a readonly transaction against all tables.
+type StateSnapshot struct {
+	tx   *mdb.Txn
+	dbis []mdb.DBI
+}
+
+// Close is used to abort the transaction and allow for cleanup
+func (s *StateSnapshot) Close() error {
+	s.tx.Abort()
+	return nil
+}
+
 // NewStateStore is used to create a new state store
 func NewStateStore() (*StateStore, error) {
 	// Create a new temp dir
@@ -475,6 +488,44 @@ func parseServiceNodes(tx *mdb.Txn, index mdb.DBI, prefix []byte) rpc.ServiceNod
 }
 
 // Snapshot is used to create a point in time snapshot
-func (s *StateStore) Snapshot() (*StateStore, error) {
-	return s, nil
+func (s *StateStore) Snapshot() (*StateSnapshot, error) {
+	// Begin a new txn
+	tx, dbis, err := s.startTxn(true, dbNodes, dbServices, dbServiceIndex)
+	if err != nil {
+		tx.Abort()
+		return nil, err
+	}
+
+	// Return the snapshot
+	snap := &StateSnapshot{
+		tx:   tx,
+		dbis: dbis,
+	}
+	return snap, nil
+}
+
+// Nodes returns all the known nodes, the slice alternates between
+// the node name and address
+func (s *StateSnapshot) Nodes() []string {
+	cursor, err := s.tx.CursorOpen(s.dbis[0])
+	if err != nil {
+		panic(fmt.Errorf("Failed to get nodes: %v", err))
+	}
+
+	var nodes []string
+	for {
+		key, val, err := cursor.Get(nil, mdb.NEXT)
+		if err == mdb.NotFound {
+			break
+		} else if err != nil {
+			panic(fmt.Errorf("Failed to get nodes: %v", err))
+		}
+		nodes = append(nodes, string(key), string(val))
+	}
+	return nodes
+}
+
+// NodeServices is used to return all the services of a given node
+func (s *StateSnapshot) NodeServices(name string) rpc.NodeServices {
+	return filterNodeServices(s.tx, s.dbis[1], name)
 }