github.com/TrueCloudLab/frostfs-api-go/v2@v2.0.0-20230228134343-196241c4e79a/rpc/client/flows.go (about)

     1  package client
     2  
     3  import (
     4  	"errors"
     5  	"io"
     6  	"sync"
     7  
     8  	"github.com/TrueCloudLab/frostfs-api-go/v2/rpc/common"
     9  	"github.com/TrueCloudLab/frostfs-api-go/v2/rpc/message"
    10  )
    11  
    12  // SendUnary initializes communication session by RPC info, performs unary RPC
    13  // and closes the session.
    14  func SendUnary(cli *Client, info common.CallMethodInfo, req, resp message.Message, opts ...CallOption) error {
    15  	rw, err := cli.Init(info, opts...)
    16  	if err != nil {
    17  		return err
    18  	}
    19  
    20  	err = rw.WriteMessage(req)
    21  	if err != nil {
    22  		return err
    23  	}
    24  
    25  	err = rw.ReadMessage(resp)
    26  	if err != nil {
    27  		return err
    28  	}
    29  
    30  	return rw.Close()
    31  }
    32  
    33  // MessageWriterCloser wraps MessageWriter
    34  // and io.Closer interfaces.
    35  type MessageWriterCloser interface {
    36  	MessageWriter
    37  	io.Closer
    38  }
    39  
    40  type clientStreamWriterCloser struct {
    41  	MessageReadWriter
    42  
    43  	resp message.Message
    44  }
    45  
    46  func (c *clientStreamWriterCloser) Close() error {
    47  	err := c.MessageReadWriter.Close()
    48  	if err != nil {
    49  		return err
    50  	}
    51  
    52  	return c.ReadMessage(c.resp)
    53  }
    54  
    55  // OpenClientStream initializes communication session by RPC info, opens client-side stream
    56  // and returns its interface.
    57  //
    58  // All stream writes must be performed before the closing. Close must be called once.
    59  func OpenClientStream(cli *Client, info common.CallMethodInfo, resp message.Message, opts ...CallOption) (MessageWriterCloser, error) {
    60  	rw, err := cli.Init(info, opts...)
    61  	if err != nil {
    62  		return nil, err
    63  	}
    64  
    65  	return &clientStreamWriterCloser{
    66  		MessageReadWriter: rw,
    67  		resp:              resp,
    68  	}, nil
    69  }
    70  
    71  // MessageReaderCloser wraps MessageReader
    72  // and io.Closer interface.
    73  type MessageReaderCloser interface {
    74  	MessageReader
    75  	io.Closer
    76  }
    77  
    78  type serverStreamReaderCloser struct {
    79  	rw MessageReadWriter
    80  
    81  	once sync.Once
    82  
    83  	req message.Message
    84  }
    85  
    86  func (s *serverStreamReaderCloser) ReadMessage(msg message.Message) error {
    87  	var err error
    88  
    89  	s.once.Do(func() {
    90  		err = s.rw.WriteMessage(s.req)
    91  	})
    92  
    93  	if err != nil {
    94  		return err
    95  	}
    96  
    97  	err = s.rw.ReadMessage(msg)
    98  	if !errors.Is(err, io.EOF) {
    99  		return err
   100  	}
   101  
   102  	err = s.rw.Close()
   103  	if err != nil {
   104  		return err
   105  	}
   106  
   107  	return io.EOF
   108  }
   109  
   110  // OpenServerStream initializes communication session by RPC info, opens server-side stream
   111  // and returns its interface.
   112  //
   113  // All stream reads must be performed before the closing. Close must be called once.
   114  func OpenServerStream(cli *Client, info common.CallMethodInfo, req message.Message, opts ...CallOption) (MessageReader, error) {
   115  	rw, err := cli.Init(info, opts...)
   116  	if err != nil {
   117  		return nil, err
   118  	}
   119  
   120  	return &serverStreamReaderCloser{
   121  		rw:  rw,
   122  		req: req,
   123  	}, nil
   124  }