From d8cd744f2fb983f8cde9fb449cbc394fc53aa874 Mon Sep 17 00:00:00 2001
From: Xiang Li <xiang.li@coreos.com>
Date: Mon, 19 Aug 2013 19:07:18 -0700
Subject: [PATCH] fix remove node

---
 command.go     |  2 +-
 raft_server.go |  5 +++++
 raft_stats.go  | 14 ++++++++++----
 transporter.go | 13 +++++++------
 4 files changed, 23 insertions(+), 11 deletions(-)

diff --git a/command.go b/command.go
index 5259311fc..060457d0c 100644
--- a/command.go
+++ b/command.go
@@ -195,7 +195,7 @@ func (c *RemoveCommand) Apply(raftServer *raft.Server) (interface{}, error) {
 	key := path.Join("_etcd/machines", c.Name)
 
 	_, err := etcdStore.Delete(key, raftServer.CommitIndex())
-	r.peersStats[c.Name] = nil
+	delete(r.peersStats, c.Name)
 
 	if err != nil {
 		return []byte{0}, err
diff --git a/raft_server.go b/raft_server.go
index 9009f64e0..6355ea86b 100644
--- a/raft_server.go
+++ b/raft_server.go
@@ -269,6 +269,11 @@ func joinByMachine(s *raft.Server, machine string, scheme string) error {
 	return fmt.Errorf("Unable to join: %v", err)
 }
 
+func (r *raftServer) Stats() []byte {
+	b, _ := json.Marshal(r.peersStats)
+	return b
+}
+
 // Register commands to raft server
 func registerCommands() {
 	raft.RegisterCommand(&JoinCommand{})
diff --git a/raft_stats.go b/raft_stats.go
index 2f15a7a80..7e2643a3d 100644
--- a/raft_stats.go
+++ b/raft_stats.go
@@ -1,7 +1,7 @@
 package main
 
 import (
-	"encoding/json"
+	"time"
 )
 
 type peerStats struct {
@@ -11,7 +11,13 @@ type peerStats struct {
 	SuccCnt    uint64  `json:"successCount"`
 }
 
-func (r *raftServer) Stats() []byte {
-	b, _ := json.Marshal(r.peersStats)
-	return b
+func (ps *peerStats) Fail() {
+	ps.FailCnt++
+}
+
+func (ps *peerStats) Succ(d time.Duration) {
+	total := float64(ps.SuccCnt) * ps.AvgLatency
+	ps.SuccCnt++
+	ps.Latency = float64(d) / (1000000.0)
+	ps.AvgLatency = (total + ps.Latency) / float64(ps.SuccCnt)
 }
diff --git a/transporter.go b/transporter.go
index 914902d5f..909b6ea68 100644
--- a/transporter.go
+++ b/transporter.go
@@ -52,7 +52,7 @@ func (t transporter) SendAppendEntriesRequest(server *raft.Server, peer *raft.Pe
 	u, _ := nameToRaftURL(peer.Name)
 	debugf("Send LogEntries to %s ", u)
 
-	thisPeerStats := r.peersStats[peer.Name]
+	thisPeerStats, ok := r.peersStats[peer.Name]
 
 	start := time.Now()
 
@@ -62,12 +62,13 @@ func (t transporter) SendAppendEntriesRequest(server *raft.Server, peer *raft.Pe
 
 	if err != nil {
 		debugf("Cannot send AppendEntriesRequest to %s: %s", u, err)
-		thisPeerStats.FailCnt++
+		if ok {
+			thisPeerStats.Fail()
+		}
 	} else {
-		total := float64(thisPeerStats.SuccCnt) * thisPeerStats.AvgLatency
-		thisPeerStats.SuccCnt++
-		thisPeerStats.Latency = float64(end.Sub(start)) / (1000000.0)
-		thisPeerStats.AvgLatency = (total + thisPeerStats.Latency) / float64(thisPeerStats.SuccCnt)
+		if ok {
+			thisPeerStats.Succ(end.Sub(start))
+		}
 	}
 
 	r.peersStats[peer.Name] = thisPeerStats