From 32ae99c225284d4430d8c98a04c3ab91b4cb361b Mon Sep 17 00:00:00 2001 From: Luther Wen Xu Date: Sat, 12 Oct 2019 23:49:28 +0800 Subject: [PATCH] go: Replace wait group with channel --- GoBot/discord/discord.go | 14 +++++--------- GoBot/main.go | 21 ++++++++++++++++----- GoBot/server/server.go | 5 +---- 3 files changed, 22 insertions(+), 18 deletions(-) diff --git a/GoBot/discord/discord.go b/GoBot/discord/discord.go index 5f07af4..07ffcea 100644 --- a/GoBot/discord/discord.go +++ b/GoBot/discord/discord.go @@ -2,19 +2,14 @@ package discord import ( "fmt" - "os" - "os/signal" "strings" - "sync" - "syscall" "github.com/bwmarrin/discordgo" "TerraOceanBot/discord/modules" ) -func StartBot(token string, wg *sync.WaitGroup) { - defer wg.Done() +func StartBot(token string, kill chan bool) { dg, err := discordgo.New("Bot " + strings.TrimSpace(token)) if err != nil { panic(err) @@ -32,10 +27,11 @@ func StartBot(token string, wg *sync.WaitGroup) { fmt.Println("Bot is now running. Press CTRL-C to exit.") go modules.ListenToVoteFinishes(dg) - sc := make(chan os.Signal, 1) - signal.Notify(sc, syscall.SIGINT, syscall.SIGTERM, os.Interrupt, os.Kill) - <-sc + + <-kill dg.Close() fmt.Println("Discordgo exited successfully.") + + <-kill } diff --git a/GoBot/main.go b/GoBot/main.go index 77244be..311c492 100644 --- a/GoBot/main.go +++ b/GoBot/main.go @@ -3,15 +3,26 @@ package main import ( "TerraOceanBot/server" "io/ioutil" - "sync" + "os" + "os/signal" + "syscall" "TerraOceanBot/discord" ) func main() { token, _ := ioutil.ReadFile("token.txt") - wg := sync.WaitGroup{} - go discord.StartBot(string(token), &wg) - server.StartServer(8080, &wg) - wg.Wait() + endGracefully := make(chan bool) + + go discord.StartBot(string(token), endGracefully) + go server.StartServer(8080) + + sc := make(chan os.Signal, 1) + signal.Notify(sc, syscall.SIGINT, syscall.SIGTERM, os.Interrupt, os.Kill) + <-sc + + //Notification to end + endGracefully <- true + //Wait for it to actually end + endGracefully <- true } diff --git a/GoBot/server/server.go b/GoBot/server/server.go index be6c8c5..6f521ac 100644 --- a/GoBot/server/server.go +++ b/GoBot/server/server.go @@ -5,7 +5,6 @@ import ( "log" "net/http" "strconv" - "sync" "github.com/gorilla/websocket" ) @@ -35,9 +34,7 @@ func newConnection(w http.ResponseWriter, r *http.Request) { } } -func StartServer(port int, wg *sync.WaitGroup) { - defer wg.Done() - +func StartServer(port int) { hostLocation := ":" + strconv.Itoa(port) http.HandleFunc("/ws", newConnection)