Start component with ctx, and wait for shutdown

internal/app/control.go

@@ -23,10 +23,11 @@ import (
 )
 
 type HttpServerEntity struct {
-	srv   *http.Server
-	ps    *session.PduSessions
-	radio *radio.Radio
-	cli   *cli.Cli
+	srv    *http.Server
+	ps     *session.PduSessions
+	radio  *radio.Radio
+	cli    *cli.Cli
+	closed chan struct{}
 }
 
 func NewHttpServerEntity(bindAddr netip.AddrPort, r *radio.Radio, ps *session.PduSessions) *HttpServerEntity {
@@ -51,14 +52,15 @@ func NewHttpServerEntity(bindAddr netip.AddrPort, r *radio.Radio, ps *session.Pd
 			Addr:    bindAddr.String(),
 			Handler: h,
 		},
-		ps:    ps,
-		radio: r,
-		cli:   c,
+		ps:     ps,
+		radio:  r,
+		cli:    c,
+		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
@@ -69,14 +71,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():
+			ctxTimeout, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
+			defer cancel()
+			if err := e.srv.Shutdown(ctxTimeout); 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
 	}
 }
 

internal/app/setup.go

@@ -7,6 +7,7 @@ package app
 
 import (
 	"context"
+	"time"
 
 	"github.com/nextmn/ue-lite/internal/config"
 	"github.com/nextmn/ue-lite/internal/radio"
@@ -36,8 +37,8 @@ func NewSetup(config *config.UEConfig) *Setup {
 	}
 }
 
-func (s *Setup) Init(ctx context.Context) error {
-	if err := s.httpServerEntity.Start(); err != nil {
+func (s *Setup) Run(ctx context.Context) error {
+	if err := s.httpServerEntity.Start(ctx); err != nil {
 		return err
 	}
 	if err := s.tunMan.Start(ctx); err != nil {
@@ -49,21 +50,14 @@ func (s *Setup) Init(ctx context.Context) error {
 	if err := s.ps.Start(ctx); err != nil {
 		return err
 	}
-	return nil
-}
-
-func (s *Setup) Run(ctx context.Context) error {
-	defer s.Exit()
-	if err := s.Init(ctx); err != nil {
-		return err
-	}
 	select {
 	case <-ctx.Done():
+		ctxShutdown, cancel := context.WithTimeout(ctx, 1*time.Second)
+		defer cancel()
+		s.ps.WaitShutdown(ctxShutdown)
+		s.radioDaemon.WaitShutdown(ctxShutdown)
+		s.tunMan.WaitShutdown(ctxShutdown)
+		s.httpServerEntity.WaitShutdown(ctxShutdown)
 		return nil
 	}
 }
-
-func (s *Setup) Exit() error {
-	s.httpServerEntity.Stop()
-	return nil
-}

internal/radio/radio_daemon.go

@@ -28,6 +28,7 @@ type RadioDaemon struct {
 	PsMan     *session.PduSessionsManager
 	UeRanAddr netip.AddrPort
 	tunMan    *tun.TunManager
+	closed    chan struct{}
 }
 
 func NewRadioDaemon(control jsonapi.ControlURI, gnbs []jsonapi.ControlURI, radio *Radio, psMan *session.PduSessionsManager, tunMan *tun.TunManager, ueRanAddr netip.AddrPort) *RadioDaemon {
@@ -38,6 +39,7 @@ func NewRadioDaemon(control jsonapi.ControlURI, gnbs []jsonapi.ControlURI, radio
 		PsMan:     psMan,
 		UeRanAddr: ueRanAddr,
 		tunMan:    tunMan,
+		closed:    make(chan struct{}),
 	}
 }
 
@@ -108,7 +110,15 @@ func (r *RadioDaemon) runUplinkDaemon(ctx context.Context, srv *net.UDPConn, ifa
 }
 
 func (r *RadioDaemon) Start(ctx context.Context) error {
-	ifacetun := r.tunMan.Tun
+	ifacetun := r.tunMan.OpenTun()
+	defer func(ctx context.Context) {
+		defer r.tunMan.CloseTun()
+		select {
+		case <-ctx.Done():
+			close(r.closed)
+			return
+		}
+	}(ctx)
 	srv, err := net.ListenUDP("udp", net.UDPAddrFromAddrPort(r.UeRanAddr))
 	if err != nil {
 		return err
@@ -143,3 +153,12 @@ func (r *RadioDaemon) Start(ctx context.Context) error {
 	}
 	return nil
 }
+
+func (r *RadioDaemon) WaitShutdown(ctx context.Context) error {
+	select {
+	case <-ctx.Done():
+		return ctx.Err()
+	case <-r.closed:
+		return nil
+	}
+}

internal/session/pdu_session.go

@@ -93,3 +93,8 @@ func (p *PduSessions) Start(ctx context.Context) error {
 	}
 	return nil
 }
+
+func (p *PduSessions) WaitShutdown(ctx context.Context) error {
+	// nothing to do
+	return nil
+}

internal/tun/tun.go

@@ -10,6 +10,7 @@ import (
 	"fmt"
 	"net/netip"
 	"strconv"
+	"sync"
 
 	"github.com/sirupsen/logrus"
 	"github.com/songgao/water"
@@ -21,36 +22,62 @@ const (
 )
 
 type TunManager struct {
-	ready bool
-	name  string
-	Tun   *water.Interface
+	ready  bool
+	name   string
+	tun    *water.Interface
+	closed chan struct{}
+	used   sync.WaitGroup
 }
 
 func NewTunManager() *TunManager {
-	return &TunManager{}
+	return &TunManager{
+		closed: make(chan struct{}),
+	}
+}
+
+// Get a tun interface.
+// Don't forget to run CloseTun when no longer in use
+func (t *TunManager) OpenTun() *water.Interface {
+	t.used.Add(1)
+	return t.tun
+}
+
+func (t *TunManager) CloseTun() {
+	t.used.Done()
 }
 
 func (t *TunManager) Start(ctx context.Context) error {
 	tun, err := newTunIface()
-	t.Tun = tun
+	t.tun = tun
 	if err != nil {
 		return err
 	}
 	t.ready = true
-	t.name = t.Tun.Name()
+	t.name = t.tun.Name()
 	go func(ctx context.Context) {
 		select {
 		case <-ctx.Done():
+			t.used.Wait() // Do not delete tun iface until all tuns are closed
 			err = runIPTables("-D", "OUTPUT", "-o", t.name, "-p", "icmp", "--icmp-type", "redirect", "-j", "DROP")
 			if err != nil {
 				logrus.WithError(err).WithFields(logrus.Fields{"interface": t.name}).Error("Error while removing iptables rules")
 				t.ready = false
+				close(t.closed)
 			}
 		}
 	}(ctx)
 	return err
 }
 
+func (t *TunManager) WaitShutdown(ctx context.Context) error {
+	select {
+	case <-ctx.Done():
+		return ctx.Err()
+	case <-t.closed:
+		return nil
+	}
+}
+
 func newTunIface() (*water.Interface, error) {
 	config := water.Config{
 		DeviceType: water.TUN,