[#327] rpc: Fix mem leak

gRPC stream must be closed by `cancel` to prevent memleak.

Signed-off-by: Dmitrii Stepanov <d.stepanov@yadro.com>

api/rpc/client/flows.go

@@ -12,18 +12,20 @@ import (
 // SendUnary initializes communication session by RPC info, performs unary RPC
 // and closes the session.
 func SendUnary(cli *Client, info common.CallMethodInfo, req, resp message.Message, opts ...CallOption) error {
-	rw, err := cli.Init(info, opts...)
+	rw, err := cli.initInternal(info, opts...)
 	if err != nil {
 		return err
 	}
 
 	err = rw.WriteMessage(req)
 	if err != nil {
+		rw.cancel()
 		return err
 	}
 
 	err = rw.ReadMessage(resp)
 	if err != nil {
+		rw.cancel()
 		return err
 	}
 
@@ -38,18 +40,28 @@ type MessageWriterCloser interface {
 }
 
 type clientStreamWriterCloser struct {
-	MessageReadWriter
-
+	sw   *streamWrapper
 	resp message.Message
 }
 
+// WriteMessage implements MessageWriterCloser.
+func (c *clientStreamWriterCloser) WriteMessage(m message.Message) error {
+	return c.sw.WriteMessage(m)
+}
+
 func (c *clientStreamWriterCloser) Close() error {
-	err := c.MessageReadWriter.Close()
+	err := c.sw.closeSend()
 	if err != nil {
+		c.sw.cancel()
 		return err
 	}
 
-	return c.ReadMessage(c.resp)
+	if err = c.sw.ReadMessage(c.resp); err != nil {
+		c.sw.cancel()
+		return err
+	}
+
+	return c.sw.Close()
 }
 
 // OpenClientStream initializes communication session by RPC info, opens client-side stream
@@ -57,14 +69,14 @@ func (c *clientStreamWriterCloser) Close() error {
 //
 // All stream writes must be performed before the closing. Close must be called once.
 func OpenClientStream(cli *Client, info common.CallMethodInfo, resp message.Message, opts ...CallOption) (MessageWriterCloser, error) {
-	rw, err := cli.Init(info, opts...)
+	rw, err := cli.initInternal(info, opts...)
 	if err != nil {
 		return nil, err
 	}
 
 	return &clientStreamWriterCloser{
-		MessageReadWriter: rw,
-		resp:              resp,
+		sw:   rw,
+		resp: resp,
 	}, nil
 }
 
@@ -76,7 +88,7 @@ type MessageReaderCloser interface {
 }
 
 type serverStreamReaderCloser struct {
-	rw MessageReadWriter
+	rw *streamWrapper
 
 	once sync.Once
 
@@ -91,11 +103,15 @@ func (s *serverStreamReaderCloser) ReadMessage(msg message.Message) error {
 	})
 
 	if err != nil {
+		s.rw.cancel()
 		return err
 	}
 
 	err = s.rw.ReadMessage(msg)
 	if !errors.Is(err, io.EOF) {
+		if err != nil {
+			s.rw.cancel()
+		}
 		return err
 	}
 
@@ -112,7 +128,7 @@ func (s *serverStreamReaderCloser) ReadMessage(msg message.Message) error {
 //
 // All stream reads must be performed before the closing. Close must be called once.
 func OpenServerStream(cli *Client, info common.CallMethodInfo, req message.Message, opts ...CallOption) (MessageReader, error) {
-	rw, err := cli.Init(info, opts...)
+	rw, err := cli.initInternal(info, opts...)
 	if err != nil {
 		return nil, err
 	}