github.com/tickoalcantara12/micro/v3@v3.0.0-20221007104245-9d75b9bcbab9/service/server/mucp/rpc_stream.go (about)

     1  // Copyright 2020 Asim Aslam
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     https://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  //
    15  // Original source: github.com/micro/go-micro/v3/server/mucp/rpc_stream.go
    16  
    17  package mucp
    18  
    19  import (
    20  	"context"
    21  	"errors"
    22  	"io"
    23  	"sync"
    24  
    25  	"github.com/tickoalcantara12/micro/v3/service/server"
    26  	"github.com/tickoalcantara12/micro/v3/util/codec"
    27  )
    28  
    29  // Implements the Streamer interface
    30  type rpcStream struct {
    31  	sync.RWMutex
    32  	id      string
    33  	closed  bool
    34  	err     error
    35  	request server.Request
    36  	codec   codec.Codec
    37  	context context.Context
    38  }
    39  
    40  func (r *rpcStream) Context() context.Context {
    41  	return r.context
    42  }
    43  
    44  func (r *rpcStream) Request() server.Request {
    45  	return r.request
    46  }
    47  
    48  func (r *rpcStream) Send(msg interface{}) error {
    49  	r.Lock()
    50  	defer r.Unlock()
    51  
    52  	resp := codec.Message{
    53  		Target:   r.request.Service(),
    54  		Method:   r.request.Method(),
    55  		Endpoint: r.request.Endpoint(),
    56  		Id:       r.id,
    57  		Type:     codec.Response,
    58  	}
    59  
    60  	if err := r.codec.Write(&resp, msg); err != nil {
    61  		r.err = err
    62  	}
    63  
    64  	return nil
    65  }
    66  
    67  func (r *rpcStream) Recv(msg interface{}) error {
    68  	req := new(codec.Message)
    69  	req.Type = codec.Request
    70  
    71  	err := r.codec.ReadHeader(req, req.Type)
    72  	r.Lock()
    73  	defer r.Unlock()
    74  	if err != nil {
    75  		// discard body
    76  		r.codec.ReadBody(nil)
    77  		r.err = err
    78  		return err
    79  	}
    80  
    81  	// check the error
    82  	if len(req.Error) > 0 {
    83  		// Check the client closed the stream
    84  		switch req.Error {
    85  		case lastStreamResponseError.Error():
    86  			// discard body
    87  			r.Unlock()
    88  			r.codec.ReadBody(nil)
    89  			r.Lock()
    90  			r.err = io.EOF
    91  			return io.EOF
    92  		default:
    93  			return errors.New(req.Error)
    94  		}
    95  	}
    96  
    97  	// we need to stay up to date with sequence numbers
    98  	r.id = req.Id
    99  	r.Unlock()
   100  	err = r.codec.ReadBody(msg)
   101  	r.Lock()
   102  	if err != nil {
   103  		r.err = err
   104  		return err
   105  	}
   106  
   107  	return nil
   108  }
   109  
   110  func (r *rpcStream) Error() error {
   111  	r.RLock()
   112  	defer r.RUnlock()
   113  	return r.err
   114  }
   115  
   116  func (r *rpcStream) Close() error {
   117  	r.Lock()
   118  	defer r.Unlock()
   119  	r.closed = true
   120  	return r.codec.Close()
   121  }