Making GenerateIdSuffix thread safe (#867)

* Making GenerateIdSuffix thread safe

* Adding comments

* Correcting test

* Correcting test

internal/helpers.go

@@ -18,29 +18,34 @@ import (
 	"fmt"
 	"strconv"
 	"strings"
+	"sync"
 
 	"github.com/GoogleCloudPlatform/spanner-migration-tool/schema"
 	"github.com/GoogleCloudPlatform/spanner-migration-tool/spanner/ddl"
 )
 
 type Counter struct {
+	counterMutex sync.Mutex
 	ObjectId string
 }
 
 var Cntr Counter
 
-func GenerateIdSuffix() string {
-
-	counter, _ := strconv.Atoi(Cntr.ObjectId)
+// Thread safe Counter to generate ids in session file.
+func (c *Counter) GenerateIdSuffix() string {
+	c.counterMutex.Lock()
+	counter, _ := strconv.Atoi(c.ObjectId)
 
 	counter = counter + 1
 
-	Cntr.ObjectId = strconv.Itoa(counter)
-	return Cntr.ObjectId
+	c.ObjectId = strconv.Itoa(counter)
+	returnVal := c.ObjectId
+	c.counterMutex.Unlock()
+	return returnVal
 }
 
 func GenerateId(idPrefix string) string {
-	idSuffix := GenerateIdSuffix()
+	idSuffix := Cntr.GenerateIdSuffix()
 	id := idPrefix + idSuffix
 	return id
 }

internal/helpers_test.go

@@ -0,0 +1,51 @@
+// Copyright 2024 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package internal
+
+import (
+	"sync"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestGenerateIdSuffix(t *testing.T) {
+	var wg sync.WaitGroup
+	tests := []struct {
+		number     int
+		expected string
+	}{
+		{10000, "10000"},
+		{20000, "30000"},
+		{10000, "40000"},
+	}
+	counter := Counter{}
+	for _, tc := range tests {
+		// Call GenerateIdSuffix n number of times parallely. 
+		for i := 0; i < tc.number; i++ {
+			// Increment the WaitGroup counter.
+			wg.Add(1)
+			go func() {
+				counter.GenerateIdSuffix()
+				// Decrement the counter when the goroutine completes.
+				defer wg.Done()
+			}()
+		}
+		// Wait for all Go routines in the tc to complete
+        wg.Wait()
+		// Assert that the counter is actually incremented to n.
+		assert.Equal(t, tc.expected, counter.ObjectId)
+	}
+}