gitee.com/liuxuezhan/go-micro-v1.18.0@v1.0.0/client/rpc_stream.go (about)

     1  package client
     2  
     3  import (
     4  	"context"
     5  	"io"
     6  	"sync"
     7  
     8  	"gitee.com/liuxuezhan/go-micro-v1.18.0/codec"
     9  )
    10  
    11  // Implements the streamer interface
    12  type rpcStream struct {
    13  	sync.RWMutex
    14  	id       string
    15  	closed   chan bool
    16  	err      error
    17  	request  Request
    18  	response Response
    19  	codec    codec.Codec
    20  	context  context.Context
    21  
    22  	// signal whether we should send EOS
    23  	sendEOS bool
    24  
    25  	// release releases the connection back to the pool
    26  	release func(err error)
    27  }
    28  
    29  func (r *rpcStream) isClosed() bool {
    30  	select {
    31  	case <-r.closed:
    32  		return true
    33  	default:
    34  		return false
    35  	}
    36  }
    37  
    38  func (r *rpcStream) Context() context.Context {
    39  	return r.context
    40  }
    41  
    42  func (r *rpcStream) Request() Request {
    43  	return r.request
    44  }
    45  
    46  func (r *rpcStream) Response() Response {
    47  	return r.response
    48  }
    49  
    50  func (r *rpcStream) Send(msg interface{}) error {
    51  	r.Lock()
    52  	defer r.Unlock()
    53  
    54  	if r.isClosed() {
    55  		r.err = errShutdown
    56  		return errShutdown
    57  	}
    58  
    59  	req := codec.Message{
    60  		Id:       r.id,
    61  		Target:   r.request.Service(),
    62  		Method:   r.request.Method(),
    63  		Endpoint: r.request.Endpoint(),
    64  		Type:     codec.Request,
    65  	}
    66  
    67  	if err := r.codec.Write(&req, msg); err != nil {
    68  		r.err = err
    69  		return err
    70  	}
    71  
    72  	return nil
    73  }
    74  
    75  func (r *rpcStream) Recv(msg interface{}) error {
    76  	r.Lock()
    77  	defer r.Unlock()
    78  
    79  	if r.isClosed() {
    80  		r.err = errShutdown
    81  		return errShutdown
    82  	}
    83  
    84  	var resp codec.Message
    85  
    86  	if err := r.codec.ReadHeader(&resp, codec.Response); err != nil {
    87  		if err == io.EOF && !r.isClosed() {
    88  			r.err = io.ErrUnexpectedEOF
    89  			return io.ErrUnexpectedEOF
    90  		}
    91  		r.err = err
    92  		return err
    93  	}
    94  
    95  	switch {
    96  	case len(resp.Error) > 0:
    97  		// We've got an error response. Give this to the request;
    98  		// any subsequent requests will get the ReadResponseBody
    99  		// error if there is one.
   100  		if resp.Error != lastStreamResponseError {
   101  			r.err = serverError(resp.Error)
   102  		} else {
   103  			r.err = io.EOF
   104  		}
   105  		if err := r.codec.ReadBody(nil); err != nil {
   106  			r.err = err
   107  		}
   108  	default:
   109  		if err := r.codec.ReadBody(msg); err != nil {
   110  			r.err = err
   111  		}
   112  	}
   113  
   114  	return r.err
   115  }
   116  
   117  func (r *rpcStream) Error() error {
   118  	r.RLock()
   119  	defer r.RUnlock()
   120  	return r.err
   121  }
   122  
   123  func (r *rpcStream) Close() error {
   124  	select {
   125  	case <-r.closed:
   126  		return nil
   127  	default:
   128  		close(r.closed)
   129  
   130  		// send the end of stream message
   131  		if r.sendEOS {
   132  			// no need to check for error
   133  			r.codec.Write(&codec.Message{
   134  				Id:       r.id,
   135  				Target:   r.request.Service(),
   136  				Method:   r.request.Method(),
   137  				Endpoint: r.request.Endpoint(),
   138  				Type:     codec.Error,
   139  				Error:    lastStreamResponseError,
   140  			}, nil)
   141  		}
   142  
   143  		err := r.codec.Close()
   144  
   145  		// release the connection
   146  		r.release(r.Error())
   147  
   148  		// return the codec error
   149  		return err
   150  	}
   151  }