From 934c28d49892537c126d125e78162049745e9f4c Mon Sep 17 00:00:00 2001
From: Yicheng Qin <qycqycqycqycqyc@gmail.com>
Date: Tue, 20 May 2014 13:12:12 -0700
Subject: [PATCH] fix(peer_server): set store and registry when setting raft
 server

New raft server needs new store and registry.
---
 etcd/etcd.go          | 8 ++++++++
 server/peer_server.go | 8 ++++++++
 server/server.go      | 8 ++++++++
 3 files changed, 24 insertions(+)

diff --git a/etcd/etcd.go b/etcd/etcd.go
index 3e6ce7ed0..f5da2116e 100644
--- a/etcd/etcd.go
+++ b/etcd/etcd.go
@@ -327,6 +327,14 @@ func (e *Etcd) runServer() {
 			e.StandbyServer.SyncCluster(peerURLs)
 			e.setMode(StandbyMode)
 		} else {
+			// Create etcd key-value store and registry.
+			e.Store = store.New()
+			e.Registry = server.NewRegistry(e.Store)
+			e.PeerServer.SetStore(e.Store)
+			e.PeerServer.SetRegistry(e.Registry)
+			e.Server.SetStore(e.Store)
+			e.Server.SetRegistry(e.Registry)
+
 			// Generate new peer server here.
 			// TODO(yichengq): raft server cannot be started after stopped.
 			// It should be removed when raft restart is implemented.
diff --git a/server/peer_server.go b/server/peer_server.go
index 0a317c769..59dbe3592 100644
--- a/server/peer_server.go
+++ b/server/peer_server.go
@@ -146,6 +146,14 @@ func (s *PeerServer) SetRaftServer(raftServer raft.Server, snapshot bool) {
 	}
 }
 
+func (s *PeerServer) SetRegistry(registry *Registry) {
+	s.registry = registry
+}
+
+func (s *PeerServer) SetStore(store store.Store) {
+	s.store = store
+}
+
 // Try all possible ways to find clusters to join
 // Include log data in -data-dir, -discovery and -peers
 //
diff --git a/server/server.go b/server/server.go
index 653420349..a7aebc31f 100644
--- a/server/server.go
+++ b/server/server.go
@@ -99,6 +99,14 @@ func (s *Server) Store() store.Store {
 	return s.store
 }
 
+func (s *Server) SetRegistry(registry *Registry) {
+	s.registry = registry
+}
+
+func (s *Server) SetStore(store store.Store) {
+	s.store = store
+}
+
 func (s *Server) installV1(r *mux.Router) {
 	s.handleFuncV1(r, "/v1/keys/{key:.*}", v1.GetKeyHandler).Methods("GET", "HEAD")
 	s.handleFuncV1(r, "/v1/keys/{key:.*}", v1.SetKeyHandler).Methods("POST", "PUT")