From f91e12fe300d6e387961e39df06e7ef6273547fc Mon Sep 17 00:00:00 2001
From: Armon Dadgar <armon.dadgar@gmail.com>
Date: Mon, 31 Mar 2014 13:41:43 -0700
Subject: [PATCH] consul: Adding FSM support for KVS operations

---
 consul/fsm.go             | 29 +++++++++++++++++++++++++++++
 consul/structs/structs.go |  1 +
 2 files changed, 30 insertions(+)

diff --git a/consul/fsm.go b/consul/fsm.go
index 4646f7e4f1..849d75d86a 100644
--- a/consul/fsm.go
+++ b/consul/fsm.go
@@ -65,6 +65,8 @@ func (c *consulFSM) Apply(log *raft.Log) interface{} {
 		return c.decodeRegister(buf[1:], log.Index)
 	case structs.DeregisterRequestType:
 		return c.applyDeregister(buf[1:], log.Index)
+	case structs.KVSRequestType:
+		return c.applyKVSOperation(buf[1:], log.Index)
 	default:
 		panic(fmt.Errorf("failed to apply request: %#v", buf))
 	}
@@ -131,6 +133,32 @@ func (c *consulFSM) applyDeregister(buf []byte, index uint64) interface{} {
 	return nil
 }
 
+func (c *consulFSM) applyKVSOperation(buf []byte, index uint64) interface{} {
+	var req structs.KVSRequest
+	if err := structs.Decode(buf, &req); err != nil {
+		panic(fmt.Errorf("failed to decode request: %v", err))
+	}
+	switch req.Op {
+	case structs.KVSSet:
+		return c.state.KVSSet(index, &req.DirEnt)
+	case structs.KVSDelete:
+		return c.state.KVSDelete(index, req.DirEnt.Key)
+	case structs.KVSDeleteTree:
+		return c.state.KVSDeleteTree(index, req.DirEnt.Key)
+	case structs.KVSCAS:
+		act, err := c.state.KVSCheckAndSet(index, &req.DirEnt)
+		if err != nil {
+			return err
+		} else {
+			return act
+		}
+	default:
+		c.logger.Printf("[WARN] consul.fsm: Invalid KVS operation '%s'", req.Op)
+		return fmt.Errorf("Invalid KVS operation '%s'", req.Op)
+	}
+	return nil
+}
+
 func (c *consulFSM) Snapshot() (raft.FSMSnapshot, error) {
 	defer func(start time.Time) {
 		c.logger.Printf("[INFO] consul.fsm: snapshot created in %v", time.Now().Sub(start))
@@ -152,6 +180,7 @@ func (c *consulFSM) Restore(old io.ReadCloser) error {
 	if err != nil {
 		return err
 	}
+	c.state.Close()
 	c.state = state
 
 	// Create a decoder
diff --git a/consul/structs/structs.go b/consul/structs/structs.go
index b4c12f5abe..5d898be9a9 100644
--- a/consul/structs/structs.go
+++ b/consul/structs/structs.go
@@ -18,6 +18,7 @@ type MessageType uint8
 const (
 	RegisterRequestType MessageType = iota
 	DeregisterRequestType
+	KVSRequestType
 )
 
 const (