From 560775abf45124fbee344d141108a23d59535d54 Mon Sep 17 00:00:00 2001
From: ubuntu <hjkja0511@gmail.com>
Date: Thu, 17 Oct 2024 09:32:33 +0000
Subject: [PATCH] fix: ctrl^C not terminate

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

diff --git a/internal/sbi/consumer/nrf_service.go b/internal/sbi/consumer/nrf_service.go
index 8648a0b..fe0ba23 100644
--- 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
index 4f7a3eb..fc17957 100644
--- 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
index add4c81..59e96a5 100644
--- 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
 }