fix: ctrl^C not terminate

free5gc · Oct 17, 2024 · 560775a · 560775a
1 parent f82f577
commit 560775a
Show file tree

Hide file tree

Showing 3 changed files with 14 additions and 6 deletions.
diff --git a/internal/sbi/consumer/nrf_service.go b/internal/sbi/consumer/nrf_service.go
@@ -144,8 +144,15 @@ func (s *nnrfService) RegisterNFInstance(ctx context.Context) (
 	var registerNfInstanceRequest Nnrf_NFManagement.RegisterNFInstanceRequest
 	registerNfInstanceRequest.NfInstanceID = &udmContext.NfId
 	registerNfInstanceRequest.NrfNfManagementNfProfile = &nfProfile
+	var res *Nnrf_NFManagement.RegisterNFInstanceResponse
 	for {
-		var res *Nnrf_NFManagement.RegisterNFInstanceResponse
+
+		select {
+		case <-ctx.Done():
+			return "", "", errors.Errorf("Context Cancel before RegisterNFInstance")
+		default:
+		}
+
 		res, err = client.NFInstanceIDDocumentApi.RegisterNFInstance(ctx, &registerNfInstanceRequest)
 
 		if err != nil || res == nil {

diff --git a/internal/sbi/server.go b/internal/sbi/server.go
@@ -28,6 +28,7 @@ type ServerUdm interface {
 
 	Consumer() *consumer.Consumer
 	Processor() *processor.Processor
+	CancelContext() context.Context
 }
 
 type Server struct {
@@ -60,7 +61,7 @@ func (s *Server) Run(traceCtx context.Context, wg *sync.WaitGroup) error {
 	logger.SBILog.Info("Starting server...")
 
 	var err error
-	_, s.Context().NfId, err = s.Consumer().RegisterNFInstance(context.Background())
+	_, s.Context().NfId, err = s.Consumer().RegisterNFInstance(s.CancelContext())
 	if err != nil {
 		logger.InitLog.Errorf("UDM register to NRF Error[%s]", err.Error())
 	}

diff --git a/pkg/service/init.go b/pkg/service/init.go
@@ -175,6 +175,10 @@ func (a *UdmApp) WaitRoutineStopped() {
 	logger.MainLog.Infof("UDM App is terminated")
 }
 
+func (a *UdmApp) CancelContext() context.Context {
+	return a.ctx
+}
+
 func (a *UdmApp) Config() *factory.Config {
 	return a.cfg
 }
@@ -183,10 +187,6 @@ func (a *UdmApp) Context() *udm_context.UDMContext {
 	return a.udmCtx
 }
 
-func (a *UdmApp) CancelContext() context.Context {
-	return a.ctx
-}
-
 func (a *UdmApp) Consumer() *consumer.Consumer {
 	return a.consumer
 }