Addressing Evan's comments

cns/cnireconciler/podinfoprovider.go

@@ -17,9 +17,14 @@ import (
 
 // NewCNIPodInfoProvider returns an implementation of cns.PodInfoByIPProvider
 // that execs out to the CNI and uses the response to build the PodInfo map.
-// if stateMigration flag is set to true it will also returns a map of containerID->EndpointInfo
-func NewCNIPodInfoProvider(stateMigration bool) (cns.PodInfoByIPProvider, map[string]*restserver.EndpointInfo, error) {
-	return newCNIPodInfoProvider(exec.New(), stateMigration)
+// if EnableStateMigration flag is set to true it will also returns a map of containerID->EndpointInfo
+func NewCNIPodInfoProvider(enableStateMigration bool) (cns.PodInfoByIPProvider, map[string]*restserver.EndpointInfo, error) {
+	podInfoByIPProvider, cniState, err := newCNIPodInfoProvider(exec.New())
+	if enableStateMigration {
+		endpointState := cniStateToCnsEndpointState(cniState)
+		return podInfoByIPProvider, endpointState, err
+	}
+	return podInfoByIPProvider, nil, err
 }
 
 func NewCNSPodInfoProvider(endpointStore store.KeyValueStore) (cns.PodInfoByIPProvider, error) {
@@ -43,7 +48,7 @@ func newCNSPodInfoProvider(endpointStore store.KeyValueStore) (cns.PodInfoByIPPr
 	}), nil
 }
 
-func newCNIPodInfoProvider(exc exec.Interface, stateMigration bool) (cns.PodInfoByIPProvider, map[string]*restserver.EndpointInfo, error) {
+func newCNIPodInfoProvider(exc exec.Interface) (cns.PodInfoByIPProvider, *api.AzureCNIState, error) {
 	cli := client.New(exc)
 	state, err := cli.GetEndpointState()
 	if err != nil {
@@ -52,15 +57,9 @@ func newCNIPodInfoProvider(exc exec.Interface, stateMigration bool) (cns.PodInfo
 	for containerID, endpointInfo := range state.ContainerInterfaces {
 		logger.Printf("state dump from CNI: [%+v], [%+v]", containerID, endpointInfo)
 	}
-	var endpointState map[string]*restserver.EndpointInfo
-	if stateMigration {
-		endpointState = cniStateToCnsEndpointState(state)
-	} else {
-		endpointState = nil
-	}
 	return cns.PodInfoByIPProviderFunc(func() (map[string]cns.PodInfo, error) {
 		return cniStateToPodInfoByIP(state)
-	}), endpointState, err
+	}), state, err
 }
 
 // cniStateToPodInfoByIP converts an AzureCNIState dumped from a CNI exec

cns/cnireconciler/podinfoprovider_test.go

@@ -76,7 +76,7 @@ func TestNewCNIPodInfoProvider(t *testing.T) {
 	for _, tt := range tests {
 		tt := tt
 		t.Run(tt.name, func(t *testing.T) {
-			got, endpointState, err := newCNIPodInfoProvider(tt.exec, true)
+			got, endpointState, err := newCNIPodInfoProvider(tt.exec)
 			if tt.wantErr {
 				assert.Error(t, err)
 				return

cns/configuration/configuration.go

@@ -35,6 +35,7 @@ type CNSConfig struct {
 	EnableCNIConflistGeneration bool
 	EnableIPAMv2                bool
 	EnablePprof                 bool
+	EnableStateMigration        bool
 	EnableSubnetScarcity        bool
 	EnableSwiftV2               bool
 	InitializeFromCNI           bool
@@ -56,7 +57,6 @@ type CNSConfig struct {
 	UseHTTPS                    bool
 	WatchPods                   bool `json:"-"`
 	WireserverIP                string
-	StateMigration              bool
 }
 
 type TelemetrySettings struct {

cns/service/main.go

@@ -1516,27 +1516,25 @@ func createOrUpdateNodeInfoCRD(ctx context.Context, restConfig *rest.Config, nod
 
 // InitializeStateFromCNS initilizes CNS Endpoint State from CNS or perform the Migration of state from statefull CNI.
 func InitializeStateFromCNS(cnsconfig *configuration.CNSConfig, endpointStateStore store.KeyValueStore) (cns.PodInfoByIPProvider, error) {
+	if cnsconfig.EnableStateMigration && !endpointStateStore.Exists() { // initilize form CNI and perform state migration
+		logger.Printf("StatelessCNI Migration is enabled")
+		logger.Printf("initializing from Statefull CNI")
+		var endpointState map[string]*restserver.EndpointInfo
+		podInfoByIPProvider, endpointState, err := cnireconciler.NewCNIPodInfoProvider(cnsconfig.EnableStateMigration)
+		if err != nil {
+			return nil, errors.Wrap(err, "failed to create CNI PodInfoProvider")
+		}
+		err = endpointStateStore.Write(restserver.EndpointStoreKey, endpointState)
+		if err != nil {
+			return nil, fmt.Errorf("failed to write endpoint state to store: %w", err)
+		}
+		return podInfoByIPProvider, nil
+	}
+	// initilize form CNS and avoid state migration
 	logger.Printf("Initializing from self managed endpoint store")
 	podInfoByIPProvider, err := cnireconciler.NewCNSPodInfoProvider(endpointStateStore) // get reference to endpoint state store from rest server
 	if err != nil {
-		if errors.Is(err, store.ErrKeyNotFound) {
-			logger.Printf("[Azure CNS] No endpoint state found, skipping initializing CNS state")
-			if cnsconfig.StateMigration {
-				logger.Printf("StatelessCNI Migration is enabled")
-				logger.Printf("initializing from Statefull CNI")
-				var endpointState map[string]*restserver.EndpointInfo
-				podInfoByIPProvider, endpointState, err = cnireconciler.NewCNIPodInfoProvider(cnsconfig.StateMigration)
-				if err != nil {
-					return nil, errors.Wrap(err, "failed to create CNI PodInfoProvider")
-				}
-				err = endpointStateStore.Write(restserver.EndpointStoreKey, endpointState)
-				if err != nil {
-					return nil, fmt.Errorf("failed to write endpoint state to store: %w", err)
-				}
-			}
-		} else {
-			return nil, errors.Wrap(err, "failed to create CNS PodInfoProvider")
-		}
+		return nil, errors.Wrap(err, "failed to create CNS PodInfoProvider")
 	}
 	return podInfoByIPProvider, nil
 }

network/endpoint.go

@@ -354,16 +354,6 @@ func GetPodNameWithoutSuffix(podName string) string {
 	return strings.Join(nameSplit, "-")
 }
 
-// GetEndpointInfoByIP returns an EndpointInfo with complete state when HNS Endpoint ID or HostVeth Name are missing on the CNS state.
-func (epInfo *EndpointInfo) GetEndpointInfoByIP(ipAddresses []net.IPNet, networkID string) (*EndpointInfo, error) {
-	// Call the platform implementation.
-	endpointInfo, err := epInfo.GetEndpointInfoByIPImpl(ipAddresses, networkID)
-	if err != nil {
-		return nil, err
-	}
-	return endpointInfo, nil
-}
-
 // IsEndpointStateInComplete returns true if both HNSEndpointID and HostVethName are missing.
 func (epInfo *EndpointInfo) IsEndpointStateIncomplete() bool {
 	if epInfo.HNSEndpointID == "" && epInfo.IfName == "" {

network/manager.go

@@ -444,7 +444,7 @@ func (nm *networkManager) GetEndpointState(networkID, endpointID string) (*Endpo
 
 	}
 	if epInfo.IsEndpointStateIncomplete() {
-		epInfo, err = epInfo.GetEndpointInfoByIP(epInfo.IPAddresses, networkID)
+		epInfo, err = epInfo.GetEndpointInfoByIPImpl(epInfo.IPAddresses, networkID)
 		if err != nil {
 			return nil, errors.Wrapf(err, "Get endpoint API returend with error")
 		}