Merge remote-tracking branch 'origin/main' into revert/component-types

Signed-off-by: Dmitri Dolguikh <dmitri.external@netbird.io>

shared/management/client/grpc.go

@@ -21,6 +21,8 @@ import (
 	"google.golang.org/grpc/connectivity"
 
 	nbgrpc "github.com/netbirdio/netbird/client/grpc"
+	"github.com/netbirdio/netbird/client/netstate"
+	"github.com/netbirdio/netbird/client/netsweep"
 	"github.com/netbirdio/netbird/client/system"
 	"github.com/netbirdio/netbird/encryption"
 	"github.com/netbirdio/netbird/shared/management/domain"
@@ -62,6 +64,13 @@ type GrpcClient struct {
 	connStateCallbackLock sync.RWMutex
 	serverURL             string
 
+	// netState gates the stream retry loop on OS-reported network
+	// availability; nil (the default) disables gating.
+	netState *netstate.State
+
+	// sweeper cuts the transport connections on network change; nil disables it.
+	sweeper *netsweep.Sweeper
+
 	// syncStreamErr holds the last Sync stream error, or nil while the stream
 	// is established and healthy. GetServerKey succeeds even when the peer
 	// cannot sync (e.g. the server returns "settings not found"), so the
@@ -111,16 +120,43 @@ func MaxRecvMsgSize() int {
 	return size
 }
 
+// Option configures optional GrpcClient behavior.
+type Option func(*GrpcClient)
+
+// WithNetworkState injects the OS network availability state that gates the
+// stream retry loop; without it gating is disabled.
+func WithNetworkState(netState *netstate.State) Option {
+	return func(c *GrpcClient) { c.netState = netState }
+}
+
+// WithSweeper injects the network change sweeper.
+func WithSweeper(sweeper *netsweep.Sweeper) Option {
+	return func(c *GrpcClient) { c.sweeper = sweeper }
+}
+
 // NewClient creates a new client to Management service
-func NewClient(ctx context.Context, addr string, ourPrivateKey wgtypes.Key, tlsEnabled bool) (*GrpcClient, error) {
-	var conn *grpc.ClientConn
+func NewClient(ctx context.Context, addr string, ourPrivateKey wgtypes.Key, tlsEnabled bool, opts ...Option) (*GrpcClient, error) {
+	// Options apply before dialing: the sweeper must wrap the first connection too.
+	c := &GrpcClient{
+		key:                   ourPrivateKey,
+		ctx:                   ctx,
+		connStateCallbackLock: sync.RWMutex{},
+		serverURL:             addr,
+	}
+	for _, opt := range opts {
+		opt(c)
+	}
 
 	var extraOpts []grpc.DialOption
 	if maxSize := MaxRecvMsgSize(); maxSize > 0 {
 		extraOpts = append(extraOpts, grpc.WithDefaultCallOptions(grpc.MaxCallRecvMsgSize(maxSize)))
 		log.Infof("management gRPC max receive message size set to %d bytes", maxSize)
 	}
+	if c.sweeper != nil {
+		extraOpts = append(extraOpts, nbgrpc.WithSweeper(c.sweeper))
+	}
 
+	var conn *grpc.ClientConn
 	operation := func() error {
 		var err error
 		conn, err = nbgrpc.CreateConnection(ctx, addr, tlsEnabled, wsproxy.ManagementComponent, extraOpts...)
@@ -136,16 +172,9 @@ func NewClient(ctx context.Context, addr string, ourPrivateKey wgtypes.Key, tlsE
 		return nil, err
 	}
 
-	realClient := proto.NewManagementServiceClient(conn)
-
-	return &GrpcClient{
-		key:                   ourPrivateKey,
-		realClient:            realClient,
-		ctx:                   ctx,
-		conn:                  conn,
-		connStateCallbackLock: sync.RWMutex{},
-		serverURL:             addr,
-	}, nil
+	c.conn = conn
+	c.realClient = proto.NewManagementServiceClient(conn)
+	return c, nil
 }
 
 // GetServerURL returns the management server URL
@@ -206,16 +235,33 @@ func (c *GrpcClient) withMgmtStream(
 	ctx context.Context,
 	handler func(ctx context.Context, serverPubKey wgtypes.Key, backOff backoff.BackOff) error,
 ) error {
-	backOff := defaultBackoff(ctx)
+	backOff := c.sweeper.QuickRetryBackoff(ctx, defaultBackoff(ctx), c.netState)
 	operation := func() error {
-		log.Debugf("management connection state %v", c.conn.GetState())
-		connState := c.conn.GetState()
+		// suspend reconnect attempts while the OS reports no usable network.
+		// Wait only errors on a cancelled context, which means shutdown, so
+		// stop the loop without reporting a failure.
+		if waited, err := c.netState.Wait(ctx); err != nil {
+			log.Debugf("management connection context has been canceled while offline, this usually indicates shutdown")
+			return nil //nolint:nilerr // a cancelled context means shutdown, not a retryable failure
+		} else if waited {
+			backOff.Reset()
+		}
 
+		connState := c.conn.GetState()
+		log.Debugf("management connection state %v", connState)
 		if connState == connectivity.Shutdown {
 			return backoff.Permanent(fmt.Errorf("connection to management has been shut down"))
-		} else if !(connState == connectivity.Ready || connState == connectivity.Idle) {
+		}
+		if !(connState == connectivity.Ready || connState == connectivity.Idle) {
+			// A dial may already be in flight (e.g. the other stream triggered
+			// it after a network change); wait for it to settle and proceed if
+			// the channel became usable, instead of burning a backoff round on
+			// a successful dial. A failed dial errors out as before.
 			c.conn.WaitForStateChange(ctx, connState)
-			return fmt.Errorf("connection to management is not ready and in %s state", connState)
+			connState = c.conn.GetState()
+			if !(connState == connectivity.Ready || connState == connectivity.Idle) {
+				return fmt.Errorf("connection to management is not ready and in %s state", connState)
+			}
 		}
 
 		serverPubKey, err := c.getServerPublicKey()
@@ -227,7 +273,7 @@ func (c *GrpcClient) withMgmtStream(
 		return handler(ctx, *serverPubKey, backOff)
 	}
 
-	err := backoff.Retry(operation, backOff)
+	err := nbgrpc.Retry(ctx, operation, backOff, c.netState)
 	if err != nil {
 		log.Warnf("exiting the Management service connection retry loop due to the unrecoverable error: %s", err)
 	}