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