From 11210fd353dca04e2efcdd7d82b40ead48de729d Mon Sep 17 00:00:00 2001
From: Armon Dadgar <armon.dadgar@gmail.com>
Date: Fri, 3 Jul 2015 14:46:30 -0700
Subject: [PATCH] nomad: change to use MemDB instead of raw iradix

---
 nomad/fsm.go         |  4 ++--
 nomad/schema.go      |  8 ++++++++
 nomad/state_store.go | 30 ++++++++++++++----------------
 3 files changed, 24 insertions(+), 18 deletions(-)
 create mode 100644 nomad/schema.go

diff --git a/nomad/fsm.go b/nomad/fsm.go
index bb851faa9..52c9fbfc2 100644
--- a/nomad/fsm.go
+++ b/nomad/fsm.go
@@ -45,7 +45,7 @@ func NewFSM(logOutput io.Writer) (*nomadFSM, error) {
 
 // Close is used to cleanup resources associated with the FSM
 func (n *nomadFSM) Close() error {
-	return n.state.Close()
+	return nil
 }
 
 // State is used to return a handle to the current state
@@ -96,5 +96,5 @@ func (s *nomadSnapshot) Persist(sink raft.SnapshotSink) error {
 }
 
 func (s *nomadSnapshot) Release() {
-	s.state.Close()
+	return
 }
diff --git a/nomad/schema.go b/nomad/schema.go
new file mode 100644
index 000000000..e71dde276
--- /dev/null
+++ b/nomad/schema.go
@@ -0,0 +1,8 @@
+package nomad
+
+import "github.com/hashicorp/go-memdb"
+
+// stateStoreSchema is used to return the schema for the state store
+func stateStoreSchema() *memdb.DBSchema {
+	return nil
+}
diff --git a/nomad/state_store.go b/nomad/state_store.go
index db205fbdc..b349e4622 100644
--- a/nomad/state_store.go
+++ b/nomad/state_store.go
@@ -1,10 +1,11 @@
 package nomad
 
 import (
+	"fmt"
 	"io"
 	"log"
 
-	"github.com/hashicorp/go-immutable-radix"
+	"github.com/hashicorp/go-memdb"
 )
 
 // The StateStore is responsible for maintaining all the Consul
@@ -14,7 +15,7 @@ import (
 // to provide write availability in the face of reads.
 type StateStore struct {
 	logger *log.Logger
-	root   *iradix.Tree
+	db     *memdb.MemDB
 }
 
 // StateSnapshot is used to provide a point-in-time snapshot
@@ -22,33 +23,30 @@ type StateSnapshot struct {
 	StateStore
 }
 
-// Close is used to abort the transaction and allow for cleanup
-func (s *StateSnapshot) Close() error {
-	return nil
-}
-
 // NewStateStore is used to create a new state store
 func NewStateStore(logOutput io.Writer) (*StateStore, error) {
+	// Create the MemDB
+	db, err := memdb.NewMemDB(stateStoreSchema())
+	if err != nil {
+		return nil, fmt.Errorf("state store setup failed: %v", err)
+	}
+
+	// Create the state store
 	s := &StateStore{
 		logger: log.New(logOutput, "", log.LstdFlags),
-		root:   iradix.New(),
+		db:     db,
 	}
 	return s, nil
 }
 
-// Close is used to safely shutdown the state store
-func (s *StateStore) Close() error {
-	return nil
-}
-
 // Snapshot is used to create a point in time snapshot. Because
-// we use an immutable radix tree, we just need to preserve the
-// pointer to the root, and we are done.
+// we use MemDB, we just need to snapshot the state of the underlying
+// database.
 func (s *StateStore) Snapshot() (*StateSnapshot, error) {
 	snap := &StateSnapshot{
 		StateStore: StateStore{
 			logger: s.logger,
-			root:   s.root,
+			db:     s.db.Snapshot(),
 		},
 	}
 	return snap, nil