From 89f5a0deec0da908f7b13ffc92e9a2a5de925a6b Mon Sep 17 00:00:00 2001
From: Agniva De Sarker <agnivade@yahoo.co.in>
Date: Fri, 8 Mar 2024 13:11:48 +0530
Subject: [PATCH] MM-56402: Initialize replica conn pool on-demand (#26410)

Previously, we would setup both pools only when
GetMasterDB was called. This was inefficient and
would waste open connections if the replica wasn't used
at all.

We fix it to initialize the pools as they are called.

https://mattermost.atlassian.net/browse/MM-56402
```release-note
NONE
```

Co-authored-by: Jesse Hallam <jesse.hallam@gmail.com>
---
 server/public/pluginapi/store.go      | 51 +++++++++++++++------------
 server/public/pluginapi/store_test.go | 13 ++-----
 2 files changed, 32 insertions(+), 32 deletions(-)

diff --git a/server/public/pluginapi/store.go b/server/public/pluginapi/store.go
index c0daee7c80..ab0adfc4d5 100644
--- a/server/public/pluginapi/store.go
+++ b/server/public/pluginapi/store.go
@@ -15,10 +15,11 @@ import (
 
 // StoreService exposes the underlying database.
 type StoreService struct {
-	initialized bool
-	api         plugin.API
-	driver      plugin.Driver
-	mutex       sync.Mutex
+	initializedMaster  bool
+	initializedReplica bool
+	api                plugin.API
+	driver             plugin.Driver
+	mutex              sync.Mutex
 
 	masterDB  *sql.DB
 	replicaDB *sql.DB
@@ -31,7 +32,7 @@ func (s *StoreService) GetMasterDB() (*sql.DB, error) {
 	s.mutex.Lock()
 	defer s.mutex.Unlock()
 
-	if err := s.initialize(); err != nil {
+	if err := s.initializeMaster(); err != nil {
 		return nil, err
 	}
 
@@ -46,7 +47,7 @@ func (s *StoreService) GetReplicaDB() (*sql.DB, error) {
 	s.mutex.Lock()
 	defer s.mutex.Unlock()
 
-	if err := s.initialize(); err != nil {
+	if err := s.initializeReplica(); err != nil {
 		return nil, err
 	}
 
@@ -62,20 +63,18 @@ func (s *StoreService) Close() error {
 	s.mutex.Lock()
 	defer s.mutex.Unlock()
 
-	if !s.initialized {
-		return nil
-	}
-
-	if err := s.masterDB.Close(); err != nil {
-		return err
-	}
-
 	if s.replicaDB != nil {
 		if err := s.replicaDB.Close(); err != nil {
 			return err
 		}
 	}
 
+	if s.masterDB != nil {
+		if err := s.masterDB.Close(); err != nil {
+			return err
+		}
+	}
+
 	return nil
 }
 
@@ -84,8 +83,8 @@ func (s *StoreService) DriverName() string {
 	return *s.api.GetConfig().SqlSettings.DriverName
 }
 
-func (s *StoreService) initialize() error {
-	if s.initialized {
+func (s *StoreService) initializeMaster() error {
+	if s.initializedMaster {
 		return nil
 	}
 
@@ -93,25 +92,33 @@ func (s *StoreService) initialize() error {
 		return errors.New("no db driver was provided")
 	}
 
-	config := s.api.GetUnsanitizedConfig()
-
 	// Set up master db
-	db := sql.OpenDB(driver.NewConnector(s.driver, true))
+	db := sql.OpenDB(driver.NewConnector(s.driver, true /* IsMaster */))
 	if err := db.Ping(); err != nil {
 		return errors.Wrap(err, "failed to connect to master db")
 	}
 	s.masterDB = db
 
+	s.initializedMaster = true
+
+	return nil
+}
+
+func (s *StoreService) initializeReplica() error {
+	if s.initializedReplica {
+		return nil
+	}
+
+	config := s.api.GetUnsanitizedConfig()
 	// Set up replica db
 	if len(config.SqlSettings.DataSourceReplicas) > 0 {
-		db := sql.OpenDB(driver.NewConnector(s.driver, false))
+		db := sql.OpenDB(driver.NewConnector(s.driver, false /* IsMaster */))
 		if err := db.Ping(); err != nil {
 			return errors.Wrap(err, "failed to connect to replica db")
 		}
 		s.replicaDB = db
 	}
 
-	s.initialized = true
-
+	s.initializedReplica = true
 	return nil
 }
diff --git a/server/public/pluginapi/store_test.go b/server/public/pluginapi/store_test.go
index 78b2b98f47..6e7ea70b9d 100644
--- a/server/public/pluginapi/store_test.go
+++ b/server/public/pluginapi/store_test.go
@@ -12,18 +12,10 @@ import (
 
 func TestStore(t *testing.T) {
 	t.Run("master db singleton", func(t *testing.T) {
-		config := &model.Config{
-			SqlSettings: model.SqlSettings{
-				DriverName: model.NewString("test"),
-				DataSource: model.NewString("TestStore-master-db"),
-			},
-		}
-
 		api := &plugintest.API{}
-		defer api.AssertExpectations(t)
-		api.On("GetUnsanitizedConfig").Return(config)
 
 		driver := &plugintest.Driver{}
+		defer driver.AssertExpectations(t)
 		driver.On("Conn", true).Return("test", nil)
 		driver.On("ConnPing", "test").Return(nil)
 		driver.On("ConnClose", "test").Return(nil)
@@ -52,6 +44,7 @@ func TestStore(t *testing.T) {
 		}
 
 		driver := &plugintest.Driver{}
+		defer driver.AssertExpectations(t)
 		driver.On("Conn", true).Return("test", nil)
 		driver.On("ConnPing", "test").Return(nil)
 		driver.On("ConnClose", "test").Return(nil)
@@ -88,7 +81,7 @@ func TestStore(t *testing.T) {
 		api.On("GetUnsanitizedConfig").Return(config)
 
 		driver := &plugintest.Driver{}
-		driver.On("Conn", true).Return("test", nil)
+		defer driver.AssertExpectations(t)
 		driver.On("Conn", false).Return("test", nil)
 		driver.On("ConnPing", "test").Return(nil)
 		driver.On("ConnClose", "test").Return(nil)