From 3f01129ddf798c4f2d5d915cb834e362a94c0632 Mon Sep 17 00:00:00 2001
From: Agniva De Sarker <agnivade@yahoo.co.in>
Date: Tue, 10 Aug 2021 10:07:35 +0530
Subject: [PATCH] MM-37165: Fix improper plugin shutdown (#18044)

* MM-37165: Fix improper plugin shutdown

This was caught from a race test failure. While the failure manifested due to a log being
written from a test after the test exited, the real reason was hidden further deeper.

What was happening is that the server would always listen for plugin requests in a
separate goroutine via `g.muxBroker.AcceptAndServe` in the `OnActivate` hook. But the
plugin shutdown process would just close the plugin connections and move on, leading
to the classic case of improperly shut down goroutines.

When this happened, an opportunity opens up in a way that the server
would still be executing a request whereas the main goroutine and therefore the parent
test has already finished. This would lead to an error like

```
{"level":"error","ts":1626451258.4141135,"caller":"mlog/sugar.go:25","msg":"pluginAPI scheduleOnce poller encountered an error but is still polling","plugin_id":"com.mattermost.plugin-incident-management","error":"ListPluginKeys: Unable to list all the plugin keys., failed to get PluginKeyValues with pluginId=com.mattermost.plugin-incident-management: sql: database is closed
```

And now, this finally calls `mlog.Error`, which finally triggers our race condition :)

To fix this, we use some basic synchronization via waitgroups and just wait for it
to finish after closing the plugin process.

https://mattermost.atlassian.net/browse/MM-37165

```release-note
NONE
```

* gofmt

```release-note
NONE
```

* split waitgroup additions

```release-note
NONE
```
---
 plugin/client_rpc.go | 24 +++++++++++++++++-------
 plugin/supervisor.go | 11 +++++++++++
 2 files changed, 28 insertions(+), 7 deletions(-)

diff --git a/plugin/client_rpc.go b/plugin/client_rpc.go
index 7f04dc406b..5446b40425 100644
--- a/plugin/client_rpc.go
+++ b/plugin/client_rpc.go
@@ -19,6 +19,7 @@ import (
 	"net/rpc"
 	"os"
 	"reflect"
+	"sync"
 
 	"github.com/dyatlov/go-opengraph/opengraph"
 	"github.com/go-sql-driver/mysql"
@@ -38,6 +39,7 @@ type hooksRPCClient struct {
 	apiImpl     API
 	driver      Driver
 	implemented [TotalHooksID]bool
+	doneWg      sync.WaitGroup
 }
 
 type hooksRPCServer struct {
@@ -240,15 +242,23 @@ type Z_OnActivateReturns struct {
 
 func (g *hooksRPCClient) OnActivate() error {
 	muxId := g.muxBroker.NextId()
-	go g.muxBroker.AcceptAndServe(muxId, &apiRPCServer{
-		impl:      g.apiImpl,
-		muxBroker: g.muxBroker,
-	})
+	g.doneWg.Add(1)
+	go func() {
+		defer g.doneWg.Done()
+		g.muxBroker.AcceptAndServe(muxId, &apiRPCServer{
+			impl:      g.apiImpl,
+			muxBroker: g.muxBroker,
+		})
+	}()
 
 	nextID := g.muxBroker.NextId()
-	go g.muxBroker.AcceptAndServe(nextID, &dbRPCServer{
-		dbImpl: g.driver,
-	})
+	g.doneWg.Add(1)
+	go func() {
+		defer g.doneWg.Done()
+		g.muxBroker.AcceptAndServe(nextID, &dbRPCServer{
+			dbImpl: g.driver,
+		})
+	}()
 
 	_args := &Z_OnActivateArgs{
 		APIMuxId:    muxId,
diff --git a/plugin/supervisor.go b/plugin/supervisor.go
index 1da4136a73..883321fbcc 100644
--- a/plugin/supervisor.go
+++ b/plugin/supervisor.go
@@ -25,6 +25,7 @@ type supervisor struct {
 	hooks       Hooks
 	implemented [TotalHooksID]bool
 	pid         int
+	hooksClient *hooksRPCClient
 }
 
 func newSupervisor(pluginInfo *model.BundleInfo, apiImpl API, driver Driver, parentLogger *mlog.Logger, metrics einterfaces.MetricsInterface) (retSupervisor *supervisor, retErr error) {
@@ -83,6 +84,11 @@ func newSupervisor(pluginInfo *model.BundleInfo, apiImpl API, driver Driver, par
 		return nil, err
 	}
 
+	c, ok := raw.(*hooksRPCClient)
+	if ok {
+		sup.hooksClient = c
+	}
+
 	sup.hooks = &hooksTimerLayer{pluginInfo.Manifest.Id, raw.(Hooks), metrics}
 
 	impl, err := sup.hooks.Implemented()
@@ -104,6 +110,11 @@ func (sup *supervisor) Shutdown() {
 	if sup.client != nil {
 		sup.client.Kill()
 	}
+
+	// Wait for API RPC server and DB RPC server to exit.
+	if sup.hooksClient != nil {
+		sup.hooksClient.doneWg.Wait()
+	}
 }
 
 func (sup *supervisor) Hooks() Hooks {