Fix stale AMF CommunicationClient causing N1N2 transfer failures

After an AMF pod restart, SMF continues sending N1N2MessageTransfer
to the old AMF pod IP, causing PDU Session Establishment to fail
with T3580 expiry. This happens because:

1. CommunicationClient bakes in the AMF pod IP at session creation
   and is never refreshed.
2. CommunicationClient is not JSON-serializable, so it's nil after
   SMF recovers SMContext from MongoDB.
3. NRF accumulates stale NF registrations across pod restarts and
   the NRF cache (15min TTL) returns stale entries on re-discovery.
4. The HTTP client has no timeout, so TCP connect to a dead pod IP
   hangs for 60s+ — longer than the T3580 window.

Fix:

- Add SendN1N2TransferWithRediscovery() that wraps all N1N2 calls
  with a 5-second context timeout. On failure, it re-discovers the
  AMF by querying NRF directly (bypassing cache), prefers an AMF
  with a different NfInstanceId than the failed one, and retries.

- Add RebuildCommunicationClient() on SMContext to reconstruct the
  HTTP client from the stored AMFProfile after DB recovery.

- Replace all 5 direct CommunicationClient.N1N2MessageTransfer call
  sites with the retry wrapper.

Verified on live cluster: after AMF pod kill, N1N2 transfer fails
in 5s, re-discovers new AMF, retries successfully. Total recovery
time 5.3 seconds vs permanent failure without the fix.

Signed-off-by: Vinod Patmanathan <vinod.patmanathan@forsway.com>

Author: Vinod Patmanathan

consumer/nf_management.go

@@ -478,6 +478,134 @@ func SendNFDiscoveryServingAMF(smContext *smfContext.SMContext) (*models.Problem
 	return nil, nil
 }
 
+// n1n2TransferTimeout is the HTTP timeout for a single N1N2MessageTransfer attempt.
+// Kept short so that retries with AMF re-discovery can happen within the UE's T3580 window.
+const n1n2TransferTimeout = 5 * time.Second
+
+// SendN1N2TransferWithRediscovery sends an N1N2MessageTransfer request using the
+// SMContext's CommunicationClient. If the request fails (including timeout), it
+// queries NRF directly (bypassing the cache) and tries every other AMF candidate
+// until one succeeds. This handles the case where NRF has multiple stale
+// registrations left behind by prior AMF pod restarts.
+func SendN1N2TransferWithRediscovery(ctx context.Context, smContext *smfContext.SMContext,
+	n1n2Request models.N1N2MessageTransferRequest,
+) (*models.N1N2MessageTransferRspData, error) {
+	if smContext.CommunicationClient == nil {
+		// Client was never built (e.g. recovered from DB with stale AMFProfile).
+		// Use the first candidate from NRF to seed it; if it fails, the retry loop below takes over.
+		if err := selectFirstAmfFromNrf(smContext); err != nil {
+			return nil, fmt.Errorf("AMF discovery failed: %w", err)
+		}
+	}
+
+	// First attempt with the currently-selected AMF
+	rspData, err := tryN1N2Transfer(ctx, smContext, n1n2Request)
+	if err == nil {
+		return rspData, nil
+	}
+	firstErr := err
+	smContext.SubPduSessLog.Warnf("N1N2Transfer failed (%v), attempting AMF re-discovery", err)
+
+	// First attempt failed — fetch all AMF candidates from NRF (bypassing cache)
+	// and try each one that isn't the NfInstanceId we already failed on.
+	candidates, discErr := fetchAmfCandidates()
+	if discErr != nil {
+		return nil, fmt.Errorf("N1N2Transfer failed (%w) and AMF re-discovery failed: %v", firstErr, discErr)
+	}
+
+	attempted := 0
+	for _, candidate := range selectableAmfCandidates(candidates, smContext.ServingNfId) {
+		if err := useAmfProfile(smContext, candidate); err != nil {
+			smContext.SubPduSessLog.Warnf("AMF candidate %s unusable: %v", candidate.NfInstanceId, err)
+			continue
+		}
+		attempted++
+		smContext.SubPduSessLog.Infof("AMF re-discovery retry %d: trying NfInstanceId %s", attempted, candidate.NfInstanceId)
+		rspData, err = tryN1N2Transfer(ctx, smContext, n1n2Request)
+		if err == nil {
+			smContext.SubPduSessLog.Infof("AMF re-discovery succeeded on attempt %d with NfInstanceId %s", attempted, candidate.NfInstanceId)
+			return rspData, nil
+		}
+		smContext.SubPduSessLog.Warnf("AMF re-discovery retry %d failed: %v", attempted, err)
+	}
+
+	if attempted == 0 {
+		return nil, fmt.Errorf("N1N2Transfer failed (%w) and no alternative AMF candidates available", firstErr)
+	}
+	return nil, fmt.Errorf("N1N2Transfer failed after %d AMF candidates; last error: %w", attempted, err)
+}
+
+// tryN1N2Transfer makes a single N1N2MessageTransfer call with a short timeout.
+func tryN1N2Transfer(ctx context.Context, smContext *smfContext.SMContext,
+	n1n2Request models.N1N2MessageTransferRequest,
+) (*models.N1N2MessageTransferRspData, error) {
+	tryCtx, cancel := context.WithTimeout(ctx, n1n2TransferTimeout)
+	defer cancel()
+	apiReq := smContext.CommunicationClient.
+		N1N2MessageCollectionCollectionAPI.
+		N1N2MessageTransfer(tryCtx, smContext.Supi).
+		N1N2MessageTransferReqData(n1n2Request.GetJsonData())
+	if binaryDataN1Message := n1n2Request.GetBinaryDataN1Message(); binaryDataN1Message != nil {
+		apiReq = apiReq.BinaryDataN1Message(binaryDataN1Message)
+	}
+	if binaryDataN2Information := n1n2Request.GetBinaryDataN2Information(); binaryDataN2Information != nil {
+		apiReq = apiReq.BinaryDataN2Information(binaryDataN2Information)
+	}
+	rspData, _, err := smContext.CommunicationClient.
+		N1N2MessageCollectionCollectionAPI.
+		N1N2MessageTransferExecute(apiReq)
+	return rspData, err
+}
+
+// selectableAmfCandidates returns candidates whose NfInstanceId differs from failedNfId,
+// preserving order. Used to filter out the AMF profile that just failed an N1N2 transfer.
+func selectableAmfCandidates(candidates []models.NFProfileDiscovery, failedNfId string) []models.NFProfileDiscovery {
+	out := make([]models.NFProfileDiscovery, 0, len(candidates))
+	for _, c := range candidates {
+		if c.NfInstanceId == failedNfId {
+			continue
+		}
+		out = append(out, c)
+	}
+	return out
+}
+
+// fetchAmfCandidates queries NRF directly (bypassing the cache) for all registered AMFs.
+func fetchAmfCandidates() ([]models.NFProfileDiscovery, error) {
+	localVarOptionals := Nnrf_NFDiscovery.ApiSearchNFInstancesRequest{}
+	ctx := context.Background()
+	result, err := SendNrfForNfInstance(ctx, smfContext.SMF_Self().NrfUri, models.NFTYPE_AMF, models.NFTYPE_SMF, localVarOptionals)
+	if err != nil {
+		return nil, fmt.Errorf("broad AMF discovery failed: %w", err)
+	}
+	if result == nil || result.NfInstances == nil || len(result.NfInstances) == 0 {
+		return nil, fmt.Errorf("broad AMF discovery returned no AMF instances")
+	}
+	return result.NfInstances, nil
+}
+
+// useAmfProfile selects the given AMF profile on the SMContext and rebuilds the CommunicationClient.
+func useAmfProfile(smContext *smfContext.SMContext, profile models.NFProfileDiscovery) error {
+	smContext.AMFProfile = deepcopy.Copy(profile).(models.NFProfileDiscovery)
+	smContext.ServingNfId = smContext.AMFProfile.NfInstanceId
+	smContext.RebuildCommunicationClient()
+	if smContext.CommunicationClient == nil {
+		return fmt.Errorf("AMF profile %s has no Namf_Communication service", profile.NfInstanceId)
+	}
+	return nil
+}
+
+// selectFirstAmfFromNrf picks the first AMF from a direct NRF query and installs it on the SMContext.
+// Used only to bootstrap a nil CommunicationClient after DB recovery.
+func selectFirstAmfFromNrf(smContext *smfContext.SMContext) error {
+	smContext.SubPduSessLog.Infof("AMF discovery: querying NRF directly (bypassing cache)")
+	candidates, err := fetchAmfCandidates()
+	if err != nil {
+		return err
+	}
+	return useAmfProfile(smContext, candidates[0])
+}
+
 func SendCreateSubscription(nrfUri string, nrfSubscriptionData models.SubscriptionData) (nrfSubData *models.SubscriptionData, problemDetails *models.ProblemDetails, err error) {
 	logger.ConsumerLog.Debugln("send Create Subscription")
 

consumer/nf_management_amf_rediscover_test.go

@@ -0,0 +1,68 @@
+// SPDX-FileCopyrightText: 2026 Forsway Solutions AB
+// SPDX-License-Identifier: Apache-2.0
+
+package consumer
+
+import (
+	"testing"
+
+	"github.com/omec-project/openapi/v2/models"
+)
+
+func amfDiscoveryProfile(nfInstanceId, apiPrefix string) models.NFProfileDiscovery {
+	return models.NFProfileDiscovery{
+		NfInstanceId: nfInstanceId,
+		NfType:       models.NFTYPE_AMF,
+		NfStatus:     models.NFSTATUS_REGISTERED,
+		NfServices: []models.NFService{
+			{ServiceName: models.SERVICENAME_NAMF_COMM, ApiPrefix: &apiPrefix},
+		},
+	}
+}
+
+func TestSelectableAmfCandidates_SkipsFailedNfInstanceId(t *testing.T) {
+	a := amfDiscoveryProfile("amf-a", "http://10.42.0.10:29518")
+	b := amfDiscoveryProfile("amf-b", "http://10.42.0.20:29518")
+	c := amfDiscoveryProfile("amf-c", "http://10.42.0.30:29518")
+
+	got := selectableAmfCandidates([]models.NFProfileDiscovery{a, b, c}, "amf-b")
+	if len(got) != 2 {
+		t.Fatalf("expected 2 candidates after filtering, got %d", len(got))
+	}
+	if got[0].NfInstanceId != "amf-a" || got[1].NfInstanceId != "amf-c" {
+		t.Errorf("expected order [amf-a, amf-c], got [%s, %s]", got[0].NfInstanceId, got[1].NfInstanceId)
+	}
+}
+
+func TestSelectableAmfCandidates_PreservesOrder(t *testing.T) {
+	a := amfDiscoveryProfile("amf-a", "http://10.42.0.10:29518")
+	b := amfDiscoveryProfile("amf-b", "http://10.42.0.20:29518")
+	c := amfDiscoveryProfile("amf-c", "http://10.42.0.30:29518")
+
+	got := selectableAmfCandidates([]models.NFProfileDiscovery{a, b, c}, "does-not-match-any")
+	if len(got) != 3 {
+		t.Fatalf("expected 3 candidates when no match, got %d", len(got))
+	}
+	for i, want := range []string{"amf-a", "amf-b", "amf-c"} {
+		if got[i].NfInstanceId != want {
+			t.Errorf("position %d: got %s, want %s", i, got[i].NfInstanceId, want)
+		}
+	}
+}
+
+func TestSelectableAmfCandidates_AllFiltered_ReturnsEmpty(t *testing.T) {
+	a := amfDiscoveryProfile("amf-x", "http://10.42.0.10:29518")
+	b := amfDiscoveryProfile("amf-x", "http://10.42.0.20:29518")
+
+	got := selectableAmfCandidates([]models.NFProfileDiscovery{a, b}, "amf-x")
+	if len(got) != 0 {
+		t.Errorf("expected empty result when every candidate matches failed NfId, got %d", len(got))
+	}
+}
+
+func TestSelectableAmfCandidates_EmptyInput_ReturnsEmpty(t *testing.T) {
+	got := selectableAmfCandidates(nil, "anything")
+	if len(got) != 0 {
+		t.Errorf("expected empty result for nil input, got %d", len(got))
+	}
+}

context/db.go

@@ -323,6 +323,7 @@ func GetSMContextByRefInDB(ref string) (smContext *SMContext) {
 			logger.DataRepoLog.Errorf("smContext unmarshall error: %v", err)
 			return nil
 		}
+		smContext.RebuildCommunicationClient()
 	} else {
 		logger.DataRepoLog.Warnf("SmContext doesn't exist with ref: %v", ref)
 		return nil

context/sm_context.go

@@ -372,6 +372,27 @@ func (smContext *SMContext) SetCreateData(createData *models.SmContextCreateData
 	smContext.ServingNfId = createData.GetServingNfId()
 }
 
+// RebuildCommunicationClient reconstructs the Namf_Communication API client
+// from the stored AMFProfile. This is needed after recovering an SMContext from
+// MongoDB, since CommunicationClient is not serializable.
+func (smContext *SMContext) RebuildCommunicationClient() {
+	if smContext.AMFProfile.NfServices == nil {
+		return
+	}
+	for _, service := range smContext.AMFProfile.NfServices {
+		if service.ServiceName == models.SERVICENAME_NAMF_COMM {
+			communicationConf := Namf_Communication.NewConfiguration()
+			serverConfig := &communicationConf.Servers[0]
+			if apiRootVar, exists := serverConfig.Variables["apiRoot"]; exists {
+				apiRootVar.DefaultValue = service.GetApiPrefix()
+				serverConfig.Variables["apiRoot"] = apiRootVar
+			}
+			smContext.CommunicationClient = Namf_Communication.NewAPIClient(communicationConf)
+			return
+		}
+	}
+}
+
 func (smContext *SMContext) BuildCreatedData() (createdData *models.SmContextCreatedData) {
 	createdData = models.NewSmContextCreatedData()
 	createdData.SNssai = smContext.Snssai

context/sm_context_rebuild_test.go

@@ -0,0 +1,107 @@
+// SPDX-FileCopyrightText: 2026 Forsway Solutions AB
+// SPDX-License-Identifier: Apache-2.0
+
+package context
+
+import (
+	"strings"
+	"testing"
+
+	"github.com/omec-project/openapi/v2/models"
+)
+
+func amfProfileWithNamfComm(apiPrefix string) models.NFProfileDiscovery {
+	commService := models.NFService{
+		ServiceName: models.SERVICENAME_NAMF_COMM,
+		ApiPrefix:   &apiPrefix,
+	}
+	return models.NFProfileDiscovery{
+		NfInstanceId: "amf-instance-1",
+		NfType:       models.NFTYPE_AMF,
+		NfStatus:     models.NFSTATUS_REGISTERED,
+		NfServices:   []models.NFService{commService},
+	}
+}
+
+func TestRebuildCommunicationClient_WithNamfComm_BuildsClientAndSetsApiRoot(t *testing.T) {
+	const apiPrefix = "http://10.42.0.42:29518"
+
+	smCtx := &SMContext{}
+	smCtx.AMFProfile = amfProfileWithNamfComm(apiPrefix)
+
+	smCtx.RebuildCommunicationClient()
+
+	if smCtx.CommunicationClient == nil {
+		t.Fatalf("expected CommunicationClient to be non-nil after RebuildCommunicationClient")
+	}
+	// Ensure the apiRoot variable was wired through to the configuration.
+	cfg := smCtx.CommunicationClient.GetConfig()
+	if cfg == nil || len(cfg.Servers) == 0 {
+		t.Fatalf("expected client configuration with at least one server")
+	}
+	apiRootVar, ok := cfg.Servers[0].Variables["apiRoot"]
+	if !ok {
+		t.Fatalf("expected apiRoot variable to be present on server[0]")
+	}
+	if apiRootVar.DefaultValue != apiPrefix {
+		t.Errorf("apiRoot DefaultValue = %q, want %q", apiRootVar.DefaultValue, apiPrefix)
+	}
+}
+
+func TestRebuildCommunicationClient_NoServices_LeavesClientNil(t *testing.T) {
+	smCtx := &SMContext{}
+	// AMFProfile.NfServices is nil here.
+
+	smCtx.RebuildCommunicationClient()
+
+	if smCtx.CommunicationClient != nil {
+		t.Errorf("expected CommunicationClient to remain nil when NfServices is nil")
+	}
+}
+
+func TestRebuildCommunicationClient_NoNamfComm_LeavesClientNil(t *testing.T) {
+	apiPrefix := "http://10.42.0.42:29518"
+	smCtx := &SMContext{}
+	smCtx.AMFProfile = models.NFProfileDiscovery{
+		NfInstanceId: "amf-instance-2",
+		NfType:       models.NFTYPE_AMF,
+		NfStatus:     models.NFSTATUS_REGISTERED,
+		NfServices: []models.NFService{
+			{ServiceName: models.SERVICENAME_NAMF_EVTS, ApiPrefix: &apiPrefix},
+		},
+	}
+
+	smCtx.RebuildCommunicationClient()
+
+	if smCtx.CommunicationClient != nil {
+		t.Errorf("expected CommunicationClient to remain nil when no namf-comm service is present")
+	}
+}
+
+// Belt-and-braces: a fresh build replaces any previously held client.
+func TestRebuildCommunicationClient_RebuildReplacesExistingClient(t *testing.T) {
+	const firstApiPrefix = "http://10.42.0.10:29518"
+	const secondApiPrefix = "http://10.42.0.20:29518"
+
+	smCtx := &SMContext{}
+	smCtx.AMFProfile = amfProfileWithNamfComm(firstApiPrefix)
+	smCtx.RebuildCommunicationClient()
+	if smCtx.CommunicationClient == nil {
+		t.Fatalf("setup: expected non-nil CommunicationClient after first build")
+	}
+	first := smCtx.CommunicationClient
+
+	smCtx.AMFProfile = amfProfileWithNamfComm(secondApiPrefix)
+	smCtx.RebuildCommunicationClient()
+
+	if smCtx.CommunicationClient == nil {
+		t.Fatalf("expected non-nil CommunicationClient after second build")
+	}
+	if smCtx.CommunicationClient == first {
+		t.Errorf("expected a fresh CommunicationClient instance after rebuild, got the same pointer")
+	}
+	apiRootVar := smCtx.CommunicationClient.GetConfig().Servers[0].Variables["apiRoot"]
+	if !strings.Contains(apiRootVar.DefaultValue, "10.42.0.20") {
+		t.Errorf("apiRoot DefaultValue = %q, expected it to reflect the second AMF endpoint", apiRootVar.DefaultValue)
+	}
+}

pfcp/handler/handler.go

@@ -13,6 +13,7 @@ import (
 
 	"github.com/omec-project/openapi/v2"
 	"github.com/omec-project/openapi/v2/models"
+	"github.com/omec-project/smf/consumer"
 	smf_context "github.com/omec-project/smf/context"
 	"github.com/omec-project/smf/factory"
 	"github.com/omec-project/smf/logger"
@@ -747,16 +748,7 @@ func HandlePfcpSessionReportRequest(msg *udp.Message) {
 					},
 				}
 
-				apiN1N2MessageTransferRequest := smContext.CommunicationClient.N1N2MessageCollectionCollectionAPI.N1N2MessageTransfer(context.Background(), smContext.Supi)
-				apiN1N2MessageTransferRequest = apiN1N2MessageTransferRequest.N1N2MessageTransferReqData(n1n2Request.GetJsonData())
-				if binaryDataN1Message := n1n2Request.GetBinaryDataN1Message(); binaryDataN1Message != nil {
-					apiN1N2MessageTransferRequest = apiN1N2MessageTransferRequest.BinaryDataN1Message(binaryDataN1Message)
-				}
-				if binaryDataN2Information := n1n2Request.GetBinaryDataN2Information(); binaryDataN2Information != nil {
-					apiN1N2MessageTransferRequest = apiN1N2MessageTransferRequest.BinaryDataN2Information(binaryDataN2Information)
-				}
-				var rspData *models.N1N2MessageTransferRspData
-				rspData, _, err = smContext.CommunicationClient.N1N2MessageCollectionCollectionAPI.N1N2MessageTransferExecute(apiN1N2MessageTransferRequest)
+				rspData, err := consumer.SendN1N2TransferWithRediscovery(context.Background(), smContext, n1n2Request)
 				if err != nil {
 					smContext.SubPfcpLog.Warnf("Send N1N2Transfer failed")
 				}