Merge pull request #18 from andy89923/fix/deregister

Fix: Graceful shutdown
free5gc · Mar 12, 2024 · ab9baa7 · ab9baa7
2 parents eb51a94 + edc1155
commit ab9baa7
Show file tree

Hide file tree

Showing 4 changed files with 32 additions and 24 deletions.
diff --git a/internal/cgf/cgf.go b/internal/cgf/cgf.go
@@ -3,14 +3,13 @@ package cgf
 
 import (
 	"bytes"
+	"context"
 	"encoding/json"
 	"os"
-	"os/signal"
 	"path/filepath"
 	"runtime/debug"
 	"strconv"
 	"sync"
-	"syscall"
 	"time"
 
 	"github.com/fclairamb/ftpserver/config"
@@ -46,7 +45,7 @@ type FtpConfig struct {
 
 var cgf *Cgf
 
-func OpenServer(wg *sync.WaitGroup) *Cgf {
+func OpenServer(ctx context.Context, wg *sync.WaitGroup) *Cgf {
 	// Arguments vars
 	cgf = new(Cgf)
 
@@ -107,7 +106,7 @@ func OpenServer(wg *sync.WaitGroup) *Cgf {
 	// Setting up the ftpserver logger
 	cgf.ftpServer.Logger = logger.FtpServerLog
 
-	go cgf.Serve(wg)
+	go cgf.Serve(ctx, wg)
 	logger.CgfLog.Info("FTP server Start")
 
 	return cgf
@@ -184,22 +183,17 @@ func SendCDR(supi string) error {
 const FTP_LOGIN_RETRY_NUMBER = 3
 const FTP_LOGIN_RETRY_WAITING_TIME = 1 * time.Second // second
 
-func (f *Cgf) Serve(wg *sync.WaitGroup) {
-	signalChannel := make(chan os.Signal, 1)
-	signal.Notify(signalChannel, os.Interrupt, syscall.SIGTERM)
-
+func (f *Cgf) Serve(ctx context.Context, wg *sync.WaitGroup) {
 	go func() {
 		defer func() {
 			if p := recover(); p != nil {
 				// Print stack for panic to log. Fatalf() will let program exit.
 				logger.InitLog.Fatalf("panic: %v\n%s", p, string(debug.Stack()))
 			}
 		}()
-
-		<-signalChannel
+		<-ctx.Done()
 		f.Terminate()
 		wg.Done()
-		os.Exit(0)
 	}()
 
 	for i := 0; i < FTP_LOGIN_RETRY_NUMBER; i++ {
@@ -222,6 +216,8 @@ func (f *Cgf) Serve(wg *sync.WaitGroup) {
 }
 
 func (f *Cgf) Terminate() {
+	logger.CgfLog.Infoln("CGF Terminating...")
+
 	f.driver.Stop()
 
 	if err := f.ftpServer.Stop(); err != nil {
@@ -247,4 +243,5 @@ func (f *Cgf) Terminate() {
 			}
 		}
 	}
+	logger.CgfLog.Infoln("CGF terminated")
 }
diff --git a/pkg/abmf/abmf.go b/pkg/abmf/abmf.go
@@ -17,6 +17,7 @@ package abmf
 
 import (
 	"bytes"
+	"context"
 	"math"
 	"strconv"
 	"sync"
@@ -39,7 +40,7 @@ import (
 
 const chargingDatasColl = "policyData.ues.chargingData"
 
-func OpenServer(wg *sync.WaitGroup) {
+func OpenServer(ctx context.Context, wg *sync.WaitGroup) {
 	// Load our custom dictionary on top of the default one, which
 	// always have the Base Protocol (RFC6733) and Credit Control
 	// Application (RFC4006).
@@ -73,12 +74,14 @@ func OpenServer(wg *sync.WaitGroup) {
 	go printErrors(mux.ErrorReports())
 	go func() {
 		defer func() {
-			logger.AcctLog.Error("ABMF server stopped")
+			logger.AcctLog.Infoln("ABMF server stopped")
 			wg.Done()
 		}()
-
-		abmfDiameter := factory.ChfConfig.Configuration.AbmfDiameter
-		addr := abmfDiameter.HostIPv4 + ":" + strconv.Itoa(abmfDiameter.Port)
+		<-ctx.Done()
+	}()
+	abmfDiameter := factory.ChfConfig.Configuration.AbmfDiameter
+	addr := abmfDiameter.HostIPv4 + ":" + strconv.Itoa(abmfDiameter.Port)
+	go func() {
 		err := diam.ListenAndServeTLS(addr, abmfDiameter.Tls.Pem, abmfDiameter.Tls.Key, mux, nil)
 		if err != nil {
 			logger.AcctLog.Errorf("ABMF server fail to listen: %V", err)

diff --git a/pkg/rf/rating.go b/pkg/rf/rating.go
@@ -17,6 +17,7 @@ package rf
 
 import (
 	"bytes"
+	"context"
 	"log"
 	"math"
 	"strconv"
@@ -40,7 +41,7 @@ import (
 
 const chargingDatasColl = "policyData.ues.chargingData"
 
-func OpenServer(wg *sync.WaitGroup) {
+func OpenServer(ctx context.Context, wg *sync.WaitGroup) {
 	// Load our custom dictionary on top of the default one, which
 	// always have the Base Protocol (RFC6733) and Credit Control
 	// Application (RFC4006).
@@ -72,12 +73,14 @@ func OpenServer(wg *sync.WaitGroup) {
 	go printErrors(mux.ErrorReports())
 	go func() {
 		defer func() {
-			logger.CgfLog.Error("Rating Function server stopped")
+			logger.CgfLog.Infoln("Rating Function server stopped")
 			wg.Done()
 		}()
-
-		rfDiameter := factory.ChfConfig.Configuration.RfDiameter
-		addr := rfDiameter.HostIPv4 + ":" + strconv.Itoa(rfDiameter.Port)
+		<-ctx.Done()
+	}()
+	rfDiameter := factory.ChfConfig.Configuration.RfDiameter
+	addr := rfDiameter.HostIPv4 + ":" + strconv.Itoa(rfDiameter.Port)
+	go func() {
 		err := diam.ListenAndServeTLS(addr, rfDiameter.Tls.Pem, rfDiameter.Tls.Key, mux, nil)
 		if err != nil {
 			log.Fatal(err)

diff --git a/pkg/service/init.go b/pkg/service/init.go
@@ -1,6 +1,7 @@
 package service
 
 import (
+	"context"
 	"fmt"
 	"io/ioutil"
 	"os"
@@ -113,14 +114,16 @@ func (c *ChfApp) Start(tlsKeyLogPath string) {
 
 	wg := sync.WaitGroup{}
 
+	ctx, cancel := context.WithCancel(context.Background())
+
 	wg.Add(1)
-	cgf.OpenServer(&wg)
+	cgf.OpenServer(ctx, &wg)
 
 	wg.Add(1)
-	rf.OpenServer(&wg)
+	rf.OpenServer(ctx, &wg)
 
 	wg.Add(1)
-	abmf.OpenServer(&wg)
+	abmf.OpenServer(ctx, &wg)
 	// Register to NRF
 	profile, err := consumer.BuildNFInstance(self)
 	if err != nil {
@@ -144,7 +147,9 @@ func (c *ChfApp) Start(tlsKeyLogPath string) {
 		}()
 
 		<-signalChannel
+		cancel()
 		c.Terminate()
+		wg.Wait()
 		os.Exit(0)
 	}()