Serialize copy-on-write marks during Clone

cow_clone_race_test.go

@@ -0,0 +1,49 @@
+package roaring
+
+import (
+	"sync"
+	"testing"
+)
+
+func TestConcurrentCloneOfCOWBitmap(t *testing.T) {
+	src := New()
+	for i := uint32(0); i < 1<<20; i += 977 {
+		src.Add(i)
+	}
+	src.SetCopyOnWrite(true)
+	card := src.GetCardinality()
+
+	const (
+		goroutines = 8
+		iterations = 100
+	)
+	start := make(chan struct{})
+	var wg sync.WaitGroup
+	for g := 0; g < goroutines; g++ {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			<-start
+			for i := 0; i < iterations; i++ {
+				_ = src.Clone()
+			}
+		}()
+	}
+	close(start)
+	wg.Wait()
+
+	if src.GetCardinality() != card {
+		t.Fatalf("cardinality changed: got %d, want %d", src.GetCardinality(), card)
+	}
+	clone := src.Clone()
+	if clone.Contains(1) {
+		t.Fatal("1 was already present")
+	}
+	clone.Add(1)
+	if src.Contains(1) {
+		t.Fatal("writing a clone changed the source")
+	}
+	if !clone.Contains(1) {
+		t.Fatal("clone did not keep its own write")
+	}
+}

parallel.go

@@ -161,11 +161,10 @@ func appenderRoutine(bitmapChan chan<- *Bitmap, resultChan <-chan keyedContainer
 		}
 	}
 	answer := &Bitmap{
-		roaringArray{
-			make([]uint16, 0, expectedKeys),
-			make([]container, 0, expectedKeys),
-			make([]bool, 0, expectedKeys),
-			false,
+		highlowcontainer: roaringArray{
+			keys:            make([]uint16, 0, expectedKeys),
+			containers:      make([]container, 0, expectedKeys),
+			needCopyOnWrite: make([]bool, 0, expectedKeys),
 		},
 	}
 	for i := range keys {

roaring.go

@@ -2174,6 +2174,9 @@ func Flip(bm *Bitmap, rangeStart, rangeEnd uint64) *Bitmap {
 // (copy-on-write requires extra care in a threaded context).
 // Calling SetCopyOnWrite(true) on a bitmap created with FromBuffer is unsafe.
 func (rb *Bitmap) SetCopyOnWrite(val bool) {
+	if val {
+		rb.highlowcontainer.ensureCowMu()
+	}
 	rb.highlowcontainer.copyOnWrite = val
 }
 

roaring64/cow_clone_race_test.go

@@ -0,0 +1,42 @@
+package roaring64
+
+import (
+	"sync"
+	"testing"
+)
+
+func TestConcurrentCloneOfCOWBitmap(t *testing.T) {
+	src := New()
+	for i := uint64(0); i < 8; i++ {
+		src.Add(i << 32)
+	}
+	src.SetCopyOnWrite(true)
+	card := src.GetCardinality()
+
+	start := make(chan struct{})
+	var wg sync.WaitGroup
+	for g := 0; g < 8; g++ {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			<-start
+			for i := 0; i < 100; i++ {
+				_ = src.Clone()
+			}
+		}()
+	}
+	close(start)
+	wg.Wait()
+
+	if src.GetCardinality() != card {
+		t.Fatalf("cardinality changed: got %d, want %d", src.GetCardinality(), card)
+	}
+	clone := src.Clone()
+	clone.Add(1)
+	if src.Contains(1) {
+		t.Fatal("writing a clone changed the source")
+	}
+	if !clone.Contains(1) {
+		t.Fatal("clone did not keep its own write")
+	}
+}

roaring64/roaring64.go

@@ -1185,6 +1185,9 @@ func Flip(rb *Bitmap, rangeStart, rangeEnd uint64) *Bitmap {
 // (copy-on-write requires extra care in a threaded context).
 // Calling SetCopyOnWrite(true) on a bitmap created with FromBuffer is unsafe.
 func (rb *Bitmap) SetCopyOnWrite(val bool) {
+	if val {
+		rb.highlowcontainer.ensureCowMu()
+	}
 	rb.highlowcontainer.copyOnWrite = val
 }
 

roaring64/roaringarray64.go

@@ -2,6 +2,7 @@ package roaring64
 
 import (
 	"errors"
+	"sync"
 
 	"github.com/RoaringBitmap/roaring/v2"
 )
@@ -11,6 +12,14 @@ type roaringArray64 struct {
 	containers      []*roaring.Bitmap
 	needCopyOnWrite []bool
 	copyOnWrite     bool
+	// cowMu guards writes to needCopyOnWrite during Clone.
+	cowMu *sync.Mutex
+}
+
+func (ra *roaringArray64) ensureCowMu() {
+	if ra.cowMu == nil {
+		ra.cowMu = &sync.Mutex{}
+	}
 }
 
 var (
@@ -160,7 +169,11 @@ func (ra *roaringArray64) clone() *roaringArray64 {
 		copy(sa.containers, ra.containers)
 		sa.needCopyOnWrite = make([]bool, len(ra.needCopyOnWrite))
 
+		ra.ensureCowMu()
+		ra.cowMu.Lock()
 		ra.markAllAsNeedingCopyOnWrite()
+		ra.cowMu.Unlock()
+		sa.ensureCowMu()
 		sa.markAllAsNeedingCopyOnWrite()
 
 		// sa.needCopyOnWrite is shared

roaring64/serialization_portable.go

@@ -95,6 +95,9 @@ func (rb *Bitmap) ReadPortableFrom(stream io.Reader) (p int64, err error) {
 		hlc.appendContainer(key, c, false)
 	}
 	hlc.copyOnWrite = rb.highlowcontainer.copyOnWrite
+	if hlc.copyOnWrite {
+		hlc.ensureCowMu()
+	}
 	rb.highlowcontainer = hlc
 	return p, nil
 }

roaringarray.go

@@ -6,6 +6,7 @@ import (
 	"errors"
 	"fmt"
 	"io"
+	"sync"
 
 	"github.com/RoaringBitmap/roaring/v2/internal"
 )
@@ -123,6 +124,15 @@ type roaringArray struct {
 	containers      []container `msg:"-"` // don't try to serialize directly.
 	needCopyOnWrite []bool
 	copyOnWrite     bool
+	// cowMu guards writes to needCopyOnWrite during Clone. Concurrent
+	// clones of one copy-on-write bitmap otherwise write that slice together.
+	cowMu *sync.Mutex `msg:"-"`
+}
+
+func (ra *roaringArray) ensureCowMu() {
+	if ra.cowMu == nil {
+		ra.cowMu = &sync.Mutex{}
+	}
 }
 
 func newRoaringArray() *roaringArray {
@@ -267,7 +277,11 @@ func (ra *roaringArray) clone() *roaringArray {
 		copy(sa.containers, ra.containers)
 		sa.needCopyOnWrite = make([]bool, len(ra.needCopyOnWrite))
 
+		ra.ensureCowMu()
+		ra.cowMu.Lock()
 		ra.markAllAsNeedingCopyOnWrite()
+		ra.cowMu.Unlock()
+		sa.ensureCowMu()
 		sa.markAllAsNeedingCopyOnWrite()
 
 		// sa.needCopyOnWrite is shared

serialization_littleendian.go

@@ -379,6 +379,7 @@ func (ra *roaringArray) frozenView(buf []byte) error {
 	ra.containers = containers
 	ra.needCopyOnWrite = needCOW
 	ra.copyOnWrite = true
+	ra.ensureCowMu()
 
 	return nil
 }