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

     1  package client
     2  
     3  import (
     4  	"context"
     5  	"time"
     6  
     7  	"github.com/TrueCloudLab/frostfs-api-go/v2/rpc/message"
     8  	"google.golang.org/grpc"
     9  )
    10  
    11  type streamWrapper struct {
    12  	grpc.ClientStream
    13  	timeout time.Duration
    14  	cancel  context.CancelFunc
    15  }
    16  
    17  func (w streamWrapper) ReadMessage(m message.Message) error {
    18  	// Can be optimized: we can create blank message here.
    19  	gm := m.ToGRPCMessage()
    20  
    21  	err := w.withTimeout(func() error {
    22  		return w.ClientStream.RecvMsg(gm)
    23  	})
    24  	if err != nil {
    25  		return err
    26  	}
    27  
    28  	return m.FromGRPCMessage(gm)
    29  }
    30  
    31  func (w streamWrapper) WriteMessage(m message.Message) error {
    32  	return w.withTimeout(func() error {
    33  		return w.ClientStream.SendMsg(m.ToGRPCMessage())
    34  	})
    35  }
    36  
    37  func (w *streamWrapper) Close() error {
    38  	return w.withTimeout(w.ClientStream.CloseSend)
    39  }
    40  
    41  func (w *streamWrapper) withTimeout(closure func() error) error {
    42  	ch := make(chan error, 1)
    43  	go func() {
    44  		ch <- closure()
    45  		close(ch)
    46  	}()
    47  
    48  	tt := time.NewTimer(w.timeout)
    49  
    50  	select {
    51  	case err := <-ch:
    52  		tt.Stop()
    53  		return err
    54  	case <-tt.C:
    55  		w.cancel()
    56  		return context.DeadlineExceeded
    57  	}
    58  }