Fix a DB connection leak on restart (#521)

* Fix a DB connection leak on restart * Properly shutting down the store if OnActivate errors * Fixing CI * Fixing CI * Update server/plugin.go Co-authored-by: Doug Lauder <[email protected]> * Addressing PR review comments * Fixing CI * Reorganizing to handled properly DB connection shutdown * Fixing CI * improving the predictability of the e2e tests * Reduce a bit of repeated code * Wait more for the plugin to be ready --------- Co-authored-by: Doug Lauder <[email protected]>
mattermost · Mar 4, 2024 · 7c615a4 · 7c615a4
1 parent eabf1ab
commit 7c615a4
Show file tree

Hide file tree

Showing 7 changed files with 69 additions and 19 deletions.
diff --git a/server/ce2e/plugin_test.go b/server/ce2e/plugin_test.go
@@ -45,6 +45,11 @@ func TestMessageHasBeenPostedNewMessageE2E(t *testing.T) {
 	err = store.SetUserInfo(user.Id, "ms-user-id", &fakeToken)
 	require.NoError(t, err)
 
+	require.Eventually(t, func() bool {
+		suggestions, _, _ := client.ListCommandAutocompleteSuggestions(context.Background(), "/msteams", team.Id)
+		return len(suggestions) > 0
+	}, 10*time.Second, 500*time.Millisecond)
+
 	t.Run("Without Channel Link", func(t *testing.T) {
 		var newPost *model.Post
 		newPost, _, err = client.CreatePost(context.Background(), &post)
@@ -100,10 +105,7 @@ func TestMessageHasBeenPostedNewMessageE2E(t *testing.T) {
 		})
 		require.NoError(t, err)
 
-		require.Eventually(t, func() bool {
-			_, _, err = client.ExecuteCommand(context.Background(), channel.Id, "/msteams link ms-team-id ms-channel-id")
-			return err == nil
-		}, 5*time.Second, 500*time.Millisecond)
+		_, _, err = client.ExecuteCommand(context.Background(), channel.Id, "/msteams link ms-team-id ms-channel-id")
 
 		var newPost *model.Post
 		newPost, _, err = client.CreatePost(context.Background(), &post)
@@ -139,10 +141,7 @@ func TestMessageHasBeenPostedNewMessageE2E(t *testing.T) {
 		err = mockClient.MockError(http.MethodPost, "/v1.0/teams/ms-team-id/channels/ms-channel-id/messages")
 		require.NoError(t, err)
 
-		require.Eventually(t, func() bool {
-			_, _, err = client.ExecuteCommand(context.Background(), channel.Id, "/msteams link ms-team-id ms-channel-id")
-			return err == nil
-		}, 5*time.Second, 500*time.Millisecond)
+		_, _, err = client.ExecuteCommand(context.Background(), channel.Id, "/msteams link ms-team-id ms-channel-id")
 
 		newPost, _, err := client.CreatePost(context.Background(), &post)
 		require.NoError(t, err)

diff --git a/server/plugin.go b/server/plugin.go
@@ -395,6 +395,12 @@ func (p *Plugin) stop(isRestart bool) {
 			p.API.LogError("failed to close metrics job", "error", err)
 		}
 	}
+
+	if !isRestart {
+		if err := p.store.Shutdown(); err != nil {
+			p.API.LogError("failed to close db connection", "error", err)
+		}
+	}
 }
 
 func (p *Plugin) restart() {
@@ -434,7 +440,7 @@ func (p *Plugin) generatePluginSecrets() error {
 	return nil
 }
 
-func (p *Plugin) OnActivate() error {
+func (p *Plugin) onActivate() error {
 	if p.clientBuilderWithToken == nil {
 		p.clientBuilderWithToken = msteams.NewTokenClient
 	}
@@ -459,27 +465,24 @@ func (p *Plugin) OnActivate() error {
 
 	p.activityHandler = handlers.New(p)
 
-	subscriptionsClusterMutex, err := cluster.NewMutex(p.API, subscriptionsClusterMutexKey)
+	p.subscriptionsClusterMutex, err = cluster.NewMutex(p.API, subscriptionsClusterMutexKey)
 	if err != nil {
 		return err
 	}
-	p.subscriptionsClusterMutex = subscriptionsClusterMutex
 
-	whitelistClusterMutex, err := cluster.NewMutex(p.API, whitelistClusterMutexKey)
+	p.whitelistClusterMutex, err = cluster.NewMutex(p.API, whitelistClusterMutexKey)
 	if err != nil {
 		return err
 	}
-	p.whitelistClusterMutex = whitelistClusterMutex
 
-	botID, err := p.apiClient.Bot.EnsureBot(&model.Bot{
+	p.userID, err = p.apiClient.Bot.EnsureBot(&model.Bot{
 		Username:    botUsername,
 		DisplayName: botDisplayName,
 		Description: "Created by the MS Teams Sync plugin.",
 	}, pluginapi.ProfileImagePath("assets/icon.png"))
 	if err != nil {
 		return err
 	}
-	p.userID = botID
 
 	if p.store == nil {
 		if p.apiClient.Store.DriverName() != model.DatabaseDriverPostgres {
@@ -498,6 +501,7 @@ func (p *Plugin) OnActivate() error {
 			func() []byte { return []byte(p.configuration.EncryptionKey) },
 		)
 		p.store = timerlayer.New(store, p.GetMetrics())
+
 		if err = p.store.Init(); err != nil {
 			return err
 		}
@@ -507,7 +511,7 @@ func (p *Plugin) OnActivate() error {
 		remoteID, err := p.API.RegisterPluginForSharedChannels(model.RegisterPluginOpts{
 			Displayname:  pluginID,
 			PluginID:     pluginID,
-			CreatorID:    botID,
+			CreatorID:    p.userID,
 			AutoShareDMs: true,
 			AutoInvited:  true,
 		})
@@ -529,7 +533,7 @@ func (p *Plugin) OnActivate() error {
 				TeamId:    linkedChannel.MattermostTeamID,
 				Home:      true,
 				ReadOnly:  false,
-				CreatorId: botID,
+				CreatorId: p.userID,
 				RemoteId:  p.remoteID,
 				ShareName: linkedChannel.MattermostChannelID,
 			})
@@ -561,15 +565,28 @@ func (p *Plugin) OnActivate() error {
 
 		p.whitelistClusterMutex.Lock()
 		defer p.whitelistClusterMutex.Unlock()
-		if err := p.store.PrefillWhitelist(); err != nil {
-			p.API.LogWarn("Error in populating the whitelist with already connected users", "error", err.Error())
+		if err2 := p.store.PrefillWhitelist(); err2 != nil {
+			p.API.LogWarn("Error in populating the whitelist with already connected users", "error", err2.Error())
 		}
 	}()
 
 	go p.start(false)
 	return nil
 }
 
+func (p *Plugin) OnActivate() error {
+	if err := p.onActivate(); err != nil {
+		p.API.LogWarn("error activating the plugin", "error", err)
+		if p.store != nil {
+			if err = p.store.Shutdown(); err != nil {
+				p.API.LogWarn("failed to close db connection", "error", err)
+			}
+		}
+		return err
+	}
+	return nil
+}
+
 func (p *Plugin) OnDeactivate() error {
 	p.stop(false)
 	return nil

diff --git a/server/plugin_test.go b/server/plugin_test.go
@@ -45,6 +45,7 @@ func newTestPlugin(t *testing.T) *Plugin {
 			return clientMock
 		},
 	}
+	plugin.store.(*storemocks.Store).On("Shutdown").Return(nil)
 	plugin.store.(*storemocks.Store).Test(t)
 
 	plugin.msteamsAppClient.(*mocks.Client).On("ClearSubscriptions").Return(nil)

diff --git a/server/store/mocks/Store.go b/server/store/mocks/Store.go
diff --git a/server/store/sqlstore/store.go b/server/store/sqlstore/store.go
@@ -50,6 +50,10 @@ func New(db *sql.DB, api plugin.API, enabledTeams func() []string, encryptionKey
 	}
 }
 
+func (s *SQLStore) Shutdown() error {
+	return s.db.Close()
+}
+
 func (s *SQLStore) createTable(tableName, columnList string) error {
 	if _, err := s.db.Exec(fmt.Sprintf("CREATE TABLE IF NOT EXISTS %s (%s)", tableName, columnList)); err != nil {
 		return err

diff --git a/server/store/store.go b/server/store/store.go
@@ -58,4 +58,5 @@ type Store interface {
 	IsUserPresentInWhitelist(userID string) (bool, error)
 	UpdateSubscriptionLastActivityAt(subscriptionID string, lastActivityAt time.Time) error
 	GetSubscriptionsLastActivityAt() (map[string]time.Time, error)
+	Shutdown() error
 }
diff --git a/server/store/timerlayer/timerlayer.go b/server/store/timerlayer/timerlayer.go