From c66c41ba203e7ed04b32227f0b7ae16ced330506 Mon Sep 17 00:00:00 2001
From: pf-lin <welin9211@gmail.com>
Date: Tue, 22 Oct 2024 11:17:34 +0000
Subject: [PATCH] fix: unable to terminate when registration with NRF is
 unsuccessful

---
 internal/sbi/consumer/nrf_service.go | 68 +++++++++++++++-------------
 internal/sbi/server.go               |  9 ----
 pkg/service/init.go                  |  8 ++++
 3 files changed, 45 insertions(+), 40 deletions(-)

diff --git a/internal/sbi/consumer/nrf_service.go b/internal/sbi/consumer/nrf_service.go
index 3aa24e07..bc7ade77 100644
--- a/internal/sbi/consumer/nrf_service.go
+++ b/internal/sbi/consumer/nrf_service.go
@@ -88,37 +88,43 @@ func (s *nnrfService) RegisterNFInstance(ctx context.Context) error {
 	}
 
 	// Check data (Use RESTful PUT)
-	for {
-		res, err = client.NFInstanceIDDocumentApi.RegisterNFInstance(ctx, registerNFInstanceRequest)
-		if err != nil || res == nil {
-			logger.ConsumerLog.Infof("SMF register to NRF Error[%s]", err.Error())
-			time.Sleep(2 * time.Second)
-			continue
-		}
-		nf = res.NrfNfManagementNfProfile
-
-		// http.StatusOK
-		if res.Location == "" {
-			// NFUpdate
-			break
-		} else { // http.StatusCreated
-			// NFRegister
-			resourceUri := res.Location
-			smfContext.NfInstanceID = resourceUri[strings.LastIndex(resourceUri, "/")+1:]
-
-			oauth2 := false
-			if nf.CustomInfo != nil {
-				v, ok := nf.CustomInfo["oauth2"].(bool)
-				if ok {
-					oauth2 = v
-					logger.MainLog.Infoln("OAuth2 setting receive from NRF:", oauth2)
-				}
+	finish := false
+	for !finish {
+		select {
+		case <-ctx.Done():
+			return fmt.Errorf("RegisterNFInstance context done")
+		default:
+			res, err = client.NFInstanceIDDocumentApi.RegisterNFInstance(ctx, registerNFInstanceRequest)
+			if err != nil || res == nil {
+				logger.ConsumerLog.Infof("SMF register to NRF Error[%s]", err.Error())
+				time.Sleep(2 * time.Second)
+				continue
 			}
-			smfContext.OAuth2Required = oauth2
-			if oauth2 && smfContext.NrfCertPem == "" {
-				logger.CfgLog.Error("OAuth2 enable but no nrfCertPem provided in config.")
+			nf = res.NrfNfManagementNfProfile
+
+			// http.StatusOK
+			if res.Location == "" {
+				// NFUpdate
+				finish = true
+			} else { // http.StatusCreated
+				// NFRegister
+				resourceUri := res.Location
+				smfContext.NfInstanceID = resourceUri[strings.LastIndex(resourceUri, "/")+1:]
+
+				oauth2 := false
+				if nf.CustomInfo != nil {
+					v, ok := nf.CustomInfo["oauth2"].(bool)
+					if ok {
+						oauth2 = v
+						logger.MainLog.Infoln("OAuth2 setting receive from NRF:", oauth2)
+					}
+				}
+				smfContext.OAuth2Required = oauth2
+				if oauth2 && smfContext.NrfCertPem == "" {
+					logger.CfgLog.Error("OAuth2 enable but no nrfCertPem provided in config.")
+				}
+				finish = true
 			}
-			break
 		}
 	}
 
@@ -153,10 +159,10 @@ func (s *nnrfService) buildNfProfile(smfContext *smf_context.SMFContext) (
 	return profile, err
 }
 
-func (s *nnrfService) RetrySendNFRegistration(maxRetry int) error {
+func (s *nnrfService) RetrySendNFRegistration(ctx context.Context, maxRetry int) error {
 	retryCount := 0
 	for retryCount < maxRetry {
-		err := s.RegisterNFInstance(context.Background())
+		err := s.RegisterNFInstance(ctx)
 		if err == nil {
 			return nil
 		}
diff --git a/internal/sbi/server.go b/internal/sbi/server.go
index 04f3371d..95b946d6 100644
--- a/internal/sbi/server.go
+++ b/internal/sbi/server.go
@@ -107,15 +107,6 @@ func newRouter(s *Server) *gin.Engine {
 }
 
 func (s *Server) Run(traceCtx context.Context, wg *sync.WaitGroup) error {
-	err := s.Consumer().RegisterNFInstance(context.Background())
-	if err != nil {
-		retry_err := s.Consumer().RetrySendNFRegistration(10)
-		if retry_err != nil {
-			logger.InitLog.Errorln(retry_err)
-			return err
-		}
-	}
-
 	wg.Add(1)
 	go s.startServer(wg)
 
diff --git a/pkg/service/init.go b/pkg/service/init.go
index 82fe2cbe..76f6797f 100644
--- a/pkg/service/init.go
+++ b/pkg/service/init.go
@@ -172,6 +172,14 @@ func (a *SmfApp) Start() {
 	a.wg.Add(1)
 	go a.listenShutDownEvent()
 
+	err = a.Consumer().RegisterNFInstance(a.ctx)
+	if err != nil {
+		retry_err := a.Consumer().RetrySendNFRegistration(a.ctx, 10)
+		if retry_err != nil {
+			logger.InitLog.Errorln(retry_err)
+		}
+	}
+
 	// Initialize PFCP server
 	a.pfcpStart(a)