Channel rather than mutex to check is closed
This commit is contained in:
parent
10d2ad0de9
commit
bffd55f500
@ -136,9 +136,12 @@ func (r *rpcClient) stream(ctx context.Context, address string, req Request) (St
|
||||
return nil, errors.InternalServerError("go.micro.client", fmt.Sprintf("Error sending request: %v", err))
|
||||
}
|
||||
|
||||
var once sync.Once
|
||||
stream := &rpcStream{
|
||||
context: ctx,
|
||||
request: req,
|
||||
once: once,
|
||||
closed: make(chan bool),
|
||||
codec: newRpcPlusCodec(msg, c, cf),
|
||||
}
|
||||
|
||||
|
@ -13,13 +13,25 @@ import (
|
||||
type rpcStream struct {
|
||||
sync.RWMutex
|
||||
seq uint64
|
||||
closed bool
|
||||
once sync.Once
|
||||
closed chan bool
|
||||
err error
|
||||
request Request
|
||||
codec clientCodec
|
||||
context context.Context
|
||||
}
|
||||
|
||||
func (r *rpcStream) isClosed() bool {
|
||||
select {
|
||||
case _, ok := <-r.closed:
|
||||
if !ok {
|
||||
return true
|
||||
}
|
||||
default:
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func (r *rpcStream) Context() context.Context {
|
||||
return r.context
|
||||
}
|
||||
@ -32,7 +44,7 @@ func (r *rpcStream) Send(msg interface{}) error {
|
||||
r.Lock()
|
||||
defer r.Unlock()
|
||||
|
||||
if r.closed {
|
||||
if r.isClosed() {
|
||||
r.err = errShutdown
|
||||
return errShutdown
|
||||
}
|
||||
@ -57,14 +69,14 @@ func (r *rpcStream) Recv(msg interface{}) error {
|
||||
r.Lock()
|
||||
defer r.Unlock()
|
||||
|
||||
if r.closed {
|
||||
if r.isClosed() {
|
||||
r.err = errShutdown
|
||||
return errShutdown
|
||||
}
|
||||
|
||||
var resp response
|
||||
if err := r.codec.ReadResponseHeader(&resp); err != nil {
|
||||
if err == io.EOF && !r.closed {
|
||||
if err == io.EOF && !r.isClosed() {
|
||||
r.err = io.ErrUnexpectedEOF
|
||||
return io.ErrUnexpectedEOF
|
||||
}
|
||||
@ -91,7 +103,7 @@ func (r *rpcStream) Recv(msg interface{}) error {
|
||||
}
|
||||
}
|
||||
|
||||
if r.err != nil && r.err != io.EOF && !r.closed {
|
||||
if r.err != nil && r.err != io.EOF && !r.isClosed() {
|
||||
log.Println("rpc: client protocol error:", r.err)
|
||||
}
|
||||
|
||||
@ -105,8 +117,8 @@ func (r *rpcStream) Error() error {
|
||||
}
|
||||
|
||||
func (r *rpcStream) Close() error {
|
||||
r.Lock()
|
||||
defer r.Unlock()
|
||||
r.closed = true
|
||||
r.once.Do(func() {
|
||||
close(r.closed)
|
||||
})
|
||||
return r.codec.Close()
|
||||
}
|
||||
|
Loading…
Reference in New Issue
Block a user