Always close tracing when a stream goes wrong

This commit is contained in:
iamqizhao 2016-02-22 18:02:10 -08:00
parent 0be94ab3f5
commit 5085c7628d
1 changed files with 9 additions and 0 deletions

View File

@ -204,6 +204,9 @@ func (cs *clientStream) SendMsg(m interface{}) (err error) {
cs.mu.Unlock() cs.mu.Unlock()
} }
defer func() { defer func() {
if err != nil {
cs.finish(err)
}
if err == nil || err == io.EOF { if err == nil || err == io.EOF {
return return
} }
@ -251,6 +254,7 @@ func (cs *clientStream) RecvMsg(m interface{}) (err error) {
} }
if err == io.EOF { if err == io.EOF {
if cs.s.StatusCode() == codes.OK { if cs.s.StatusCode() == codes.OK {
cs.finish(err)
return nil return nil
} }
return Errorf(cs.s.StatusCode(), cs.s.StatusDesc()) 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) { func (cs *clientStream) CloseSend() (err error) {
err = cs.t.Write(cs.s, nil, &transport.Options{Last: true}) 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 { if err == nil || err == io.EOF {
return return
} }