gitee.com/sasukebo/go-micro/v4@v4.7.1/client/rpc_stream.go (about)

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