malloc: various changes and optimizations (matrixorigin#16990)

malloc: add checked allocator

malloc: wrap checked allocator in default allocator

malloc: add hints

malloc: remove checked deallocator

malloc: add DoNotReuse hint to checked deallocator

malloc: add stack info to checked allocator

malloc: more tests for check allocator

malloc: add allocator config

malloc: check double free before calling upstream in CheckedAllocator deallocation

malloc: do not use default allocator in tests to avoid dependency loop

malloc: fix stacktrace

malloc: more detail infos for checked allocator panic messages

malloc: show pointer address in checked allocator panic messages

malloc: show allocated size in checked allocator panic messages

malloc: fix checked allocator finalizer

fileservice: add fuzzFS

fileservice: refine TestFuzzingDiskS3

malloc: add NoClear hint

malloc: refine ClassAllocator tests

malloc: add FixedSizeAllocator

malloc: add FixedSizeSyncPoolAllocator

malloc: use fixedSizeSyncPoolAllocator as go allocator by default

fileservice: fix missing release in tests

pkg/common/malloc/allocator.go

@@ -17,9 +17,9 @@ package malloc
 import "unsafe"
 
 type Allocator interface {
-	Allocate(size uint64) (unsafe.Pointer, Deallocator, error)
+	Allocate(size uint64, hint Hints) (unsafe.Pointer, Deallocator, error)
 }
 
 type Deallocator interface {
-	Deallocate(unsafe.Pointer)
+	Deallocate(ptr unsafe.Pointer, hint Hints)
 }

pkg/common/malloc/allocator_bench_test.go

@@ -33,11 +33,11 @@ func benchmarkAllocator(
 		allcator := newAllocator()
 		b.ResetTimer()
 		for i := 0; i < b.N; i++ {
-			ptr, dec, err := allcator.Allocate(n)
+			ptr, dec, err := allcator.Allocate(n, NoHints)
 			if err != nil {
 				b.Fatal(err)
 			}
-			dec.Deallocate(ptr)
+			dec.Deallocate(ptr, NoHints)
 		}
 	})
 
@@ -46,11 +46,11 @@ func benchmarkAllocator(
 		b.ResetTimer()
 		b.RunParallel(func(pb *testing.PB) {
 			for pb.Next() {
-				ptr, dec, err := allcator.Allocate(n)
+				ptr, dec, err := allcator.Allocate(n, NoHints)
 				if err != nil {
 					b.Fatal(err)
 				}
-				dec.Deallocate(ptr)
+				dec.Deallocate(ptr, NoHints)
 			}
 		})
 	})

pkg/common/malloc/allocator_test.go

@@ -29,7 +29,7 @@ func testAllocator(
 	t.Run("allocate", func(t *testing.T) {
 		allocator := newAllocator()
 		for i := uint64(1); i < 128*MB; i = uint64(math.Ceil(float64(i) * 1.1)) {
-			ptr, dec, err := allocator.Allocate(i)
+			ptr, dec, err := allocator.Allocate(i, NoHints)
 			if err != nil {
 				t.Fatal(err)
 			}
@@ -42,7 +42,7 @@ func testAllocator(
 			for i := range slice {
 				slice[i] = byte(i)
 			}
-			dec.Deallocate(ptr)
+			dec.Deallocate(ptr, NoHints)
 		}
 	})
 

pkg/common/malloc/c_allocator.go

@@ -30,14 +30,16 @@ func NewCAllocator() *CAllocator {
 
 var _ Allocator = new(CAllocator)
 
-func (c *CAllocator) Allocate(size uint64) (unsafe.Pointer, Deallocator, error) {
+func (c *CAllocator) Allocate(size uint64, hints Hints) (unsafe.Pointer, Deallocator, error) {
 	ptr := C.malloc(C.ulong(size))
-	clear(unsafe.Slice((*byte)(ptr), size))
+	if hints&NoClear == 0 {
+		clear(unsafe.Slice((*byte)(ptr), size))
+	}
 	return ptr, c, nil
 }
 
 var _ Deallocator = new(CAllocator)
 
-func (c *CAllocator) Deallocate(ptr unsafe.Pointer) {
+func (c *CAllocator) Deallocate(ptr unsafe.Pointer, hints Hints) {
 	C.free(ptr)
 }

pkg/common/malloc/chain_deallocator.go

@@ -23,9 +23,9 @@ type chainDeallocator []Deallocator
 
 var _ Deallocator = &chainDeallocator{}
 
-func (c *chainDeallocator) Deallocate(ptr unsafe.Pointer) {
+func (c *chainDeallocator) Deallocate(ptr unsafe.Pointer, hints Hints) {
 	for i := len(*c) - 1; i >= 0; i-- {
-		(*c)[i].Deallocate(ptr)
+		(*c)[i].Deallocate(ptr, hints)
 	}
 	*c = (*c)[:0]
 	chainDeallocatorPool.Put(c)

pkg/common/malloc/checked_allocator.go

@@ -0,0 +1,103 @@
+// Copyright 2024 Matrix Origin
+//
+// 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 malloc
+
+import (
+	"fmt"
+	"runtime"
+	"sync"
+	"sync/atomic"
+	"unsafe"
+)
+
+type CheckedAllocator struct {
+	upstream Allocator
+	fraction uint32
+	funcPool sync.Pool
+}
+
+type checkedAllocatorArgs struct {
+	deallocated *atomic.Bool
+	deallocator Deallocator
+	stackID     uint64
+	size        uint64
+}
+
+func NewCheckedAllocator(upstream Allocator, fraction uint32) *CheckedAllocator {
+	ret := &CheckedAllocator{
+		upstream: upstream,
+		fraction: fraction,
+	}
+
+	ret.funcPool = sync.Pool{
+		New: func() any {
+			argumented := new(argumentedFuncDeallocator[checkedAllocatorArgs])
+			argumented.fn = func(ptr unsafe.Pointer, hints Hints, args checkedAllocatorArgs) {
+
+				if !args.deallocated.CompareAndSwap(false, true) {
+					panic(fmt.Sprintf(
+						"double free: address %p, size %v, allocated at %s",
+						ptr,
+						args.size,
+						stackInfo(args.stackID),
+					))
+				}
+
+				hints |= DoNotReuse
+				args.deallocator.Deallocate(ptr, hints)
+
+				ret.funcPool.Put(argumented)
+			}
+			return argumented
+		},
+	}
+
+	return ret
+}
+
+var _ Allocator = new(CheckedAllocator)
+
+func (c *CheckedAllocator) Allocate(size uint64, hints Hints) (unsafe.Pointer, Deallocator, error) {
+	ptr, dec, err := c.upstream.Allocate(size, hints)
+	if err != nil {
+		return nil, nil, err
+	}
+
+	if fastrand()%c.fraction > 0 {
+		return ptr, dec, nil
+	}
+
+	stackID := getStacktraceID(0)
+	deallocated := new(atomic.Bool) // this will not be GC until deallocator is called
+	runtime.SetFinalizer(deallocated, func(deallocated *atomic.Bool) {
+		if !deallocated.Load() {
+			panic(fmt.Sprintf(
+				"missing free: address %p, size %v, allocated at %s",
+				ptr,
+				size,
+				stackInfo(stackID),
+			))
+		}
+	})
+
+	fn := c.funcPool.Get().(*argumentedFuncDeallocator[checkedAllocatorArgs])
+	fn.SetArgument(checkedAllocatorArgs{
+		deallocated: deallocated,
+		deallocator: dec,
+		stackID:     stackID,
+		size:        size,
+	})
+	return ptr, fn, nil
+}

pkg/common/malloc/checked_allocator_test.go

@@ -0,0 +1,118 @@
+// Copyright 2024 Matrix Origin
+//
+// 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 malloc
+
+import (
+	"fmt"
+	"runtime"
+	"strings"
+	"testing"
+	"unsafe"
+)
+
+func TestCheckedAllocator(t *testing.T) {
+	testAllocator(t, func() Allocator {
+		return NewCheckedAllocator(
+			NewClassAllocator(NewFixedSizeMmapAllocator),
+			1,
+		)
+	})
+
+	// missing free
+	t.Run("missing free", func(t *testing.T) {
+		allocator := NewCheckedAllocator(
+			NewClassAllocator(NewFixedSizeMmapAllocator),
+			1,
+		)
+		ptr, dec, err := allocator.Allocate(42, NoHints)
+		if err != nil {
+			t.Fatal(err)
+		}
+		// comment the following line to trigger a missing-free panic
+		// this panic will be raised in SetFinalizer func so it's not recoverable and not testable
+		dec.Deallocate(ptr, NoHints)
+		_ = ptr
+		_ = dec
+		runtime.GC()
+	})
+
+	// double free
+	t.Run("double free", func(t *testing.T) {
+		defer func() {
+			p := recover()
+			if p == nil {
+				t.Fatal("should panic")
+			}
+			msg := fmt.Sprintf("%v", p)
+			if !strings.Contains(msg, "double free") {
+				t.Fatalf("got %v", msg)
+			}
+		}()
+		allocator := NewCheckedAllocator(
+			NewClassAllocator(NewFixedSizeMmapAllocator),
+			1,
+		)
+		ptr, dec, err := allocator.Allocate(42, NoHints)
+		if err != nil {
+			t.Fatal(err)
+		}
+		dec.Deallocate(ptr, NoHints)
+		dec.Deallocate(ptr, NoHints)
+	})
+
+	// use after free
+	t.Run("use after free", func(t *testing.T) {
+		allocator := NewCheckedAllocator(
+			NewClassAllocator(NewFixedSizeMmapAllocator),
+			1,
+		)
+		ptr, dec, err := allocator.Allocate(42, NoHints)
+		if err != nil {
+			t.Fatal(err)
+		}
+		slice := unsafe.Slice((*byte)(ptr), 42)
+		for i := range slice {
+			slice[i] = byte(i)
+		}
+		dec.Deallocate(ptr, NoHints)
+		// zero or segfault
+		//for i := range slice {
+		//	if slice[i] != 0 {
+		//		t.Fatal("should zero")
+		//	}
+		//}
+	})
+
+}
+
+func BenchmarkCheckedAllocator(b *testing.B) {
+	for _, n := range benchNs {
+		benchmarkAllocator(b, func() Allocator {
+			return NewCheckedAllocator(
+				NewClassAllocator(NewFixedSizeMmapAllocator),
+				1,
+			)
+		}, n)
+	}
+}
+
+func FuzzCheckedAllocator(f *testing.F) {
+	fuzzAllocator(f, func() Allocator {
+		return NewCheckedAllocator(
+			NewClassAllocator(NewFixedSizeMmapAllocator),
+			1,
+		)
+	})
+}