fix: clone first bitmap instead of storing reference in ThreadSafeKindBitmap Or - BED-7656

graph/types.go

@@ -136,10 +136,6 @@ func NewThreadSafeKindBitmap() *ThreadSafeKindBitmap {
 	}
 }
 
-func (s ThreadSafeKindBitmap) Count(kinds ...Kind) uint64 {
-	return s.Get(kinds...).Cardinality()
-}
-
 func (s ThreadSafeKindBitmap) Get(kinds ...Kind) cardinality.Duplex[uint64] {
 	s.rwLock.RLock()
 	defer s.rwLock.RUnlock()
@@ -168,7 +164,7 @@ func (s ThreadSafeKindBitmap) Or(kind Kind, other cardinality.Duplex[uint64]) {
 	if kindBitmap, hasKind := s.bitmaps[kind.String()]; hasKind {
 		kindBitmap.Or(other)
 	} else {
-		s.bitmaps[kind.String()] = other
+		s.bitmaps[kind.String()] = other.Clone()
 	}
 }
 
@@ -206,10 +202,7 @@ func (s ThreadSafeKindBitmap) Add(kind Kind, value uint64) {
 	if kindBitmap, hasKind := s.bitmaps[kind.String()]; hasKind {
 		kindBitmap.Add(value)
 	} else {
-		kindBitmap = cardinality.NewBitmap64()
-		kindBitmap.Add(value)
-
-		s.bitmaps[kind.String()] = kindBitmap
+		s.bitmaps[kind.String()] = cardinality.NewBitmap64With(value)
 	}
 }
 
@@ -219,14 +212,10 @@ func (s ThreadSafeKindBitmap) CheckedAdd(kind Kind, value uint64) bool {
 
 	if kindBitmap, hasKind := s.bitmaps[kind.String()]; hasKind {
 		return kindBitmap.CheckedAdd(value)
-	} else {
-		kindBitmap = cardinality.NewBitmap64()
-		kindBitmap.Add(value)
-
-		s.bitmaps[kind.String()] = kindBitmap
-
-		return true
 	}
+
+	s.bitmaps[kind.String()] = cardinality.NewBitmap64With(value)
+	return true
 }
 
 // IndexedSlice is a structure maps a comparable key to a value that implements size.Sizable.

graph/types_test.go

@@ -0,0 +1,92 @@
+package graph_test
+
+import (
+	"context"
+	"math/rand"
+	"strconv"
+	"sync"
+	"testing"
+	"time"
+
+	"github.com/specterops/dawgs/cardinality"
+	"github.com/specterops/dawgs/graph"
+)
+
+func generateKinds(numKinds int) graph.Kinds {
+	var kinds graph.Kinds
+
+	for kindIdx := range numKinds {
+		kinds = kinds.Add(graph.StringKind("Kind" + strconv.Itoa(kindIdx+1)))
+	}
+
+	return kinds
+}
+
+func Test_ThreadSafeKindBitmap_ConcurrentAccess(t *testing.T) {
+	var (
+		instance = graph.NewThreadSafeKindBitmap()
+		scratch  = cardinality.NewBitmap64()
+		kinds    = generateKinds(1_00)
+		workerWG = &sync.WaitGroup{}
+	)
+
+	ctx, done := context.WithCancel(context.Background())
+	defer done()
+
+	workerWG.Add(1)
+
+	go func() {
+		defer workerWG.Done()
+
+		iteration := 0
+
+		for {
+			for range 1_000 {
+				scratch.Add(rand.Uint64() % 1_000_000)
+			}
+
+			kind := kinds[iteration%len(kinds)]
+			iteration += 1
+
+			instance.Or(kind, scratch)
+			scratch.Clear()
+
+			select {
+			case <-ctx.Done():
+				return
+
+			default:
+			}
+		}
+	}()
+
+	for range 4 {
+		workerWG.Add(1)
+
+		go func() {
+			defer workerWG.Done()
+
+			for {
+				for i := range 10_000 {
+					kind := kinds[i%len(kinds)]
+
+					for range 100_000 {
+						instance.Cardinality(kind)
+
+						select {
+						case <-ctx.Done():
+							return
+
+						default:
+						}
+					}
+				}
+			}
+		}()
+	}
+
+	time.Sleep(time.Millisecond * 250)
+	done()
+
+	workerWG.Wait()
+}