feat: register procedure graceful shutdown

free5gc · May 9, 2024 · 4488ffb · 4488ffb
1 parent 8a74e48
commit 4488ffb
Show file tree

Hide file tree

Showing 2 changed files with 45 additions and 38 deletions.
diff --git a/internal/sbi/consumer/nrf_service.go b/internal/sbi/consumer/nrf_service.go
@@ -40,7 +40,7 @@ func (ns *NrfService) buildNFProfile(context *nssf_context.NSSFContext) (profile
 	return
 }
 
-func (ns *NrfService) SendRegisterNFInstance(nssfCtx *nssf_context.NSSFContext) (
+func (ns *NrfService) SendRegisterNFInstance(ctx context.Context, nssfCtx *nssf_context.NSSFContext) (
 	resourceNrfUri string, retrieveNfInstanceId string, err error,
 ) {
 	nfInstanceId := nssfCtx.NfId
@@ -52,44 +52,51 @@ func (ns *NrfService) SendRegisterNFInstance(nssfCtx *nssf_context.NSSFContext)
 
 	var res *http.Response
 	var nf models.NfProfile
-	for {
-		nf, res, err = apiClient.NFInstanceIDDocumentApi.RegisterNFInstance(context.TODO(), nfInstanceId, profile)
-		if err != nil || res == nil {
-			// TODO : add log
-			logger.ConsumerLog.Errorf("NSSF register to NRF Error[%s]", err.Error())
-			time.Sleep(2 * time.Second)
-			continue
-		}
-		defer func() {
-			if resCloseErr := res.Body.Close(); resCloseErr != nil {
-				logger.ConsumerLog.Errorf("NFInstanceIDDocumentApi response body cannot close: %+v", resCloseErr)
+	finish := false
+	for !finish {
+		select {
+		case <-ctx.Done():
+			return "", "", fmt.Errorf("context done")
+
+		default:
+			nf, res, err = apiClient.NFInstanceIDDocumentApi.RegisterNFInstance(ctx, nfInstanceId, profile)
+			if err != nil || res == nil {
+				// TODO : add log
+				logger.ConsumerLog.Errorf("NSSF register to NRF Error[%s]", err.Error())
+				time.Sleep(2 * time.Second)
+				continue
 			}
-		}()
-		status := res.StatusCode
-		if status == http.StatusOK {
-			// NFUpdate
-			break
-		} else if status == http.StatusCreated {
-			// NFRegister
-			resourceUri := res.Header.Get("Location")
-			resourceNrfUri = resourceUri[:strings.Index(resourceUri, "/nnrf-nfm/")]
-			retrieveNfInstanceId = resourceUri[strings.LastIndex(resourceUri, "/")+1:]
+			defer func() {
+				if resCloseErr := res.Body.Close(); resCloseErr != nil {
+					logger.ConsumerLog.Errorf("NFInstanceIDDocumentApi response body cannot close: %+v", resCloseErr)
+				}
+			}()
+			status := res.StatusCode
+			if status == http.StatusOK {
+				// NFUpdate
+				finish = true
+			} else if status == http.StatusCreated {
+				// NFRegister
+				resourceUri := res.Header.Get("Location")
+				resourceNrfUri = resourceUri[:strings.Index(resourceUri, "/nnrf-nfm/")]
+				retrieveNfInstanceId = 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)
+				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)
+					}
 				}
+				nssf_context.GetSelf().OAuth2Required = oauth2
+				if oauth2 && nssf_context.GetSelf().NrfCertPem == "" {
+					logger.CfgLog.Error("OAuth2 enable but no nrfCertPem provided in config.")
+				}
+				finish = true
+			} else {
+				fmt.Println("NRF return wrong status code", status)
 			}
-			nssf_context.GetSelf().OAuth2Required = oauth2
-			if oauth2 && nssf_context.GetSelf().NrfCertPem == "" {
-				logger.CfgLog.Error("OAuth2 enable but no nrfCertPem provided in config.")
-			}
-			break
-		} else {
-			fmt.Println("NRF return wrong status code", status)
 		}
 	}
 	return resourceNrfUri, retrieveNfInstanceId, err

diff --git a/pkg/service/init.go b/pkg/service/init.go
@@ -115,11 +115,11 @@ func (a *NssfApp) SetReportCaller(reportCaller bool) {
 	logger.Log.SetReportCaller(reportCaller)
 }
 
-func (a *NssfApp) registerToNrf() error {
+func (a *NssfApp) registerToNrf(ctx context.Context) error {
 	nssfContext := a.nssfCtx
 
 	var err error
-	_, nssfContext.NfId, err = a.consumer.SendRegisterNFInstance(nssfContext)
+	_, nssfContext.NfId, err = a.consumer.SendRegisterNFInstance(ctx, nssfContext)
 	if err != nil {
 		return fmt.Errorf("failed to register NSSF to NRF: %s", err.Error())
 	}
@@ -148,7 +148,7 @@ func (a *NssfApp) Start() {
 		cancel() // Notify each goroutine and wait them stopped
 	}()
 
-	err := a.registerToNrf()
+	err := a.registerToNrf(ctx)
 	if err != nil {
 		logger.MainLog.Errorf("register to NRF failed: %+v", err)
 	} else {