Using a waitgroup to wait for all programs to exit

cmd/gluetun/main.go

@@ -7,6 +7,7 @@ import (
 	"os/signal"
 	"regexp"
 	"strings"
+	"sync"
 	"syscall"
 	"time"
 
@@ -164,29 +165,25 @@ func _main(background context.Context, args []string) int {
 	restartPublicIP := make(chan struct{})
 	restartTinyproxy := make(chan struct{})
 	restartShadowsocks := make(chan struct{})
-	openvpnDone := make(chan struct{})
+	wg := &sync.WaitGroup{}
-	unboundDone := make(chan struct{})
-	serverDone := make(chan struct{})
-	tinyproxyDone := make(chan struct{})
-	shadowsocksDone := make(chan struct{})
 
 	openvpnLooper := openvpn.NewLooper(ovpnConf, allSettings.OpenVPN, logger, streamMerger, fatalOnError, uid, gid)
 	// wait for restartOpenvpn
-	go openvpnLooper.Run(ctx, restartOpenvpn, openvpnDone)
+	go openvpnLooper.Run(ctx, restartOpenvpn, wg)
 
 	unboundLooper := dns.NewLooper(dnsConf, allSettings.DNS, logger, streamMerger, uid, gid)
 	// wait for restartUnbound
-	go unboundLooper.Run(ctx, restartUnbound, unboundDone)
+	go unboundLooper.Run(ctx, restartUnbound, wg)
 
 	publicIPLooper := publicip.NewLooper(client, logger, fileManager, allSettings.System.IPStatusFilepath, uid, gid)
 	go publicIPLooper.Run(ctx, restartPublicIP)
 	go publicIPLooper.RunRestartTicker(ctx, restartPublicIP)
 
 	tinyproxyLooper := tinyproxy.NewLooper(tinyProxyConf, firewallConf, allSettings.TinyProxy, logger, streamMerger, uid, gid)
-	go tinyproxyLooper.Run(ctx, restartTinyproxy, tinyproxyDone)
+	go tinyproxyLooper.Run(ctx, restartTinyproxy, wg)
 
 	shadowsocksLooper := shadowsocks.NewLooper(shadowsocksConf, firewallConf, allSettings.ShadowSocks, allSettings.DNS, logger, streamMerger, uid, gid)
-	go shadowsocksLooper.Run(ctx, restartShadowsocks, shadowsocksDone)
+	go shadowsocksLooper.Run(ctx, restartShadowsocks, wg)
 
 	if allSettings.TinyProxy.Enabled {
 		<-restartTinyproxy
@@ -218,7 +215,7 @@ func _main(background context.Context, args []string) int {
 	}()
 
 	httpServer := server.New("0.0.0.0:8000", logger, restartOpenvpn, restartUnbound)
-	go httpServer.Run(ctx, serverDone)
+	go httpServer.Run(ctx, wg)
 
 	// Start openvpn for the first time
 	restartOpenvpn <- struct{}{}
@@ -256,11 +253,7 @@ func _main(background context.Context, args []string) int {
 		logger.Error(err)
 		exitStatus = 1
 	}
-	<-serverDone
+	wg.Wait()
-	<-unboundDone
-	<-openvpnDone
-	<-tinyproxyDone
-	<-shadowsocksDone
 	return exitStatus
 }
 

internal/dns/loop.go

@@ -3,6 +3,7 @@ package dns
 import (
 	"context"
 	"net"
+	"sync"
 	"time"
 
 	"github.com/qdm12/golibs/command"
@@ -12,7 +13,7 @@ import (
 )
 
 type Looper interface {
-	Run(ctx context.Context, restart <-chan struct{}, done chan<- struct{})
+	Run(ctx context.Context, restart <-chan struct{}, wg *sync.WaitGroup)
 	RunRestartTicker(ctx context.Context, restart chan<- struct{})
 }
 
@@ -43,12 +44,13 @@ func (l *looper) attemptingRestart(err error) {
 	time.Sleep(10 * time.Second)
 }
 
-func (l *looper) Run(ctx context.Context, restart <-chan struct{}, done chan<- struct{}) {
+func (l *looper) Run(ctx context.Context, restart <-chan struct{}, wg *sync.WaitGroup) {
+	wg.Add(1)
+	defer wg.Done()
 	l.fallbackToUnencryptedDNS()
 	select {
 	case <-restart:
 	case <-ctx.Done():
-		close(done)
 		return
 	}
 	_, unboundCancel := context.WithCancel(ctx)
@@ -59,13 +61,11 @@ func (l *looper) Run(ctx context.Context, restart <-chan struct{}, done chan<- s
 			case <-restart:
 			case <-ctx.Done():
 				unboundCancel()
-				close(done)
 				return
 			}
 		}
 		if ctx.Err() == context.Canceled {
 			unboundCancel()
-			close(done)
 			return
 		}
 
@@ -121,7 +121,6 @@ func (l *looper) Run(ctx context.Context, restart <-chan struct{}, done chan<- s
 			l.logger.Warn("context canceled: exiting loop")
 			unboundCancel()
 			close(waitError)
-			close(done)
 			return
 		case <-restart: // triggered restart
 			unboundCancel()

internal/openvpn/loop.go

@@ -2,6 +2,7 @@ package openvpn
 
 import (
 	"context"
+	"sync"
 	"time"
 
 	"github.com/qdm12/golibs/command"
@@ -11,7 +12,7 @@ import (
 )
 
 type Looper interface {
-	Run(ctx context.Context, restart <-chan struct{}, done chan<- struct{})
+	Run(ctx context.Context, restart <-chan struct{}, wg *sync.WaitGroup)
 }
 
 type looper struct {
@@ -37,11 +38,12 @@ func NewLooper(conf Configurator, settings settings.OpenVPN, logger logging.Logg
 	}
 }
 
-func (l *looper) Run(ctx context.Context, restart <-chan struct{}, done chan<- struct{}) {
+func (l *looper) Run(ctx context.Context, restart <-chan struct{}, wg *sync.WaitGroup) {
+	wg.Add(1)
+	defer wg.Done()
 	select {
 	case <-restart:
 	case <-ctx.Done():
-		close(done)
 		return
 	}
 	for {
@@ -69,7 +71,6 @@ func (l *looper) Run(ctx context.Context, restart <-chan struct{}, done chan<- s
 			l.logger.Warn("context canceled: exiting loop")
 			openvpnCancel()
 			close(waitError)
-			close(done)
 			return
 		case <-restart: // triggered restart
 			l.logger.Info("restarting")

internal/server/server.go

@@ -4,13 +4,14 @@ import (
 	"context"
 	"fmt"
 	"net/http"
+	"sync"
 	"time"
 
 	"github.com/qdm12/golibs/logging"
 )
 
 type Server interface {
-	Run(ctx context.Context, serverDone chan struct{})
+	Run(ctx context.Context, wg *sync.WaitGroup)
 }
 
 type server struct {
@@ -29,10 +30,11 @@ func New(address string, logger logging.Logger, restartOpenvpn, restartUnbound c
 	}
 }
 
-func (s *server) Run(ctx context.Context, serverDone chan struct{}) {
+func (s *server) Run(ctx context.Context, wg *sync.WaitGroup) {
+	wg.Add(1)
 	server := http.Server{Addr: s.address, Handler: s.makeHandler()}
 	go func() {
-		defer close(serverDone)
+		defer wg.Done()
 		<-ctx.Done()
 		s.logger.Warn("context canceled: exiting loop")
 		shutdownCtx, cancel := context.WithTimeout(context.Background(), 2*time.Second)

internal/shadowsocks/loop.go

@@ -2,6 +2,7 @@ package shadowsocks
 
 import (
 	"context"
+	"sync"
 	"time"
 
 	"github.com/qdm12/golibs/command"
@@ -12,7 +13,7 @@ import (
 )
 
 type Looper interface {
-	Run(ctx context.Context, restart <-chan struct{}, done chan<- struct{})
+	Run(ctx context.Context, restart <-chan struct{}, wg *sync.WaitGroup)
 }
 
 type looper struct {
@@ -46,11 +47,12 @@ func NewLooper(conf Configurator, firewallConf firewall.Configurator, settings s
 	}
 }
 
-func (l *looper) Run(ctx context.Context, restart <-chan struct{}, done chan<- struct{}) {
+func (l *looper) Run(ctx context.Context, restart <-chan struct{}, wg *sync.WaitGroup) {
+	wg.Add(1)
+	defer wg.Done()
 	select {
 	case <-restart:
 	case <-ctx.Done():
-		close(done)
 		return
 	}
 	for {
@@ -97,7 +99,6 @@ func (l *looper) Run(ctx context.Context, restart <-chan struct{}, done chan<- s
 			l.logger.Warn("context canceled: exiting loop")
 			shadowsocksCancel()
 			close(waitError)
-			close(done)
 			return
 		case <-restart: // triggered restart
 			l.logger.Info("restarting")

internal/tinyproxy/loop.go

@@ -2,6 +2,7 @@ package tinyproxy
 
 import (
 	"context"
+	"sync"
 	"time"
 
 	"github.com/qdm12/golibs/command"
@@ -12,7 +13,7 @@ import (
 )
 
 type Looper interface {
-	Run(ctx context.Context, restart <-chan struct{}, done chan<- struct{})
+	Run(ctx context.Context, restart <-chan struct{}, wg *sync.WaitGroup)
 }
 
 type looper struct {
@@ -44,11 +45,12 @@ func NewLooper(conf Configurator, firewallConf firewall.Configurator, settings s
 	}
 }
 
-func (l *looper) Run(ctx context.Context, restart <-chan struct{}, done chan<- struct{}) {
+func (l *looper) Run(ctx context.Context, restart <-chan struct{}, wg *sync.WaitGroup) {
+	wg.Add(1)
+	defer wg.Done()
 	select {
 	case <-restart:
 	case <-ctx.Done():
-		close(done)
 		return
 	}
 	for {
@@ -89,7 +91,6 @@ func (l *looper) Run(ctx context.Context, restart <-chan struct{}, done chan<- s
 			l.logger.Warn("context canceled: exiting loop")
 			tinyproxyCancel()
 			close(waitError)
-			close(done)
 			return
 		case <-restart: // triggered restart
 			l.logger.Info("restarting")