[client] Fix/grpc retry

flow/client/client.go

@@ -14,7 +14,6 @@ import (
 	log "github.com/sirupsen/logrus"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/codes"
-	"google.golang.org/grpc/connectivity"
 	"google.golang.org/grpc/credentials"
 	"google.golang.org/grpc/credentials/insecure"
 	"google.golang.org/grpc/keepalive"
@@ -26,11 +25,15 @@ import (
 	"github.com/netbirdio/netbird/util/wsproxy"
 )
 
+var ErrClientClosed = errors.New("client is closed")
+
 type GRPCClient struct {
 	realClient proto.FlowServiceClient
 	clientConn *grpc.ClientConn
 	stream     proto.FlowService_EventsClient
-	streamMu   sync.Mutex
+	opts       []grpc.DialOption
+	closed     bool       // prevent creating conn in the middle of the Close
+	mu         sync.Mutex // protects clientConn, realClient, stream, and closed
 }
 
 func NewClient(addr, payload, signature string, interval time.Duration) (*GRPCClient, error) {
@@ -73,28 +76,71 @@ func NewClient(addr, payload, signature string, interval time.Duration) (*GRPCCl
 	return &GRPCClient{
 		realClient: proto.NewFlowServiceClient(conn),
 		clientConn: conn,
+		opts:       opts,
 	}, nil
 }
 
 func (c *GRPCClient) Close() error {
-	c.streamMu.Lock()
-	defer c.streamMu.Unlock()
-
+	c.mu.Lock()
+	c.closed = true
 	c.stream = nil
-	if err := c.clientConn.Close(); err != nil && !errors.Is(err, context.Canceled) {
+	conn := c.clientConn
+	c.clientConn = nil
+	c.mu.Unlock()
+
+	if err := conn.Close(); err != nil && !errors.Is(err, context.Canceled) {
 		return fmt.Errorf("close client connection: %w", err)
 	}
 
 	return nil
 }
 
+func (c *GRPCClient) Send(event *proto.FlowEvent) error {
+	c.mu.Lock()
+	stream := c.stream
+	c.mu.Unlock()
+
+	if stream == nil {
+		return errors.New("stream not initialized")
+	}
+
+	if err := stream.Send(event); err != nil {
+		return fmt.Errorf("send flow event: %w", err)
+	}
+
+	return nil
+}
+
 func (c *GRPCClient) Receive(ctx context.Context, interval time.Duration, msgHandler func(msg *proto.FlowEventAck) error) error {
 	backOff := defaultBackoff(ctx, interval)
 	operation := func() error {
-		if err := c.establishStreamAndReceive(ctx, msgHandler); err != nil {
-			if s, ok := status.FromError(err); ok && s.Code() == codes.Canceled {
-				return fmt.Errorf("receive: %w: %w", err, context.Canceled)
+		stream, err := c.establishStream(ctx)
+		if err != nil {
+			if isCancellation(err) {
+				return backoff.Permanent(err)
+			}
+			return err
+		}
+
+		// we have a successful connection, reset the backoff so that if receive fails later,
+		// the next retry starts with a short delay instead of continuing the already-elapsed timer
+		backOff.Reset()
+
+		if err := c.receive(stream, msgHandler); err != nil {
+			if isCancellation(err) {
+				return backoff.Permanent(err)
+			}
+			// RST_STREAM/PROTOCOL_ERROR — connection is corrupt, recreate immediately
+			if s, ok := status.FromError(err); ok && s.Code() == codes.Internal {
+				log.Warnf("connection corrupt, attempting reconnection: %v", err)
+				if err := c.recreateConnection(); err != nil {
+					log.Errorf("recreate connection: %v", err)
+					return err
+				}
+				log.Infof("connection recreated successfully")
+				return fmt.Errorf("connection recreated, re-establishing stream")
 			}
+
 			log.Errorf("receive failed: %v", err)
 			return fmt.Errorf("receive: %w", err)
 		}
@@ -108,30 +154,69 @@ func (c *GRPCClient) Receive(ctx context.Context, interval time.Duration, msgHan
 	return nil
 }
 
-func (c *GRPCClient) establishStreamAndReceive(ctx context.Context, msgHandler func(msg *proto.FlowEventAck) error) error {
-	if c.clientConn.GetState() == connectivity.Shutdown {
-		return errors.New("connection to flow receiver has been shut down")
+func (c *GRPCClient) recreateConnection() error {
+	c.mu.Lock()
+	if c.closed {
+		c.mu.Unlock()
+		return backoff.Permanent(ErrClientClosed)
 	}
 
-	stream, err := c.realClient.Events(ctx, grpc.WaitForReady(true))
+	conn, err := grpc.NewClient(c.clientConn.Target(), c.opts...)
 	if err != nil {
-		return fmt.Errorf("create event stream: %w", err)
+		c.mu.Unlock()
+		return fmt.Errorf("create new connection: %w", err)
+	}
+
+	old := c.clientConn
+	c.clientConn = conn
+	c.realClient = proto.NewFlowServiceClient(conn)
+	c.stream = nil
+	c.mu.Unlock()
+
+	_ = old.Close()
+
+	return nil
+}
+
+func (c *GRPCClient) establishStream(ctx context.Context) (proto.FlowService_EventsClient, error) {
+	c.mu.Lock()
+	if c.closed {
+		c.mu.Unlock()
+		return nil, backoff.Permanent(ErrClientClosed)
 	}
+	cl := c.realClient
+	c.mu.Unlock()
 
-	err = stream.Send(&proto.FlowEvent{IsInitiator: true})
+	// open stream outside the lock — blocking operation
+	stream, err := cl.Events(ctx)
 	if err != nil {
-		log.Infof("failed to send initiator message to flow receiver but will attempt to continue. Error: %s", err)
+		return nil, fmt.Errorf("create event stream: %w", err)
+	}
+	streamReady := false
+	defer func() {
+		if !streamReady {
+			_ = stream.CloseSend()
+		}
+	}()
+
+	if err = stream.Send(&proto.FlowEvent{IsInitiator: true}); err != nil {
+		return nil, fmt.Errorf("send initiator: %w", err)
 	}
 
 	if err = checkHeader(stream); err != nil {
-		return fmt.Errorf("check header: %w", err)
+		return nil, fmt.Errorf("check header: %w", err)
 	}
 
-	c.streamMu.Lock()
+	c.mu.Lock()
+	if c.closed {
+		c.mu.Unlock()
+		return nil, backoff.Permanent(ErrClientClosed)
+	}
 	c.stream = stream
-	c.streamMu.Unlock()
+	c.mu.Unlock()
+	streamReady = true
 
-	return c.receive(stream, msgHandler)
+	return stream, nil
 }
 
 func (c *GRPCClient) receive(stream proto.FlowService_EventsClient, msgHandler func(msg *proto.FlowEventAck) error) error {
@@ -178,18 +263,12 @@ func defaultBackoff(ctx context.Context, interval time.Duration) backoff.BackOff
 	}, ctx)
 }
 
-func (c *GRPCClient) Send(event *proto.FlowEvent) error {
-	c.streamMu.Lock()
-	stream := c.stream
-	c.streamMu.Unlock()
-
-	if stream == nil {
-		return errors.New("stream not initialized")
+func isCancellation(err error) bool {
+	if errors.Is(err, context.Canceled) {
+		return true
 	}
-
-	if err := stream.Send(event); err != nil {
-		return fmt.Errorf("send flow event: %w", err)
+	if s, ok := status.FromError(err); ok {
+		return s.Code() == codes.Canceled
 	}
-
-	return nil
+	return false
 }