From 5085c7628d2cd711008f044be871270199c19d9c Mon Sep 17 00:00:00 2001
From: iamqizhao <toqizhao@gmail.com>
Date: Mon, 22 Feb 2016 18:02:10 -0800
Subject: [PATCH] Always close tracing when a stream goes wrong

---
 stream.go | 9 +++++++++
 1 file changed, 9 insertions(+)

diff --git a/stream.go b/stream.go
index ea685cc15..0cacdf9de 100644
--- a/stream.go
+++ b/stream.go
@@ -204,6 +204,9 @@ func (cs *clientStream) SendMsg(m interface{}) (err error) {
 		cs.mu.Unlock()
 	}
 	defer func() {
+		if err != nil {
+			cs.finish(err)
+		}
 		if err == nil || err == io.EOF {
 			return
 		}
@@ -251,6 +254,7 @@ func (cs *clientStream) RecvMsg(m interface{}) (err error) {
 		}
 		if err == io.EOF {
 			if cs.s.StatusCode() == codes.OK {
+				cs.finish(err)
 				return nil
 			}
 			return Errorf(cs.s.StatusCode(), cs.s.StatusDesc())
@@ -272,6 +276,11 @@ func (cs *clientStream) RecvMsg(m interface{}) (err error) {
 
 func (cs *clientStream) CloseSend() (err error) {
 	err = cs.t.Write(cs.s, nil, &transport.Options{Last: true})
+	defer func() {
+		if err != nil {
+			cs.finish(err)
+		}
+	}()
 	if err == nil || err == io.EOF {
 		return
 	}