malloc: add profiler

pkg/common/malloc/chain_deallocator.go

@@ -0,0 +1,63 @@
+// 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 (
+	"sync"
+	"unsafe"
+)
+
+type chainDeallocator []Deallocator
+
+var _ Deallocator = &chainDeallocator{}
+
+func (c *chainDeallocator) Deallocate(ptr unsafe.Pointer) {
+	for i := len(*c) - 1; i >= 0; i-- {
+		(*c)[i].Deallocate(ptr)
+	}
+	*c = (*c)[:0]
+	chainDeallocatorPool.Put(c)
+}
+
+func ChainDeallocator(dec1 Deallocator, dec2 Deallocator) Deallocator {
+	if chain1, ok := dec1.(*chainDeallocator); ok {
+		if chain2, ok := dec2.(*chainDeallocator); ok {
+			*chain1 = append(*chain1, *chain2...)
+			*chain2 = (*chain2)[:0]
+			chainDeallocatorPool.Put(chain2)
+			return chain1
+		} else {
+			*chain1 = append(*chain1, dec2)
+			return chain1
+		}
+	}
+	if chain2, ok := dec2.(*chainDeallocator); ok {
+		chain := chainDeallocatorPool.Get().(*chainDeallocator)
+		*chain = append(*chain, dec1)
+		*chain = append(*chain, *chain2...)
+		*chain2 = (*chain2)[:0]
+		chainDeallocatorPool.Put(chain2)
+		return chain
+	}
+	chain := chainDeallocatorPool.Get().(*chainDeallocator)
+	*chain = append(*chain, dec1, dec2)
+	return chain
+}
+
+var chainDeallocatorPool = sync.Pool{
+	New: func() any {
+		return &chainDeallocator{}
+	},
+}

pkg/common/malloc/checked_deallocator_test.go

@@ -22,7 +22,7 @@ import (
 )
 
 func TestCheckedDeallocator(t *testing.T) {
-	allocator := NewClassAllocator(1)
+	allocator := NewClassAllocator(ptrTo(uint32(1)))
 
 	func() {
 		defer func() {

pkg/common/malloc/class_allocator.go

@@ -52,7 +52,7 @@ type Class struct {
 }
 
 func NewClassAllocator(
-	checkFraction uint32,
+	checkFraction *uint32,
 ) *ClassAllocator {
 	ret := &ClassAllocator{}
 
@@ -90,7 +90,7 @@ func (c *ClassAllocator) Allocate(size uint64) (unsafe.Pointer, Deallocator) {
 
 type fixedSizeMmapAllocator struct {
 	size          uint64
-	checkFraction uint32
+	checkFraction *uint32
 	// buffer1 buffers objects
 	buffer1 chan unsafe.Pointer
 	// buffer2 buffers MADV_DONTNEED objects
@@ -99,7 +99,7 @@ type fixedSizeMmapAllocator struct {
 
 func newFixedSizedAllocator(
 	size uint64,
-	checkFraction uint32,
+	checkFraction *uint32,
 ) *fixedSizeMmapAllocator {
 	// if size is larger than smallClassCap, num1 will be zero, buffer1 will be empty
 	num1 := smallClassCap / size
@@ -119,9 +119,9 @@ func newFixedSizedAllocator(
 var _ Allocator = new(fixedSizeMmapAllocator)
 
 func (f *fixedSizeMmapAllocator) Allocate(_ uint64) (ptr unsafe.Pointer, dec Deallocator) {
-	if f.checkFraction > 0 {
+	if f.checkFraction != nil {
 		defer func() {
-			if fastrand()%f.checkFraction == 0 {
+			if fastrand()%*f.checkFraction == 0 {
 				dec = newCheckedDeallocator(dec)
 			}
 		}()
@@ -167,8 +167,8 @@ var _ Deallocator = new(fixedSizeMmapAllocator)
 
 func (f *fixedSizeMmapAllocator) Deallocate(ptr unsafe.Pointer) {
 
-	if f.checkFraction > 0 &&
-		fastrand()%f.checkFraction == 0 {
+	if f.checkFraction != nil &&
+		fastrand()%*f.checkFraction == 0 {
 		// do not reuse to detect use-after-free
 		if err := unix.Munmap(
 			unsafe.Slice((*byte)(ptr), f.size),

pkg/common/malloc/class_allocator_test.go

@@ -18,6 +18,6 @@ import "testing"
 
 func TestClassAllocator(t *testing.T) {
 	testAllocator(t, func() Allocator {
-		return NewClassAllocator(1)
+		return NewClassAllocator(ptrTo(uint32(1)))
 	})
 }

pkg/common/malloc/config.go

@@ -24,19 +24,46 @@ import (
 type Config struct {
 	// CheckFraction controls the fraction of checked deallocations
 	// On average, 1 / fraction of deallocations will be checked for double free or missing free
-	CheckFraction uint32 `toml:"check-fraction"`
+	CheckFraction *uint32 `toml:"check-fraction"`
+
+	// FullStackFraction controls the fraction of full stack collecting in profile sampling
+	// On average, 1 / fraction of allocations will be sampled with full stack information
+	FullStackFraction *uint32 `toml:"full-stack-fraction"`
 
 	// EnableMetrics indicates whether to expose metrics to prometheus
-	EnableMetrics bool `toml:"enable-metrics"`
+	EnableMetrics *bool `toml:"enable-metrics"`
 }
 
 var defaultConfig = func() *atomic.Pointer[Config] {
 	ret := new(atomic.Pointer[Config])
-	ret.Store(new(Config))
+
+	// default config
+	ret.Store(&Config{
+		CheckFraction:     ptrTo(uint32(65536)),
+		EnableMetrics:     ptrTo(true),
+		FullStackFraction: ptrTo(uint32(16)),
+	})
+
 	return ret
 }()
 
-func SetDefaultConfig(c Config) {
-	defaultConfig.Store(&c)
-	logutil.Info("malloc: set default config", zap.Any("config", c))
+func SetDefaultConfig(delta Config) {
+
+	// read
+	config := *defaultConfig.Load()
+
+	// set
+	if delta.CheckFraction != nil {
+		config.CheckFraction = delta.CheckFraction
+	}
+	if delta.EnableMetrics != nil {
+		config.EnableMetrics = delta.EnableMetrics
+	}
+	if delta.FullStackFraction != nil {
+		config.FullStackFraction = delta.FullStackFraction
+	}
+
+	// update
+	defaultConfig.Store(&config)
+	logutil.Info("malloc: set default config", zap.Any("config", delta))
 }

pkg/common/malloc/default.go

@@ -15,6 +15,7 @@
 package malloc
 
 import (
+	"net/http"
 	"os"
 	"runtime"
 	"strings"
@@ -26,17 +27,22 @@ func NewDefault(config *Config) (allocator Allocator) {
 		config = &c
 	}
 
-	var metrics Metrics
-	if config.EnableMetrics {
-		go metrics.startExport()
-	}
+	defer func() {
+		if config.FullStackFraction != nil && *config.FullStackFraction > 0 {
+			allocator = NewProfileAllocator(
+				allocator,
+				globalProfiler,
+				*config.FullStackFraction,
+			)
+		}
+	}()
 
 	switch strings.TrimSpace(strings.ToLower(os.Getenv("MO_MALLOC"))) {
 
 	case "c":
 		allocator = NewCAllocator()
-		if config.EnableMetrics {
-			allocator = NewMetricsAllocator(allocator, &metrics)
+		if config.EnableMetrics != nil && *config.EnableMetrics {
+			allocator = NewMetricsAllocator(allocator)
 		}
 		return allocator
 
@@ -46,8 +52,8 @@ func NewDefault(config *Config) (allocator Allocator) {
 			func() Allocator {
 				var ret Allocator
 				ret = NewPureGoClassAllocator(256 * MB)
-				if config.EnableMetrics {
-					ret = NewMetricsAllocator(ret, &metrics)
+				if config.EnableMetrics != nil && *config.EnableMetrics {
+					ret = NewMetricsAllocator(ret)
 				}
 				return ret
 			},
@@ -59,12 +65,20 @@ func NewDefault(config *Config) (allocator Allocator) {
 			func() Allocator {
 				var ret Allocator
 				ret = NewClassAllocator(config.CheckFraction)
-				if config.EnableMetrics {
-					ret = NewMetricsAllocator(ret, &metrics)
+				if config.EnableMetrics != nil && *config.EnableMetrics {
+					ret = NewMetricsAllocator(ret)
 				}
 				return ret
 			},
 		)
 
 	}
 }
+
+var globalProfiler = NewProfiler[HeapSampleValues]()
+
+func init() {
+	http.HandleFunc("/debug/pprof/malloc", func(w http.ResponseWriter, req *http.Request) {
+		globalProfiler.Write(w)
+	})
+}

pkg/common/malloc/func_deallocator.go

@@ -0,0 +1,40 @@
+// 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 "unsafe"
+
+type FuncDeallocator func(ptr unsafe.Pointer)
+
+var _ Deallocator = FuncDeallocator(nil)
+
+func (f FuncDeallocator) Deallocate(ptr unsafe.Pointer) {
+	f(ptr)
+}
+
+type argumentedFuncDeallocator[T any] struct {
+	argument T
+	fn       func(unsafe.Pointer, T)
+}
+
+func (a *argumentedFuncDeallocator[T]) SetArgument(arg T) {
+	a.argument = arg
+}
+
+var _ Deallocator = &argumentedFuncDeallocator[int]{}
+
+func (a *argumentedFuncDeallocator[T]) Deallocate(ptr unsafe.Pointer) {
+	a.fn(ptr, a.argument)
+}

pkg/common/malloc/metrics_allocator.go

@@ -16,31 +16,31 @@ package malloc
 
 import (
 	"sync"
-	"sync/atomic"
-	"time"
 	"unsafe"
 
-	"github.com/matrixorigin/matrixone/pkg/logutil"
 	metric "github.com/matrixorigin/matrixone/pkg/util/metric/v2"
-	"go.uber.org/zap"
 )
 
 type MetricsAllocator struct {
 	upstream Allocator
-	metrics  *Metrics
-	inUse    sync.Map // unsafe.Pointer -> size
+	funcPool sync.Pool
 }
 
-type Metrics struct {
-	AllocateBytesDelta atomic.Uint64
-	FreeBytesDelta     atomic.Uint64
-}
-
-func NewMetricsAllocator(upstream Allocator, metrics *Metrics) *MetricsAllocator {
-	return &MetricsAllocator{
+func NewMetricsAllocator(upstream Allocator) *MetricsAllocator {
+	ret := &MetricsAllocator{
 		upstream: upstream,
-		metrics:  metrics,
 	}
+	ret.funcPool = sync.Pool{
+		New: func() any {
+			argumented := new(argumentedFuncDeallocator[uint64])
+			argumented.fn = func(_ unsafe.Pointer, size uint64) {
+				metric.MallocCounterFreeBytes.Add(float64(size))
+				ret.funcPool.Put(argumented)
+			}
+			return argumented
+		},
+	}
+	return ret
 }
 
 type AllocateInfo struct {
@@ -51,44 +51,9 @@ type AllocateInfo struct {
 var _ Allocator = new(MetricsAllocator)
 
 func (m *MetricsAllocator) Allocate(size uint64) (unsafe.Pointer, Deallocator) {
-	m.metrics.AllocateBytesDelta.Add(size)
+	metric.MallocCounterAllocateBytes.Add(float64(size))
 	ptr, dec := m.upstream.Allocate(size)
-	m.inUse.Store(ptr, AllocateInfo{
-		Deallocator: dec,
-		Size:        size,
-	})
-	return ptr, m
-}
-
-var _ Deallocator = new(MetricsAllocator)
-
-func (m *MetricsAllocator) Deallocate(ptr unsafe.Pointer) {
-	v, ok := m.inUse.LoadAndDelete(ptr)
-	if !ok {
-		panic("double free")
-	}
-	info := v.(AllocateInfo)
-	m.metrics.FreeBytesDelta.Add(info.Size)
-	info.Deallocator.Deallocate(ptr)
-}
-
-func (m *Metrics) startExport() {
-	var sumAllocateBytes, sumFreeBytes, lastSumAllocateBytes uint64
-	for range time.NewTicker(time.Second).C {
-		allocateBytes := m.AllocateBytesDelta.Swap(0)
-		freeBytes := m.FreeBytesDelta.Swap(0)
-
-		sumAllocateBytes += allocateBytes
-		sumFreeBytes += freeBytes
-		if sumAllocateBytes-lastSumAllocateBytes > (1 << 30) {
-			logutil.Info("malloc stats",
-				zap.Any("allocate", sumAllocateBytes),
-				zap.Any("free", sumFreeBytes),
-			)
-			lastSumAllocateBytes = sumAllocateBytes
-		}
-
-		metric.MallocCounterAllocateBytes.Add(float64(allocateBytes))
-		metric.MallocCounterFreeBytes.Add(float64(freeBytes))
-	}
+	fn := m.funcPool.Get().(*argumentedFuncDeallocator[uint64])
+	fn.SetArgument(size)
+	return ptr, ChainDeallocator(dec, fn)
 }