Start components with ctx, and wait for shutdown

nextmn · Dec 12, 2024 · 23c768b · 23c768b
1 parent 0a518d3
commit 23c768b
Show file tree

Hide file tree

Showing 4 changed files with 86 additions and 34 deletions.
diff --git a/internal/app/control.go b/internal/app/control.go
@@ -22,9 +22,10 @@ import (
 )
 
 type HttpServerEntity struct {
-	srv   *http.Server
-	ps    *session.PduSessions
-	radio *radio.Radio
+	srv    *http.Server
+	ps     *session.PduSessions
+	radio  *radio.Radio
+	closed chan struct{}
 }
 
 func NewHttpServerEntity(bindAddr netip.AddrPort, r *radio.Radio, ps *session.PduSessions) *HttpServerEntity {
@@ -45,13 +46,14 @@ func NewHttpServerEntity(bindAddr netip.AddrPort, r *radio.Radio, ps *session.Pd
 			Addr:    bindAddr.String(),
 			Handler: h,
 		},
-		ps:    ps,
-		radio: r,
+		ps:     ps,
+		radio:  r,
+		closed: make(chan struct{}),
 	}
 	return &e
 }
 
-func (e *HttpServerEntity) Start() error {
+func (e *HttpServerEntity) Start(ctx context.Context) error {
 	l, err := net.Listen("tcp", e.srv.Addr)
 	if err != nil {
 		return err
@@ -62,14 +64,26 @@ func (e *HttpServerEntity) Start() error {
 			logrus.WithError(err).Error("Http Server error")
 		}
 	}(l)
+	go func(ctx context.Context) {
+		defer close(e.closed)
+		select {
+		case <-ctx.Done():
+			ctxShutdown, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
+			defer cancel()
+			if err := e.srv.Shutdown(ctxShutdown); err != nil {
+				logrus.WithError(err).Info("HTTP Server Shutdown")
+			}
+		}
+	}(ctx)
 	return nil
 }
 
-func (e *HttpServerEntity) Stop() {
-	ctx, cancel := context.WithTimeout(context.Background(), 1*time.Second) // context.Background() is already Done()
-	defer cancel()
-	if err := e.srv.Shutdown(ctx); err != nil {
-		logrus.WithError(err).Info("HTTP Server Shutdown")
+func (e *HttpServerEntity) WaitShutdown(ctx context.Context) error {
+	select {
+	case <-ctx.Done():
+		return ctx.Err()
+	case <-e.closed:
+		return nil
 	}
 }
 

diff --git a/internal/app/setup.go b/internal/app/setup.go
@@ -7,8 +7,10 @@ package app
 
 import (
 	"context"
+	"time"
 
 	"github.com/nextmn/gnb-lite/internal/config"
+	"github.com/nextmn/gnb-lite/internal/gtp"
 	"github.com/nextmn/gnb-lite/internal/radio"
 	"github.com/nextmn/gnb-lite/internal/session"
 )
@@ -19,6 +21,7 @@ type Setup struct {
 	radio            *radio.Radio
 	rDaemon          *radio.RadioDaemon
 	psMan            *session.PduSessionsManager
+	gtp              *gtp.Gtp
 }
 
 func NewSetup(config *config.GNBConfig) *Setup {
@@ -32,33 +35,30 @@ func NewSetup(config *config.GNBConfig) *Setup {
 		radio:            r,
 		rDaemon:          rDaemon,
 		psMan:            psMan,
+		gtp:              gtp.NewGtp(config.Gtp, psMan, rDaemon),
 	}
 }
 func (s *Setup) Init(ctx context.Context) error {
+	return nil
+}
+
+func (s *Setup) Run(ctx context.Context) error {
 	if err := s.rDaemon.Start(ctx); err != nil {
 		return err
 	}
-	if err := s.StartGtpUProtocolEntity(ctx, s.config.Gtp); err != nil {
+	if err := s.gtp.Start(ctx); err != nil {
 		return err
 	}
-	if err := s.httpServerEntity.Start(); err != nil {
-		return err
-	}
-	return nil
-}
-
-func (s *Setup) Run(ctx context.Context) error {
-	defer s.Exit()
-	if err := s.Init(ctx); err != nil {
+	if err := s.httpServerEntity.Start(ctx); err != nil {
 		return err
 	}
 	select {
 	case <-ctx.Done():
+		ctxShutdown, cancel := context.WithTimeout(ctx, 1*time.Second)
+		defer cancel()
+		s.httpServerEntity.WaitShutdown(ctxShutdown)
+		s.gtp.WaitShutdown(ctxShutdown)
+		s.rDaemon.WaitShutdown(ctxShutdown)
 		return nil
 	}
 }
-
-func (s *Setup) Exit() error {
-	s.httpServerEntity.Stop()
-	return nil
-}
diff --git a/internal/app/gtp.go → internal/gtp/gtp.go b/internal/app/gtp.go → internal/gtp/gtp.go
@@ -3,7 +3,7 @@
 // found in the LICENSE file.
 // SPDX-License-Identifier: MIT
 
-package app
+package gtp
 
 import (
 	"context"
@@ -18,17 +18,34 @@ import (
 	"github.com/wmnsk/go-gtp/gtpv1/message"
 )
 
+type Gtp struct {
+	ipAddr  netip.Addr
+	psMan   *session.PduSessionsManager
+	rDaemon *radio.RadioDaemon
+	closed  chan struct{}
+}
+
 const GTPU_PORT = 2152
 
-func (s *Setup) StartGtpUProtocolEntity(ctx context.Context, ipAddress netip.Addr) error {
-	logrus.WithFields(logrus.Fields{"listen-addr": ipAddress}).Info("Creating new GTP-U Protocol Entity")
-	laddr := net.UDPAddrFromAddrPort(netip.AddrPortFrom(ipAddress, GTPU_PORT))
+func NewGtp(ipAddr netip.Addr, psMan *session.PduSessionsManager, rDaemon *radio.RadioDaemon) *Gtp {
+	return &Gtp{
+		ipAddr:  ipAddr,
+		psMan:   psMan,
+		rDaemon: rDaemon,
+		closed:  make(chan struct{}),
+	}
+}
+
+func (gtp *Gtp) Start(ctx context.Context) error {
+	logrus.WithFields(logrus.Fields{"listen-addr": gtp.ipAddr}).Info("Creating new GTP-U Protocol Entity")
+	laddr := net.UDPAddrFromAddrPort(netip.AddrPortFrom(gtp.ipAddr, GTPU_PORT))
 	uConn := gtpv1.NewUPlaneConn(laddr)
 	uConn.DisableErrorIndication()
 	uConn.AddHandler(message.MsgTypeTPDU, func(c gtpv1.Conn, senderAddr net.Addr, msg message.Message) error {
-		return tpduHandler(c, senderAddr, msg, s.psMan, s.rDaemon)
+		return gtp.tpduHandler(c, senderAddr, msg)
 	})
 	go func(ctx context.Context) error {
+		defer close(gtp.closed)
 		defer uConn.Close()
 		if err := uConn.ListenAndServe(ctx); err != nil {
 			logrus.WithError(err).Trace("GTP uConn closed")
@@ -42,12 +59,21 @@ func (s *Setup) StartGtpUProtocolEntity(ctx context.Context, ipAddress netip.Add
 }
 
 // handle GTP PDU (Downlink)
-func tpduHandler(c gtpv1.Conn, senderAddr net.Addr, msg message.Message, psMan *session.PduSessionsManager, rDaemon *radio.RadioDaemon) error {
+func (gtp *Gtp) tpduHandler(c gtpv1.Conn, senderAddr net.Addr, msg message.Message) error {
 	teid := msg.TEID()
-	ue, err := psMan.GetUECtrl(teid)
+	ue, err := gtp.psMan.GetUECtrl(teid)
 	if err != nil {
 		return err
 	}
 	packet := msg.(*message.TPDU).Decapsulate()
-	return rDaemon.WriteDownlink(packet, ue)
+	return gtp.rDaemon.WriteDownlink(packet, ue)
+}
+
+func (gtp *Gtp) WaitShutdown(ctx context.Context) error {
+	select {
+	case <-ctx.Done():
+		return ctx.Err()
+	case <-gtp.closed:
+		return nil
+	}
 }
diff --git a/internal/radio/radio_daemon.go b/internal/radio/radio_daemon.go
@@ -28,6 +28,7 @@ type RadioDaemon struct {
 	gnbRanAddr         netip.AddrPort
 	PduSessionsManager *session.PduSessionsManager
 	srv                *net.UDPConn
+	closed             chan struct{}
 }
 
 func NewRadioDaemon(radio *Radio, psMan *session.PduSessionsManager, gnbRanAddr netip.AddrPort) *RadioDaemon {
@@ -36,6 +37,7 @@ func NewRadioDaemon(radio *Radio, psMan *session.PduSessionsManager, gnbRanAddr
 		radio:              radio,
 		PduSessionsManager: psMan,
 		gnbRanAddr:         gnbRanAddr,
+		closed:             make(chan struct{}),
 	}
 }
 
@@ -95,8 +97,18 @@ func (r *RadioDaemon) Start(ctx context.Context) error {
 		return nil
 	}(ctx, srv)
 	go func(ctx context.Context, srv *net.UDPConn) {
+		defer close(r.closed)
 		defer srv.Close()
 		r.runUplinkDaemon(ctx, srv)
 	}(ctx, srv)
 	return nil
 }
+
+func (r *RadioDaemon) WaitShutdown(ctx context.Context) error {
+	select {
+	case <-ctx.Done():
+		return ctx.Err()
+	case <-r.closed:
+		return nil
+	}
+}