From a46e20f92ae33d2481c187fdd776e6b3e2c1939f Mon Sep 17 00:00:00 2001
From: Gyu-Ho Lee <gyuhox@gmail.com>
Date: Tue, 9 Feb 2016 09:43:06 -0800
Subject: [PATCH] etcd-tester: close gRPC connection when canceling

Currently gRPC connection just gets recreated
for every Stress call. When Stress ends or gets
canceled, gRPC connection must also be closed.

For https://github.com/coreos/etcd/issues/4464.
---
 tools/functional-tester/etcd-tester/stresser.go | 16 +++++++++++++---
 1 file changed, 13 insertions(+), 3 deletions(-)

diff --git a/tools/functional-tester/etcd-tester/stresser.go b/tools/functional-tester/etcd-tester/stresser.go
index 434cba72e..dc3c81712 100644
--- a/tools/functional-tester/etcd-tester/stresser.go
+++ b/tools/functional-tester/etcd-tester/stresser.go
@@ -57,6 +57,7 @@ type stresser struct {
 	success int
 
 	cancel func()
+	conn   *grpc.ClientConn
 }
 
 func (s *stresser) Stress() error {
@@ -64,10 +65,14 @@ func (s *stresser) Stress() error {
 	if err != nil {
 		return fmt.Errorf("%v (%s)", err, s.Endpoint)
 	}
-	kvc := pb.NewKVClient(conn)
-
 	ctx, cancel := context.WithCancel(context.Background())
+
+	s.mu.Lock()
+	s.conn = conn
 	s.cancel = cancel
+	s.mu.Unlock()
+
+	kvc := pb.NewKVClient(conn)
 
 	for i := 0; i < s.N; i++ {
 		go func(i int) {
@@ -97,7 +102,12 @@ func (s *stresser) Stress() error {
 }
 
 func (s *stresser) Cancel() {
-	s.cancel()
+	s.mu.Lock()
+	cancel, conn := s.cancel, s.conn
+	s.mu.Unlock()
+	cancel()
+	// TODO: wait for all routines to exit by adding a waitGroup before calling conn.Close()
+	conn.Close()
 }
 
 func (s *stresser) Report() (success int, failure int) {