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

     1  package server
     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  	codec   codec.Codec
    17  	context context.Context
    18  	id      string
    19  	sync.RWMutex
    20  	closed bool
    21  }
    22  
    23  func (r *rpcStream) Context() context.Context {
    24  	return r.context
    25  }
    26  
    27  func (r *rpcStream) Request() Request {
    28  	return r.request
    29  }
    30  
    31  func (r *rpcStream) Send(msg interface{}) error {
    32  	r.Lock()
    33  	defer r.Unlock()
    34  
    35  	resp := codec.Message{
    36  		Target:   r.request.Service(),
    37  		Method:   r.request.Method(),
    38  		Endpoint: r.request.Endpoint(),
    39  		Id:       r.id,
    40  		Type:     codec.Response,
    41  	}
    42  
    43  	if err := r.codec.Write(&resp, msg); err != nil {
    44  		r.err = err
    45  	}
    46  
    47  	return nil
    48  }
    49  
    50  func (r *rpcStream) Recv(msg interface{}) error {
    51  	req := new(codec.Message)
    52  	req.Type = codec.Request
    53  
    54  	err := r.codec.ReadHeader(req, req.Type)
    55  	r.Lock()
    56  	defer r.Unlock()
    57  	if err != nil {
    58  		// discard body
    59  		r.codec.ReadBody(nil)
    60  		r.err = err
    61  		return err
    62  	}
    63  
    64  	// check the error
    65  	if len(req.Error) > 0 {
    66  		// Check the client closed the stream
    67  		switch req.Error {
    68  		case errLastStreamResponse.Error():
    69  			// discard body
    70  			r.Unlock()
    71  			r.codec.ReadBody(nil)
    72  			r.Lock()
    73  			r.err = io.EOF
    74  			return io.EOF
    75  		default:
    76  			return errors.New(req.Error)
    77  		}
    78  	}
    79  
    80  	// we need to stay up to date with sequence numbers
    81  	r.id = req.Id
    82  	r.Unlock()
    83  	err = r.codec.ReadBody(msg)
    84  	r.Lock()
    85  	if err != nil {
    86  		r.err = err
    87  		return err
    88  	}
    89  
    90  	return nil
    91  }
    92  
    93  func (r *rpcStream) Error() error {
    94  	r.RLock()
    95  	defer r.RUnlock()
    96  	return r.err
    97  }
    98  
    99  func (r *rpcStream) Close() error {
   100  	r.Lock()
   101  	defer r.Unlock()
   102  	r.closed = true
   103  	return r.codec.Close()
   104  }