From 81d3012650df0f4c590d196d3df5f390a8bdfcda Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Daniel=20Garc=C3=ADa?= <dgarcia@infoblox.com>
Date: Wed, 16 Aug 2023 01:06:41 -0500
Subject: [PATCH] acquire writelock when updating DSN

---
 dsnexec/pkg/dsnexec/dsnexec.go       | 10 ++++---
 dsnexec/pkg/dsnexec/dsnexec_test.go  | 44 ++++++++++++++++++++++++++++
 dsnexec/pkg/dsnexec/nulldb/nulldb.go | 23 +++++++++++++++
 3 files changed, 73 insertions(+), 4 deletions(-)
 create mode 100644 dsnexec/pkg/dsnexec/nulldb/nulldb.go

diff --git a/dsnexec/pkg/dsnexec/dsnexec.go b/dsnexec/pkg/dsnexec/dsnexec.go
index e9f07d68..a015066c 100644
--- a/dsnexec/pkg/dsnexec/dsnexec.go
+++ b/dsnexec/pkg/dsnexec/dsnexec.go
@@ -55,15 +55,17 @@ func (w *Handler) Exec() error {
 
 // UpdateDSN updates the dsn for a source.
 func (w *Handler) UpdateDSN(path, content string) error {
-	w.l.RLock()
-	defer w.l.RUnlock()
-
-	// update the local state
+	// acquire a write lock and update the config
+	w.l.Lock()
 	w.config.Sources[path] = DBConnInfo{
 		Driver: w.config.Sources[path].Driver,
 		DSN:    content,
 	}
+	w.l.Unlock()
 
+	// acquire a read lock and execute
+	w.l.RLock()
+	defer w.l.RUnlock()
 	if err := w.exec(); err != nil {
 		return fmt.Errorf("failed initial execute: %v", err)
 	}
diff --git a/dsnexec/pkg/dsnexec/dsnexec_test.go b/dsnexec/pkg/dsnexec/dsnexec_test.go
index 8e2a42e5..a7780802 100644
--- a/dsnexec/pkg/dsnexec/dsnexec_test.go
+++ b/dsnexec/pkg/dsnexec/dsnexec_test.go
@@ -5,11 +5,55 @@ import (
 	"database/sql/driver"
 	"fmt"
 	"reflect"
+	"sync"
 	"testing"
 
+	"github.com/google/uuid"
 	"github.com/infobloxopen/db-controller/dsnexec/pkg/tdb"
 )
 
+func TestHandler_UpdateDSN_Concurrency(t *testing.T) {
+	w := &Handler{
+		config: Config{
+			Sources: map[string]DBConnInfo{
+				"test": {
+					Driver: "postgres",
+					DSN:    "postgres://user:pass@localhost:5432/dbname?sslmode=disable",
+				},
+			},
+			Commands: []Command{
+				{
+					Command: "select 1",
+				},
+			},
+		},
+	}
+
+	// setup the destination DSN to use the test driver
+	w.config.Destination = DBConnInfo{
+		Driver: "nulldb",
+		DSN:    fmt.Sprintf("nulldb://%s", t.Name()),
+	}
+
+	// launch 100 goroutines to update the DSN concurrently
+	c := make(chan struct{})
+	wg := sync.WaitGroup{}
+	for i := 0; i < 100; i++ {
+		wg.Add(1)
+		go func() {
+			<-c
+			for j := 0; j < 100; j++ {
+				dsn := fmt.Sprintf("postgres://user:%s@localhost:5432/dbname?sslmode=disable", uuid.New())
+				w.UpdateDSN("test", dsn)
+			}
+			wg.Done()
+		}()
+	}
+	// launch the goroutines
+	close(c)
+	wg.Wait()
+}
+
 func TestHandler_UpdateDSN(t *testing.T) {
 	type fields struct {
 		config Config
diff --git a/dsnexec/pkg/dsnexec/nulldb/nulldb.go b/dsnexec/pkg/dsnexec/nulldb/nulldb.go
new file mode 100644
index 00000000..6cdab1c9
--- /dev/null
+++ b/dsnexec/pkg/dsnexec/nulldb/nulldb.go
@@ -0,0 +1,23 @@
+package nulldb
+
+import (
+	"database/sql"
+	"database/sql/driver"
+	"fmt"
+)
+
+var (
+	defaultDriver *d
+)
+
+type d struct {
+}
+
+func init() {
+	defaultDriver = &d{}
+	sql.Register("nulldb", defaultDriver)
+}
+
+func (d *d) Open(name string) (driver.Conn, error) {
+	return nil, fmt.Errorf("unsupported Open in nulldb")
+}