syncmap

Author: FGasper

internal/verifier/compare.go

@@ -6,6 +6,7 @@ import (
 	"time"
 
 	"github.com/10gen/migration-verifier/internal/types"
+	"github.com/10gen/migration-verifier/syncmap"
 	"github.com/pkg/errors"
 	"go.mongodb.org/mongo-driver/bson"
 	"go.mongodb.org/mongo-driver/mongo"
@@ -15,6 +16,8 @@ import (
 
 const readTimeout = 10 * time.Minute
 
+type docCacheMap = syncmap.SyncMap[string, bson.Raw]
+
 func (verifier *Verifier) FetchAndCompareDocuments(
 	givenCtx context.Context,
 	task *VerificationTask,
@@ -73,8 +76,8 @@ func (verifier *Verifier) compareDocsFromChannels(
 
 	namespace := task.QueryFilter.Namespace
 
-	srcCache := map[string]bson.Raw{}
-	dstCache := map[string]bson.Raw{}
+	srcCache := &docCacheMap{}
+	dstCache := &docCacheMap{}
 
 	// This is the core document-handling logic. It either:
 	//
@@ -84,7 +87,7 @@ func (verifier *Verifier) compareDocsFromChannels(
 	handleNewDoc := func(doc bson.Raw, isSrc bool) error {
 		mapKey := getMapKey(doc, mapKeyFieldNames)
 
-		var ourMap, theirMap map[string]bson.Raw
+		var ourMap, theirMap *docCacheMap
 
 		if isSrc {
 			ourMap = srcCache
@@ -95,22 +98,22 @@ func (verifier *Verifier) compareDocsFromChannels(
 		}
 		// See if we've already cached a document with this
 		// mapKey from the other channel.
-		theirDoc, exists := theirMap[mapKey]
+		theirDoc, exists := theirMap.Load(mapKey)
 
 		// If there is no such cached document, then cache the newly-received
 		// document in our map then proceed to the next document.
 		//
 		// (We'll remove the cache entry when/if the other channel yields a
 		// document with the same mapKey.)
 		if !exists {
-			ourMap[mapKey] = doc
+			ourMap.Store(mapKey, doc)
 			return nil
 		}
 
 		// We have two documents! First we remove the cache entry. This saves
 		// memory, but more importantly, it lets us know, once we exhaust the
 		// channels, which documents were missing on one side or the other.
-		delete(theirMap, mapKey)
+		theirMap.Delete(mapKey)
 
 		// Now we determine which document came from whom.
 		var srcDoc, dstDoc bson.Raw
@@ -227,9 +230,9 @@ func (verifier *Verifier) compareDocsFromChannels(
 	// missing on the other side. We add results for those.
 
 	// We might as well pre-grow the slice:
-	results = slices.Grow(results, len(srcCache)+len(dstCache))
+	results = slices.Grow(results, srcCache.Len()+dstCache.Len())
 
-	for _, doc := range srcCache {
+	srcCache.Range(func(_ string, doc bson.Raw) bool {
 		results = append(
 			results,
 			VerificationResult{
@@ -240,9 +243,11 @@ func (verifier *Verifier) compareDocsFromChannels(
 				dataSize:  len(doc),
 			},
 		)
-	}
 
-	for _, doc := range dstCache {
+		return true
+	})
+
+	dstCache.Range(func(_ string, doc bson.Raw) bool {
 		results = append(
 			results,
 			VerificationResult{
@@ -253,7 +258,9 @@ func (verifier *Verifier) compareDocsFromChannels(
 				dataSize:  len(doc),
 			},
 		)
-	}
+
+		return true
+	})
 
 	return results, docCount, byteCount, nil
 }

syncmap/syncmap.go

@@ -0,0 +1,87 @@
+package syncmap
+
+// This package wraps sync.Map to provide type safety via generics.
+
+// Surprisingly, as of May 2023 there doesn’t appear to be a “go-to”
+// public implementation of this. Initially we considered gopkg.in/typ.v4,
+// but this largely reimplements sync.Map, which seems unideal for
+// long-term maintenance and seems to lack wide adoption.
+// github.com/puzpuzpuz/xsync is widely used but, because it attempts
+// to iterate on sync.Map’s internals, requires map keys to be strings.
+//
+// Assumedly Go’s standard library will eventually add something like
+// this, but for now this is a small, useful enough effort for us to
+// have for ourselves.
+//
+// See this module’s tests for example usage.
+
+import (
+	"sync"
+)
+
+// SyncMap is a generics-powered replacement for sync.Map.
+type SyncMap[KT comparable, VT any] struct {
+	internalMap sync.Map
+}
+
+// Load behaves like sync.Map’s corresponding method, but
+// if the value doesn’t exist then this returns the value type’s
+// zero-value.
+func (m *SyncMap[KT, VT]) Load(key KT) (VT, bool) {
+	return m.returnLoad(m.internalMap.Load(key))
+}
+
+// Delete behaves like sync.Map’s corresponding method.
+func (m *SyncMap[KT, VT]) Delete(key KT) {
+	m.internalMap.Delete(key)
+}
+
+// Store behaves like sync.Map’s corresponding method.
+func (m *SyncMap[KT, VT]) Store(key KT, value VT) {
+	m.internalMap.Store(key, value)
+}
+
+// LoadAndDelete behaves like sync.Map’s corresponding method.
+func (m *SyncMap[KT, VT]) LoadAndDelete(key KT) (VT, bool) {
+	return m.returnLoad(m.internalMap.LoadAndDelete(key))
+}
+
+// LoadOrStore behaves like sync.Map’s corresponding method.
+func (m *SyncMap[KT, VT]) LoadOrStore(key KT, value VT) (VT, bool) {
+	actual, loaded := m.internalMap.LoadOrStore(key, value)
+
+	return actual.(VT), loaded
+}
+
+// Range behaves like sync.Map’s corresponding method.
+func (m *SyncMap[KT, VT]) Range(f func(KT, VT) bool) {
+	m.internalMap.Range(func(key, value any) bool {
+		return f(key.(KT), value.(VT))
+	})
+}
+
+// ----------------------------------------------------------------------
+
+// Len is a convenience method that returns the number of elements
+// the map stores. It has no counterpart in sync.Map.
+// This method is O(n); i.e., larger maps take longer.
+func (m *SyncMap[KT, VT]) Len() int {
+	mapLen := 0
+
+	m.internalMap.Range(func(_, _ any) bool {
+		mapLen++
+		return true
+	})
+
+	return mapLen
+}
+
+// ----------------------------------------------------------------------
+
+func (_ *SyncMap[KT, VT]) returnLoad(val any, loaded bool) (VT, bool) {
+	if !loaded {
+		return *new(VT), false
+	}
+
+	return val.(VT), true
+}

syncmap/syncmap_test.go

@@ -0,0 +1,104 @@
+package syncmap
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/suite"
+)
+
+type mySuite struct {
+	suite.Suite
+}
+
+func TestUnitTestSuite(t *testing.T) {
+	suite.Run(t, &mySuite{})
+}
+
+func (s *mySuite) TestSyncMap_NonexistentString() {
+	strMap := SyncMap[string, string]{}
+
+	nadaStr, myOk := strMap.Load("foo")
+	s.Assert().Empty(nadaStr, "Load() of string when key is nonexistent")
+	s.Assert().False(myOk, "Load() indicates nonexistence")
+}
+
+func (s *mySuite) TestSyncMap_Range() {
+	keys := []int32{}
+	vals := []int64{}
+	otherMap := SyncMap[int32, int64]{}
+	otherMap.Store(1, 11)
+	otherMap.Store(2, 22)
+	otherMap.Store(3, 33)
+
+	otherMap.Range(func(k int32, v int64) bool {
+		keys = append(keys, k)
+		vals = append(vals, v)
+		return true
+	})
+
+	s.Assert().ElementsMatch(
+		[]int32{1, 2, 3},
+		keys,
+		"Range(): keys as expected",
+	)
+
+	s.Assert().ElementsMatch(
+		[]int64{11, 22, 33},
+		vals,
+		"Range(): values as expected",
+	)
+}
+
+func (s *mySuite) TestSyncMap_Len() {
+	theMap := SyncMap[string, []byte]{}
+
+	theMap.Store("foo", []byte{})
+	theMap.Store("bar", []byte{})
+	theMap.Store("baz", []byte{})
+	theMap.Store("qux", []byte{})
+
+	s.Assert().Equal(4, theMap.Len(), "Len() gives the size")
+}
+
+func (s *mySuite) TestSyncMap_LoadOrStore() {
+	theMap := SyncMap[string, []byte]{}
+
+	theVal, ok := theMap.LoadOrStore("foo", []byte{2})
+	s.Assert().Equal([]byte{2}, theVal, "LoadOrStore() returned new value if no old one exists")
+	s.Assert().False(ok, "LoadOrStore() indicates nonexistence")
+
+	theVal, ok = theMap.LoadOrStore("foo", []byte{3})
+	s.Assert().Equal([]byte{2}, theVal, "LoadOrStore() returned old value")
+	s.Assert().True(ok, "LoadOrStore() indicates existence")
+}
+
+func (s *mySuite) TestSyncMap_LoadAndDelete() {
+	theMap := SyncMap[string, []byte]{}
+
+	theVal, ok := theMap.LoadAndDelete("foo")
+	s.Assert().Nil(theVal, "LoadAndDelete() returns nil on nonexistence")
+	s.Assert().False(ok, "boolean indicates nonexistence")
+
+	theMap.Store("foo", []byte{1})
+
+	theVal, ok = theMap.LoadAndDelete("foo")
+	s.Assert().Equal([]byte{1}, theVal, "LoadAndDelete() returned Store()d value")
+	s.Assert().True(ok, "boolean indicates existence")
+}
+
+func (s *mySuite) TestSyncMap_Basics() {
+	theMap := SyncMap[string, []byte]{}
+
+	nada, ok := theMap.Load("foo")
+	s.Assert().Nil(nada, "Load() of array when key is nonexistent")
+	s.Assert().False(ok, "Load() indicates nonexistence")
+
+	theMap.Store("foo", []byte{1})
+	theVal, ok := theMap.Load("foo")
+	s.Assert().Equal([]byte{1}, theVal, "Load() returned Store()d value")
+	s.Assert().True(ok, "Load() indicates existence")
+
+	theMap.Delete("foo")
+	_, ok = theMap.Load("foo")
+	s.Assert().False(ok, "Load() indicates nonexistence after Delete()")
+}