From 5e69ce099f521aa49fc267c62235c003eae530ff Mon Sep 17 00:00:00 2001
From: Chris <ccbrown112@gmail.com>
Date: Tue, 3 Oct 2017 10:53:53 -0500
Subject: Goroutine wranglin (#7556)

* goroutine wranglin

* synchronize WebConn.WritePump
---
 app/app.go | 53 +++++++++++++++++++++++++++++++++++++++++++++--------
 1 file changed, 45 insertions(+), 8 deletions(-)

(limited to 'app/app.go')

diff --git a/app/app.go b/app/app.go
index 7974ab44f..d0d5bb4e0 100644
--- a/app/app.go
+++ b/app/app.go
@@ -7,7 +7,9 @@ import (
 	"io/ioutil"
 	"net/http"
 	"sync"
-	"time"
+	"sync/atomic"
+
+	l4g "github.com/alecthomas/log4go"
 
 	"github.com/mattermost/mattermost-server/einterfaces"
 	ejobs "github.com/mattermost/mattermost-server/einterfaces/jobs"
@@ -18,6 +20,9 @@ import (
 )
 
 type App struct {
+	goroutineCount      int32
+	goroutineExitSignal chan struct{}
+
 	Srv *Server
 
 	PluginEnv              *pluginenv.Environment
@@ -43,7 +48,8 @@ type App struct {
 }
 
 var globalApp App = App{
-	Jobs: &jobs.JobServer{},
+	goroutineExitSignal: make(chan struct{}, 1),
+	Jobs:                &jobs.JobServer{},
 }
 
 var appCount = 0
@@ -61,7 +67,8 @@ func New() *App {
 			panic("Only one App should exist at a time. Did you forget to call Shutdown()?")
 		}
 		app := &App{
-			Jobs: &jobs.JobServer{},
+			goroutineExitSignal: make(chan struct{}, 1),
+			Jobs:                &jobs.JobServer{},
 		}
 		app.initEnterprise()
 		return app
@@ -76,12 +83,19 @@ func New() *App {
 func (a *App) Shutdown() {
 	appCount--
 	if appCount == 0 {
-		// XXX: This is to give all of our runaway goroutines time to complete.
-		//      We should wrangle them up and remove this.
-		time.Sleep(time.Second)
-
 		if a.Srv != nil {
-			a.StopServer()
+			l4g.Info(utils.T("api.server.stop_server.stopping.info"))
+
+			a.Srv.GracefulServer.Stop(TIME_TO_WAIT_FOR_CONNECTIONS_TO_CLOSE_ON_SERVER_SHUTDOWN)
+			a.Srv.Store.Close()
+			a.HubStop()
+
+			a.ShutDownPlugins()
+			a.WaitForGoroutines()
+
+			a.Srv = nil
+
+			l4g.Info(utils.T("api.server.stop_server.stopped.info"))
 		}
 	}
 }
@@ -211,6 +225,29 @@ func (a *App) Config() *model.Config {
 	return utils.Cfg
 }
 
+// Go creates a goroutine, but maintains a record of it to ensure that execution completes before
+// the app is destroyed.
+func (a *App) Go(f func()) {
+	atomic.AddInt32(&a.goroutineCount, 1)
+
+	go func() {
+		f()
+
+		atomic.AddInt32(&a.goroutineCount, -1)
+		select {
+		case a.goroutineExitSignal <- struct{}{}:
+		default:
+		}
+	}()
+}
+
+// WaitForGoroutines blocks until all goroutines created by App.Go exit.
+func (a *App) WaitForGoroutines() {
+	for atomic.LoadInt32(&a.goroutineCount) != 0 {
+		<-a.goroutineExitSignal
+	}
+}
+
 func CloseBody(r *http.Response) {
 	if r.Body != nil {
 		ioutil.ReadAll(r.Body)
-- 
cgit v1.2.3-1-g7c22