go: Replace wait group with channel

master
Luther Wen Xu 2019-10-12 23:49:28 +07:00
parent 38580f8910
commit 32ae99c225
Signed by untrusted user: chanbakjsd
GPG Key ID: B7D77E3E9D102B70
3 changed files with 22 additions and 18 deletions

@ -2,19 +2,14 @@ package discord
import ( import (
"fmt" "fmt"
"os"
"os/signal"
"strings" "strings"
"sync"
"syscall"
"github.com/bwmarrin/discordgo" "github.com/bwmarrin/discordgo"
"TerraOceanBot/discord/modules" "TerraOceanBot/discord/modules"
) )
func StartBot(token string, wg *sync.WaitGroup) { func StartBot(token string, kill chan bool) {
defer wg.Done()
dg, err := discordgo.New("Bot " + strings.TrimSpace(token)) dg, err := discordgo.New("Bot " + strings.TrimSpace(token))
if err != nil { if err != nil {
panic(err) panic(err)
@ -32,10 +27,11 @@ func StartBot(token string, wg *sync.WaitGroup) {
fmt.Println("Bot is now running. Press CTRL-C to exit.") fmt.Println("Bot is now running. Press CTRL-C to exit.")
go modules.ListenToVoteFinishes(dg) go modules.ListenToVoteFinishes(dg)
sc := make(chan os.Signal, 1)
signal.Notify(sc, syscall.SIGINT, syscall.SIGTERM, os.Interrupt, os.Kill) <-kill
<-sc
dg.Close() dg.Close()
fmt.Println("Discordgo exited successfully.") fmt.Println("Discordgo exited successfully.")
<-kill
} }

@ -3,15 +3,26 @@ package main
import ( import (
"TerraOceanBot/server" "TerraOceanBot/server"
"io/ioutil" "io/ioutil"
"sync" "os"
"os/signal"
"syscall"
"TerraOceanBot/discord" "TerraOceanBot/discord"
) )
func main() { func main() {
token, _ := ioutil.ReadFile("token.txt") token, _ := ioutil.ReadFile("token.txt")
wg := sync.WaitGroup{} endGracefully := make(chan bool)
go discord.StartBot(string(token), &wg)
server.StartServer(8080, &wg) go discord.StartBot(string(token), endGracefully)
wg.Wait() 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
} }

@ -5,7 +5,6 @@ import (
"log" "log"
"net/http" "net/http"
"strconv" "strconv"
"sync"
"github.com/gorilla/websocket" "github.com/gorilla/websocket"
) )
@ -35,9 +34,7 @@ func newConnection(w http.ResponseWriter, r *http.Request) {
} }
} }
func StartServer(port int, wg *sync.WaitGroup) { func StartServer(port int) {
defer wg.Done()
hostLocation := ":" + strconv.Itoa(port) hostLocation := ":" + strconv.Itoa(port)
http.HandleFunc("/ws", newConnection) http.HandleFunc("/ws", newConnection)