From 21524473610fc8605caab0033ea2eccf10764854 Mon Sep 17 00:00:00 2001
From: Yicheng Qin <qycqycqycqycqyc@gmail.com>
Date: Thu, 5 Mar 2015 15:50:37 -0800
Subject: [PATCH] tools/functional-tester: add cleanup rpc

---
 tools/functional-tester/etcd-agent/agent.go   | 31 +++++++++++++++++--
 .../etcd-agent/client/client.go               |  6 ++++
 tools/functional-tester/etcd-agent/rpc.go     |  5 +++
 .../functional-tester/etcd-tester/cluster.go  |  9 ++++++
 tools/functional-tester/etcd-tester/tester.go |  4 ++-
 5 files changed, 52 insertions(+), 3 deletions(-)

diff --git a/tools/functional-tester/etcd-agent/agent.go b/tools/functional-tester/etcd-agent/agent.go
index 688efa262..ee4c6b495 100644
--- a/tools/functional-tester/etcd-agent/agent.go
+++ b/tools/functional-tester/etcd-agent/agent.go
@@ -15,10 +15,12 @@
 package main
 
 import (
+	"fmt"
 	"net"
 	"os"
 	"os/exec"
 	"path"
+	"time"
 )
 
 type Agent struct {
@@ -70,13 +72,27 @@ func (a *Agent) restart() error {
 	return a.cmd.Start()
 }
 
+func (a *Agent) cleanup() error {
+	a.stop()
+	a.logfile.Close()
+	if err := archiveLogAndDataDir("etcd.log", a.dataDir()); err != nil {
+		return err
+	}
+	f, err := os.Create("etcd.log")
+	a.logfile = f
+	return err
+}
+
 // terminate stops the exiting etcd process the agent started
 // and removes the data dir.
 func (a *Agent) terminate() error {
 	a.stop()
-	args := a.cmd.Args
+	return os.RemoveAll(a.dataDir())
+}
 
+func (a *Agent) dataDir() string {
 	datadir := path.Join(a.cmd.Path, "*.etcd")
+	args := a.cmd.Args
 	// only parse the simple case like "-data-dir /var/lib/etcd"
 	for i, arg := range args {
 		if arg == "-data-dir" {
@@ -84,5 +100,16 @@ func (a *Agent) terminate() error {
 			break
 		}
 	}
-	return os.RemoveAll(datadir)
+	return datadir
+}
+
+func archiveLogAndDataDir(log string, datadir string) error {
+	dir := fmt.Sprint(time.Now().Format(time.RFC3339))
+	if err := os.Mkdir(dir, 0700); err != nil {
+		return err
+	}
+	if err := os.Rename(log, path.Join(dir, log)); err != nil {
+		return err
+	}
+	return os.Rename(datadir, path.Join(dir, datadir))
 }
diff --git a/tools/functional-tester/etcd-agent/client/client.go b/tools/functional-tester/etcd-agent/client/client.go
index 88f956840..11704ec0a 100644
--- a/tools/functional-tester/etcd-agent/client/client.go
+++ b/tools/functional-tester/etcd-agent/client/client.go
@@ -24,6 +24,8 @@ type Agent interface {
 	Stop() error
 	// Restart restarts the existing etcd the agent stopped.
 	Restart() (int, error)
+	// Cleanup stops the exiting etcd the agent started and archived its data dir.
+	Cleanup() error
 	// Terminate stops the exiting etcd the agent started and removes its data dir.
 	Terminate() error
 	// Isoloate isolates the network of etcd
@@ -65,6 +67,10 @@ func (a *agent) Restart() (int, error) {
 	return pid, nil
 }
 
+func (a *agent) Cleanup() error {
+	return a.rpcClient.Call("Agent.RPCCleanup", struct{}{}, nil)
+}
+
 func (a *agent) Terminate() error {
 	return a.rpcClient.Call("Agent.RPCTerminate", struct{}{}, nil)
 }
diff --git a/tools/functional-tester/etcd-agent/rpc.go b/tools/functional-tester/etcd-agent/rpc.go
index de4c2c92a..288a9dac1 100644
--- a/tools/functional-tester/etcd-agent/rpc.go
+++ b/tools/functional-tester/etcd-agent/rpc.go
@@ -56,6 +56,11 @@ func (a *Agent) RPCRestart(args struct{}, pid *int) error {
 	return nil
 }
 
+func (a *Agent) RPCCleanup(args struct{}, reply *struct{}) error {
+	log.Printf("rpc: cleanup etcd")
+	return a.cleanup()
+}
+
 func (a *Agent) RPCTerminate(args struct{}, reply *struct{}) error {
 	log.Printf("rpc: terminate etcd")
 	return a.terminate()
diff --git a/tools/functional-tester/etcd-tester/cluster.go b/tools/functional-tester/etcd-tester/cluster.go
index dcbbe7011..25b851666 100644
--- a/tools/functional-tester/etcd-tester/cluster.go
+++ b/tools/functional-tester/etcd-tester/cluster.go
@@ -117,6 +117,15 @@ func (c *cluster) WaitHealth() error {
 	return err
 }
 
+func (c *cluster) Cleanup() error {
+	for _, a := range c.Agents {
+		if err := a.Cleanup(); err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
 func (c *cluster) Terminate() {
 	for _, a := range c.Agents {
 		a.Terminate()
diff --git a/tools/functional-tester/etcd-tester/tester.go b/tools/functional-tester/etcd-tester/tester.go
index a23cbe9e6..81ae3a0a1 100644
--- a/tools/functional-tester/etcd-tester/tester.go
+++ b/tools/functional-tester/etcd-tester/tester.go
@@ -59,6 +59,8 @@ func (tt *tester) runLoop() {
 
 func (tt *tester) cleanup(i, j int) error {
 	log.Printf("etcd-tester: [round#%d case#%d] cleaning up...", i, j)
-	tt.cluster.Terminate()
+	if err := tt.cluster.Cleanup(); err != nil {
+		return err
+	}
 	return tt.cluster.Bootstrap()
 }