Fix exiting too fast for servers to start

toxuin · Jun 7, 2021 · 823e2ce · 823e2ce
1 parent a8b60a3
commit 823e2ce
Show file tree

Hide file tree

Showing 4 changed files with 37 additions and 13 deletions.
diff --git a/main.go b/main.go
@@ -8,6 +8,7 @@ import (
 	"github.com/toxuin/alarmserver/servers/ftp"
 	"github.com/toxuin/alarmserver/servers/hikvision"
 	"github.com/toxuin/alarmserver/servers/hisilicon"
+	"sync"
 )
 
 var config *conf.Config
@@ -23,6 +24,8 @@ func main() {
 		config.Printout()
 	}
 
+	processesWaitGroup := sync.WaitGroup{}
+
 	// INIT BUSES
 	mqttBus := mqtt.Bus{Debug: config.Debug}
 	if config.Mqtt.Enabled {
@@ -53,6 +56,7 @@ func main() {
 		// START HISILICON ALARM SERVER
 		hisiliconServer := hisilicon.Server{
 			Debug:          config.Debug,
+			WaitGroup:      &processesWaitGroup,
 			Port:           config.Hisilicon.Port,
 			MessageHandler: messageHandler,
 		}
@@ -66,6 +70,7 @@ func main() {
 		// START HIKVISION ALARM SERVER
 		hikvisionServer := hikvision.Server{
 			Debug:          config.Debug,
+			WaitGroup:      &processesWaitGroup,
 			Cameras:        &config.Hikvision.Cams,
 			MessageHandler: messageHandler,
 		}
@@ -79,6 +84,7 @@ func main() {
 		// START FTP SERVER
 		ftpServer := ftp.Server{
 			Debug:          config.Debug,
+			WaitGroup:      &processesWaitGroup,
 			Port:           config.Ftp.Port,
 			AllowFiles:     config.Ftp.AllowFiles,
 			RootPath:       config.Ftp.RootPath,
@@ -90,4 +96,6 @@ func main() {
 			fmt.Println("STARTED FTP SERVER")
 		}
 	}
+
+	processesWaitGroup.Wait()
 }
diff --git a/servers/ftp/server.go b/servers/ftp/server.go
@@ -3,10 +3,12 @@ package ftp
 import (
 	"fmt"
 	"goftp.io/server/v2"
+	"sync"
 )
 
 type Server struct {
 	Debug          bool
+	WaitGroup      *sync.WaitGroup
 	Port           int
 	AllowFiles     bool
 	RootPath       string
@@ -32,6 +34,9 @@ func (serv *Server) Start() {
 		serv.Password = "root"
 	}
 
+	defer serv.WaitGroup.Done()
+	serv.WaitGroup.Add(1)
+
 	eventChannel := make(chan Event, 5)
 
 	// START MESSAGE PROCESSOR

diff --git a/servers/hikvision/server.go b/servers/hikvision/server.go
@@ -24,6 +24,7 @@ type HikEvent struct {
 
 type Server struct {
 	Debug          bool
+	WaitGroup      *sync.WaitGroup
 	Cameras        *[]HikCamera
 	MessageHandler func(topic string, data string)
 }
@@ -87,23 +88,28 @@ func (server *Server) Start() {
 		}
 	}
 
-	waitGroup := sync.WaitGroup{}
+	cameraWaitGroup := sync.WaitGroup{}
 	eventChannel := make(chan HikEvent, 5)
 
 	// START ALL CAMERA LISTENERS
 	for _, camera := range *server.Cameras {
-		server.addCamera(&waitGroup, &camera, eventChannel)
+		server.addCamera(&cameraWaitGroup, &camera, eventChannel)
 	}
 
 	// START MESSAGE PROCESSOR
-	go func(waitGroup *sync.WaitGroup, channel <-chan HikEvent) {
-		defer waitGroup.Done()
+	go func(camWaitGroup *sync.WaitGroup, channel <-chan HikEvent) {
+		// WAIT GROUP FOR INDIVIDUAL CAMERAS
+		defer camWaitGroup.Done()
+
+		// EXTERNAL WAIT GROUP FOR PROCESSES
+		defer server.WaitGroup.Done()
+		server.WaitGroup.Add(1)
 		for {
 			event := <-channel
 			go server.MessageHandler(event.Camera.Name+"/"+event.Type, event.Message)
 		}
-	}(&waitGroup, eventChannel)
-	waitGroup.Add(1)
+	}(&cameraWaitGroup, eventChannel)
+	cameraWaitGroup.Add(1)
 
-	waitGroup.Wait()
+	cameraWaitGroup.Wait()
 }
diff --git a/servers/hisilicon/server.go b/servers/hisilicon/server.go
@@ -8,6 +8,7 @@ import (
 	"net"
 	"strconv"
 	"strings"
+	"sync"
 )
 
 // Converts 0x1704A8C0 to 192.168.4.23
@@ -36,6 +37,13 @@ func hexIpToCIDR(hexAddr string) string {
 	return ipAddr
 }
 
+type Server struct {
+	Debug          bool
+	WaitGroup      *sync.WaitGroup
+	Port           string
+	MessageHandler func(topic string, data string)
+}
+
 func (server *Server) handleTcpConnection(conn net.Conn) {
 	defer conn.Close()
 
@@ -84,12 +92,6 @@ func (server *Server) handleTcpConnection(conn net.Conn) {
 	server.MessageHandler(serialId+"/"+event, string(jsonBytes))
 }
 
-type Server struct {
-	Debug          bool
-	Port           string
-	MessageHandler func(topic string, data string)
-}
-
 func (server *Server) Start() {
 	if server.Port == "" {
 		server.Port = "15002" // DEFAULT PORT
@@ -102,6 +104,9 @@ func (server *Server) Start() {
 	}
 
 	go func() {
+		defer server.WaitGroup.Done()
+		server.WaitGroup.Add(1)
+
 		// START TCP SERVER
 		tcpListener, err := net.Listen("tcp4", ":"+server.Port)
 		if err != nil {