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

     1  package server
     2  
     3  import (
     4  	"context"
     5  	"errors"
     6  	"io"
     7  	"sync"
     8  
     9  	"gitee.com/liuxuezhan/go-micro-v1.18.0/codec"
    10  )
    11  
    12  // Implements the Streamer interface
    13  type rpcStream struct {
    14  	sync.RWMutex
    15  	id      string
    16  	closed  bool
    17  	err     error
    18  	request Request
    19  	codec   codec.Codec
    20  	context context.Context
    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  	r.Lock()
    52  	defer r.Unlock()
    53  
    54  	req := new(codec.Message)
    55  	req.Type = codec.Request
    56  
    57  	if err := r.codec.ReadHeader(req, req.Type); 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 lastStreamResponseError.Error():
    69  			// discard body
    70  			r.codec.ReadBody(nil)
    71  			r.err = io.EOF
    72  			return io.EOF
    73  		default:
    74  			return errors.New(req.Error)
    75  		}
    76  	}
    77  
    78  	// we need to stay up to date with sequence numbers
    79  	r.id = req.Id
    80  	if err := r.codec.ReadBody(msg); err != nil {
    81  		r.err = err
    82  		return err
    83  	}
    84  
    85  	return nil
    86  }
    87  
    88  func (r *rpcStream) Error() error {
    89  	r.RLock()
    90  	defer r.RUnlock()
    91  	return r.err
    92  }
    93  
    94  func (r *rpcStream) Close() error {
    95  	r.Lock()
    96  	defer r.Unlock()
    97  	r.closed = true
    98  	return r.codec.Close()
    99  }