go-micro.dev/v5@v5.12.0/client/rpc_stream.go (about)

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