github.com/micro/go-micro/v2@v2.9.1/util/stream/stream.go (about)

     1  // Package stream encapsulates streams within streams
     2  package stream
     3  
     4  import (
     5  	"context"
     6  	"sync"
     7  
     8  	"github.com/micro/go-micro/v2/client"
     9  	"github.com/micro/go-micro/v2/codec"
    10  	"github.com/micro/go-micro/v2/metadata"
    11  	"github.com/micro/go-micro/v2/server"
    12  )
    13  
    14  type Stream interface {
    15  	Context() context.Context
    16  	SendMsg(interface{}) error
    17  	RecvMsg(interface{}) error
    18  	Close() error
    19  }
    20  
    21  type stream struct {
    22  	Stream
    23  
    24  	sync.RWMutex
    25  	err     error
    26  	request *request
    27  }
    28  
    29  type request struct {
    30  	client.Request
    31  	context context.Context
    32  }
    33  
    34  func (r *request) Codec() codec.Reader {
    35  	return r.Request.Codec().(codec.Reader)
    36  }
    37  
    38  func (r *request) Header() map[string]string {
    39  	md, _ := metadata.FromContext(r.context)
    40  	return md
    41  }
    42  
    43  func (r *request) Read() ([]byte, error) {
    44  	return nil, nil
    45  }
    46  
    47  func (s *stream) Request() server.Request {
    48  	return s.request
    49  }
    50  
    51  func (s *stream) Send(v interface{}) error {
    52  	err := s.Stream.SendMsg(v)
    53  	if err != nil {
    54  		s.Lock()
    55  		s.err = err
    56  		s.Unlock()
    57  	}
    58  	return err
    59  }
    60  
    61  func (s *stream) Recv(v interface{}) error {
    62  	err := s.Stream.RecvMsg(v)
    63  	if err != nil {
    64  		s.Lock()
    65  		s.err = err
    66  		s.Unlock()
    67  	}
    68  	return err
    69  }
    70  
    71  func (s *stream) Error() error {
    72  	s.RLock()
    73  	defer s.RUnlock()
    74  	return s.err
    75  }
    76  
    77  // New returns a new encapsulated stream
    78  // Proto stream within a server.Stream
    79  func New(service, endpoint string, req interface{}, s Stream) server.Stream {
    80  	return &stream{
    81  		Stream: s,
    82  		request: &request{
    83  			context: s.Context(),
    84  			Request: client.DefaultClient.NewRequest(service, endpoint, req),
    85  		},
    86  	}
    87  }