github.com/annwntech/go-micro/v2@v2.9.5/client/rpc_stream.go (about)

     1  package client
     2  
     3  import (
     4  	"context"
     5  	"io"
     6  	"sync"
     7  
     8  	"github.com/annwntech/go-micro/v2/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  	r.Unlock()
    87  	err := r.codec.ReadHeader(&resp, codec.Response)
    88  	r.Lock()
    89  	if err != nil {
    90  		if err == io.EOF && !r.isClosed() {
    91  			r.err = io.ErrUnexpectedEOF
    92  			return io.ErrUnexpectedEOF
    93  		}
    94  		r.err = err
    95  		return err
    96  	}
    97  
    98  	switch {
    99  	case len(resp.Error) > 0:
   100  		// We've got an error response. Give this to the request;
   101  		// any subsequent requests will get the ReadResponseBody
   102  		// error if there is one.
   103  		if resp.Error != lastStreamResponseError {
   104  			r.err = serverError(resp.Error)
   105  		} else {
   106  			r.err = io.EOF
   107  		}
   108  		r.Unlock()
   109  		err = r.codec.ReadBody(nil)
   110  		r.Lock()
   111  		if err != nil {
   112  			r.err = err
   113  		}
   114  	default:
   115  		r.Unlock()
   116  		err = r.codec.ReadBody(msg)
   117  		r.Lock()
   118  		if err != nil {
   119  			r.err = err
   120  		}
   121  	}
   122  
   123  	return r.err
   124  }
   125  
   126  func (r *rpcStream) Error() error {
   127  	r.RLock()
   128  	defer r.RUnlock()
   129  	return r.err
   130  }
   131  
   132  func (r *rpcStream) Close() error {
   133  	r.Lock()
   134  
   135  	select {
   136  	case <-r.closed:
   137  		r.Unlock()
   138  		return nil
   139  	default:
   140  		close(r.closed)
   141  		r.Unlock()
   142  
   143  		// send the end of stream message
   144  		if r.sendEOS {
   145  			// no need to check for error
   146  			r.codec.Write(&codec.Message{
   147  				Id:       r.id,
   148  				Target:   r.request.Service(),
   149  				Method:   r.request.Method(),
   150  				Endpoint: r.request.Endpoint(),
   151  				Type:     codec.Error,
   152  				Error:    lastStreamResponseError,
   153  			}, nil)
   154  		}
   155  
   156  		err := r.codec.Close()
   157  
   158  		// release the connection
   159  		r.release(r.Error())
   160  
   161  		// return the codec error
   162  		return err
   163  	}
   164  }