Refactoring Bytes to check double free and code cleanup

pkg/fileservice/bytes.go

@@ -25,16 +25,17 @@ import (
 type Bytes struct {
 	bytes       []byte
 	deallocator malloc.Deallocator
-	deallocated uint32
-	_refs       atomic.Int32
-	refs        *atomic.Int32
+	refs        atomic.Int32
 }
 
 func (b *Bytes) Size() int64 {
 	return int64(len(b.bytes))
 }
 
 func (b *Bytes) Bytes() []byte {
+	if b.refs.Load() <= 0 {
+		return nil
+	}
 	return b.bytes
 }
 
@@ -44,24 +45,20 @@ func (b *Bytes) Slice(length int) fscache.Data {
 }
 
 func (b *Bytes) Retain() {
-	if b.refs != nil {
-		b.refs.Add(1)
-	}
+	b.refs.Add(1)
 }
 
 func (b *Bytes) Release() {
-	if b.refs != nil {
-		if n := b.refs.Add(-1); n == 0 {
-			if b.deallocator != nil &&
-				atomic.CompareAndSwapUint32(&b.deallocated, 0, 1) {
-				b.deallocator.Deallocate(malloc.NoHints)
-			}
-		}
-	} else {
-		if b.deallocator != nil &&
-			atomic.CompareAndSwapUint32(&b.deallocated, 0, 1) {
+	n := b.refs.Add(-1)
+	if n == 0 {
+		// set bytes to nil
+		b.bytes = nil
+		if b.deallocator != nil {
 			b.deallocator.Deallocate(malloc.NoHints)
+			b.deallocator = nil
 		}
+	} else if n < 0 {
+		panic("Bytes.Release: double free")
 	}
 }
 
@@ -80,8 +77,7 @@ func (b *bytesAllocator) allocateCacheData(size int, hints malloc.Hints) fscache
 		bytes:       slice,
 		deallocator: dec,
 	}
-	bytes._refs.Store(1)
-	bytes.refs = &bytes._refs
+	bytes.refs.Store(1)
 	return bytes
 }
 

pkg/fileservice/bytes_test.go

@@ -15,7 +15,9 @@
 package fileservice
 
 import (
+	"sync"
 	"testing"
+	"time"
 
 	"github.com/matrixorigin/matrixone/pkg/common/malloc"
 	"github.com/stretchr/testify/assert"
@@ -29,6 +31,75 @@ func TestBytes(t *testing.T) {
 			bytes:       bytes,
 			deallocator: deallocator,
 		}
+		bs.refs.Store(1)
 		bs.Release()
 	})
 }
+
+func TestBytesError(t *testing.T) {
+	t.Run("Bytes double free", func(t *testing.T) {
+		bytes, deallocator, err := ioAllocator().Allocate(42, malloc.NoHints)
+		assert.Nil(t, err)
+		bs := &Bytes{
+			bytes:       bytes,
+			deallocator: deallocator,
+		}
+		bs.refs.Store(1)
+
+		// deallocate memory
+		bs.Release()
+
+		// nil pointer
+		ptr := bs.Bytes()
+		assert.Nil(t, ptr)
+
+		// double free
+		assert.Panics(t, func() { bs.Release() }, "double free")
+	})
+
+	t.Run("Bytes nil deallocator", func(t *testing.T) {
+		bs := &Bytes{
+			bytes:       []byte("123"),
+			deallocator: nil,
+		}
+		bs.refs.Store(1)
+
+		// deallocate memory
+		bs.Release()
+
+		// nil pointer
+		ptr := bs.Bytes()
+		assert.Nil(t, ptr)
+
+		assert.Panics(t, func() { bs.Release() }, "double free")
+	})
+}
+
+func TestBytesConcurrent(t *testing.T) {
+	bs := &Bytes{
+		bytes:       []byte("123"),
+		deallocator: nil,
+	}
+	bs.refs.Store(1)
+	nthread := 5
+	var wg sync.WaitGroup
+	for i := 0; i < nthread; i++ {
+		wg.Add(1)
+		go func(id int) {
+			defer wg.Done()
+
+			bs.Retain()
+
+			time.Sleep(1 * time.Millisecond)
+
+			bs.Release()
+		}(i)
+	}
+
+	wg.Wait()
+
+	bs.Release()
+
+	// double free
+	assert.Panics(t, func() { bs.Release() }, "double free")
+}