618 lines
15 KiB
Go
618 lines
15 KiB
Go
// Copyright 2025 PingCAP, Inc.
|
|
//
|
|
// 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 memory
|
|
|
|
import (
|
|
"fmt"
|
|
"math"
|
|
"math/rand"
|
|
"strings"
|
|
"sync"
|
|
"sync/atomic"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/stretchr/testify/require"
|
|
)
|
|
|
|
func (p *ResourcePool) setLimit(newLimit int64) {
|
|
atomic.StoreInt64(&p.limit, newLimit)
|
|
}
|
|
|
|
func TestPoolAllocations(t *testing.T) {
|
|
maxs := []int64{1, 9, 10, 11, 99, 100, 101, 0}
|
|
factors := []int{1, 2, 10, 10000}
|
|
poolAllocSizes := []int64{1, 2, 9, 10, 11, 100}
|
|
preBudgets := []int64{0, 1, 2, 9, 10, 11, 100}
|
|
|
|
rnd := rand.New(rand.NewSource(1))
|
|
|
|
var pool *ResourcePool
|
|
|
|
m := NewResourcePoolDefault("test", 0)
|
|
m.StartNoReserved(nil)
|
|
accs := make([]Budget, 4)
|
|
for i := range accs {
|
|
accs[i] = m.CreateBudget()
|
|
}
|
|
|
|
checkInvariants := func() {
|
|
var sum int64
|
|
fail := false
|
|
for accI := range accs {
|
|
if accs[accI].used < 0 {
|
|
t.Errorf("budget %d < 0: %d", accI, accs[accI].used)
|
|
fail = true
|
|
}
|
|
sum += accs[accI].Capacity()
|
|
}
|
|
if m.mu.allocated < 0 {
|
|
t.Errorf("pool size < 0: %d", m.mu.allocated)
|
|
fail = true
|
|
}
|
|
if sum != m.mu.allocated {
|
|
t.Errorf("total budget sum %d different from pool size %d", sum, m.mu.allocated)
|
|
fail = true
|
|
}
|
|
if m.mu.budget.used < 0 {
|
|
t.Errorf("pool budget < 0: %d", m.mu.budget.used)
|
|
fail = true
|
|
}
|
|
avail := m.mu.budget.Capacity() + m.reserved
|
|
if sum > avail {
|
|
t.Errorf("total budget sum %d greater than total pool budget %d", sum, avail)
|
|
fail = true
|
|
}
|
|
if pool.mu.allocated > pool.reserved {
|
|
t.Errorf("pool cur %d exceeds max %d", pool.mu.allocated, pool.reserved)
|
|
fail = true
|
|
}
|
|
if m.mu.budget.Capacity() != pool.mu.allocated {
|
|
t.Errorf("pool budget %d different from pool cur %d", m.mu.budget.used, pool.mu.allocated)
|
|
fail = true
|
|
}
|
|
|
|
if fail {
|
|
t.Fatal("invariants not preserved")
|
|
}
|
|
}
|
|
|
|
const numBudgetOps = 200
|
|
|
|
for _, max := range maxs {
|
|
pool = NewResourcePoolDefault("test", 1)
|
|
pool.Start(nil, (max))
|
|
|
|
for _, hf := range factors {
|
|
pool.maxUnusedBlocks = int64(hf)
|
|
|
|
for _, pb := range preBudgets {
|
|
mmax := pb + max
|
|
|
|
for _, pa := range poolAllocSizes {
|
|
m = NewResourcePoolDefault("test", pa)
|
|
m.Start(pool, (pb))
|
|
|
|
for range numBudgetOps {
|
|
accI := rnd.Intn(len(accs))
|
|
switch rnd.Intn(3) {
|
|
case 0:
|
|
sz := randomSize(rnd, mmax)
|
|
checkInvariants()
|
|
accs[accI].Grow(sz)
|
|
checkInvariants()
|
|
case 1:
|
|
checkInvariants()
|
|
accs[accI].Clear()
|
|
checkInvariants()
|
|
case 2:
|
|
osz := rnd.Int63n(accs[accI].used + 1)
|
|
nsz := randomSize(rnd, mmax)
|
|
checkInvariants()
|
|
accs[accI].resize(osz, nsz)
|
|
checkInvariants()
|
|
}
|
|
}
|
|
|
|
for accI := range accs {
|
|
checkInvariants()
|
|
accs[accI].Clear()
|
|
checkInvariants()
|
|
}
|
|
|
|
m.Stop()
|
|
if pool.mu.allocated != 0 {
|
|
t.Fatalf("pool not empty after pool close: %d", pool.mu.allocated)
|
|
}
|
|
}
|
|
}
|
|
}
|
|
pool.Stop()
|
|
}
|
|
}
|
|
|
|
func TestBudget(t *testing.T) {
|
|
p := NewResourcePoolDefault("test", 1)
|
|
p.Start(nil, (100))
|
|
p.allocAlignSize = 1
|
|
p.maxUnusedBlocks = 1
|
|
|
|
a1 := p.CreateBudget()
|
|
a2 := p.CreateBudget()
|
|
if err := a1.Grow(10); err != nil {
|
|
t.Fatalf("pool refused allocation: %v", err)
|
|
}
|
|
|
|
if err := a2.Grow(30); err != nil {
|
|
t.Fatalf("pool refused allocation: %v", err)
|
|
}
|
|
|
|
if err := a1.Grow(61); err == nil {
|
|
t.Fatalf("pool accepted excessive allocation")
|
|
}
|
|
|
|
if err := a2.Grow(61); err == nil {
|
|
t.Fatalf("pool accepted excessive allocation")
|
|
}
|
|
|
|
a1.Clear()
|
|
|
|
if err := a2.Grow(61); err != nil {
|
|
t.Fatalf("pool refused allocation: %v", err)
|
|
}
|
|
|
|
if err := a2.resize(50, 60); err == nil {
|
|
t.Fatalf("pool accepted excessive allocation")
|
|
}
|
|
|
|
if err := a1.resize(0, 5); err != nil {
|
|
t.Fatalf("pool refused allocation: %v", err)
|
|
}
|
|
|
|
if err := a2.resize(a2.used, 40); err != nil {
|
|
t.Fatalf("pool refused reset + allocation: %v", err)
|
|
}
|
|
|
|
a1.Clear()
|
|
a2.Clear()
|
|
|
|
if p.mu.allocated != 0 {
|
|
t.Fatal("closing spans leaves bytes in pool")
|
|
}
|
|
|
|
if m2 := a1.Pool(); m2 != p {
|
|
t.Fatalf("a1.Pool() returned %v, wanted %v", m2, &p)
|
|
}
|
|
|
|
p.Stop()
|
|
}
|
|
|
|
func TestNilBudget(t *testing.T) {
|
|
var ba *Budget
|
|
_ = ba.Used()
|
|
_ = ba.Pool()
|
|
_ = ba.Capacity()
|
|
ba.Empty()
|
|
ba.Clear()
|
|
ba.Clear()
|
|
require.Nil(t, ba.resize(0, 10))
|
|
require.Nil(t, ba.ResizeTo(10))
|
|
require.Nil(t, ba.Grow(10))
|
|
ba.Shrink(10)
|
|
}
|
|
|
|
func TestResourcePool(t *testing.T) {
|
|
p := NewResourcePoolDefault("test", 1)
|
|
p.Start(nil, (100))
|
|
p.maxUnusedBlocks = 1
|
|
|
|
if err := p.allocate(10); err != nil {
|
|
t.Fatalf("pool refused small allocation: %v", err)
|
|
}
|
|
if err := p.allocate(91); err == nil {
|
|
t.Fatalf("pool accepted excessive allocation: %v", err)
|
|
}
|
|
if err := p.allocate(90); err != nil {
|
|
t.Fatalf("pool refused top allocation: %v", err)
|
|
}
|
|
if p.mu.allocated != 100 {
|
|
t.Fatalf("incorrect allocation: got %d, expected %d", p.mu.allocated, 100)
|
|
}
|
|
|
|
p.release(90)
|
|
if p.mu.allocated != 10 {
|
|
t.Fatalf("incorrect allocation: got %d, expected %d", p.mu.allocated, 10)
|
|
}
|
|
if p.mu.maxAllocated != 100 {
|
|
t.Fatalf("incorrect max allocation: got %d, expected %d", p.mu.maxAllocated, 100)
|
|
}
|
|
if p.MaxAllocated() != 100 {
|
|
t.Fatalf("incorrect MaximumBytes(): got %d, expected %d", p.mu.maxAllocated, 100)
|
|
}
|
|
|
|
p.release(10)
|
|
if p.mu.allocated != 0 {
|
|
t.Fatalf("incorrect allocation: got %d, expected %d", p.mu.allocated, 0)
|
|
}
|
|
|
|
limitedPool := newResourcePoolWithLimit(
|
|
"testlimit", 10, 1)
|
|
limitedPool.StartNoReserved(p)
|
|
|
|
if err := limitedPool.allocate(10); err != nil {
|
|
t.Fatalf("limited pool refused small allocation: %v", err)
|
|
}
|
|
if err := limitedPool.allocate(1); err == nil {
|
|
t.Fatal("limited pool allowed allocation over limit")
|
|
}
|
|
|
|
limitedPool.release(10)
|
|
|
|
limitedPool.Stop()
|
|
p.Stop()
|
|
}
|
|
|
|
func newResourcePoolWithLimit(name string, limiit, allocAlignSize int64) *ResourcePool {
|
|
return NewResourcePool(
|
|
newPoolUID(), name, limiit, allocAlignSize, DefMaxUnusedBlocks, PoolActions{})
|
|
}
|
|
|
|
func TestMemoryAllocationEdgeCases(t *testing.T) {
|
|
m := NewResourcePoolDefault("test", 1e9)
|
|
m.Start(nil, (1e9))
|
|
|
|
a := m.CreateBudget()
|
|
if err := a.Grow(1); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if err := a.Grow(math.MaxInt64); err == nil {
|
|
t.Fatalf("expected error, but found success")
|
|
}
|
|
|
|
a.Clear()
|
|
m.Stop()
|
|
}
|
|
|
|
func TestMultiSharedGauge(t *testing.T) {
|
|
minAllocation := int64(1000)
|
|
|
|
parent := NewResourcePoolDefault("root", minAllocation)
|
|
parent.Start(nil, (100000))
|
|
|
|
child := parent.NewResourcePoolInheritWithLimit("child", 20000)
|
|
child.StartNoReserved(parent)
|
|
|
|
acc := child.CreateBudget()
|
|
require.NoError(t, acc.Grow(100))
|
|
|
|
require.Equal(t, minAllocation, parent.Allocated())
|
|
}
|
|
|
|
func TestActions(t *testing.T) {
|
|
{
|
|
root := NewResourcePoolDefault("root", 1000)
|
|
root.Start(nil, math.MaxInt64)
|
|
p1 := NewResourcePoolDefault("p1", 666)
|
|
p1.StartNoReserved(root)
|
|
require.NoError(t, p1.ExplicitReserve(1002))
|
|
require.Equal(t, p1.Capacity(), root.allocAlignSize*2)
|
|
require.True(t, p1.mu.budget.explicitReserved == 1002)
|
|
pb1 := p1.CreateBudget()
|
|
require.NoError(t, pb1.Reserve(123))
|
|
require.Equal(t, p1.Capacity(), root.allocAlignSize*2)
|
|
require.Equal(t, pb1.Capacity(), p1.allocAlignSize)
|
|
require.True(t, pb1.explicitReserved == 123)
|
|
require.NoError(t, pb1.Grow(123))
|
|
pb1.Shrink(123)
|
|
require.Equal(t, p1.Capacity(), root.allocAlignSize*2)
|
|
pb1.Clear()
|
|
}
|
|
{
|
|
name := "root"
|
|
noteActionTriggerCnt := 0
|
|
root := NewResourcePool(
|
|
newPoolUID(),
|
|
name, math.MaxInt64,
|
|
1,
|
|
DefMaxUnusedBlocks,
|
|
PoolActions{
|
|
NoteAction: NoteAction{
|
|
Threshold: 101,
|
|
CB: func(s NoteActionState) {
|
|
require.Equal(t, s.Pool.name, name)
|
|
noteActionTriggerCnt += 1
|
|
},
|
|
},
|
|
},
|
|
)
|
|
root.Start(nil /* pool */, (math.MaxInt64))
|
|
root.maxUnusedBlocks = 0
|
|
|
|
b := root.CreateBudget()
|
|
require.NoError(t, b.Grow(100))
|
|
|
|
require.Equal(t, b.Used(), int64(100))
|
|
require.Equal(t, noteActionTriggerCnt, 0)
|
|
|
|
require.NoError(t, b.Grow(2))
|
|
require.Equal(t, b.Used(), int64(102))
|
|
require.Equal(t, noteActionTriggerCnt, 1)
|
|
}
|
|
|
|
{
|
|
name := "root"
|
|
outOfCapacityActionNum := 0
|
|
root := NewResourcePoolDefault(
|
|
name,
|
|
1,
|
|
)
|
|
root.StartNoReserved(nil)
|
|
|
|
b := root.CreateBudget()
|
|
|
|
{
|
|
err := b.Grow(1)
|
|
require.Error(t, err)
|
|
}
|
|
|
|
root.SetOutOfCapacityAction(func(s OutOfCapacityActionArgs) error {
|
|
require.Equal(t, s.Pool.name, name)
|
|
s.Pool.forceAddCap(s.Request)
|
|
outOfCapacityActionNum++
|
|
return nil
|
|
})
|
|
|
|
require.NoError(t, b.Grow(1))
|
|
|
|
require.Equal(t, outOfCapacityActionNum, 1)
|
|
|
|
require.NoError(t, b.Grow(10))
|
|
|
|
require.Equal(t, outOfCapacityActionNum, 2)
|
|
|
|
root.ApproxAvailable()
|
|
|
|
require.Equal(t, root.mu.budget.used, b.Used())
|
|
|
|
b.Clear()
|
|
|
|
require.NoError(t, b.Grow(5))
|
|
|
|
require.Equal(t, outOfCapacityActionNum, 2)
|
|
|
|
b.Clear()
|
|
|
|
outOfLimitActionCnt := 0
|
|
root.SetOutOfLimitAction(func(r *ResourcePool) error {
|
|
outOfLimitActionCnt++
|
|
root.setLimit(root.capacity())
|
|
return fmt.Errorf("")
|
|
})
|
|
root.SetLimit(1)
|
|
require.True(t, root.Limit() == 1)
|
|
require.Error(t, b.Grow(2))
|
|
require.Equal(t, outOfLimitActionCnt, 1)
|
|
err := b.Grow(root.capacity() - root.allocated())
|
|
require.NoError(t, err)
|
|
require.Equal(t, outOfLimitActionCnt, 1)
|
|
}
|
|
}
|
|
|
|
func genPool(
|
|
name string, parent *ResourcePool,
|
|
) *ResourcePool {
|
|
var reservedBytes int64
|
|
if parent == nil {
|
|
reservedBytes = math.MaxInt64
|
|
}
|
|
return getPoolImpl(name, parent, reservedBytes)
|
|
}
|
|
|
|
func getPoolImpl(
|
|
name string, parent *ResourcePool, reservedBytes int64,
|
|
) *ResourcePool {
|
|
m := NewResourcePoolDefault(name, 1)
|
|
m.Start(parent, (reservedBytes))
|
|
return m
|
|
}
|
|
|
|
func getPoolUsed(
|
|
t *testing.T,
|
|
name string,
|
|
parent *ResourcePool,
|
|
usedBytes, reservedBytes int64,
|
|
) *ResourcePool {
|
|
m := getPoolImpl(name, parent, reservedBytes)
|
|
if usedBytes != 0 {
|
|
acc := m.CreateBudget()
|
|
if err := acc.Grow(usedBytes); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
}
|
|
return m
|
|
}
|
|
|
|
func TestResourcePoolTree(t *testing.T) {
|
|
export := func(p *ResourcePool) string {
|
|
var pools []ResourcePoolState
|
|
_ = p.Traverse(func(pool ResourcePoolState) error {
|
|
pools = append(pools, pool)
|
|
return nil
|
|
})
|
|
var sb strings.Builder
|
|
for _, e := range pools {
|
|
for range e.Level {
|
|
sb.WriteString("-")
|
|
}
|
|
sb.WriteString(e.Name + "\n")
|
|
}
|
|
return sb.String()
|
|
}
|
|
|
|
parent := genPool("parent", nil)
|
|
child1 := genPool("child1", parent)
|
|
child2 := genPool("child2", parent)
|
|
require.Equal(t, "parent\n-child2\n-child1\n", export(parent))
|
|
require.Equal(t, "child1\n", export(child1))
|
|
require.Equal(t, "child2\n", export(child2))
|
|
|
|
grandchild1 := genPool("grandchild1", child1)
|
|
grandchild2 := genPool("grandchild2", child2)
|
|
require.Equal(t, "parent\n-child2\n--grandchild2\n-child1\n--grandchild1\n", export(parent))
|
|
require.Equal(t, "child1\n-grandchild1\n", export(child1))
|
|
require.Equal(t, "child2\n-grandchild2\n", export(child2))
|
|
|
|
grandchild2.Stop()
|
|
child2.Stop()
|
|
|
|
require.Equal(t, "parent\n-child1\n--grandchild1\n", export(parent))
|
|
require.Equal(t, "child1\n-grandchild1\n", export(child1))
|
|
|
|
grandchild1.Stop()
|
|
child1.Stop()
|
|
|
|
require.Equal(t, "parent\n", export(parent))
|
|
parent.Stop()
|
|
}
|
|
|
|
func TestResourcePoolUsedFromReserved(t *testing.T) {
|
|
root := genPool("root", nil)
|
|
const usedBytes = int64(1 << 10)
|
|
child := getPoolUsed(t, "child", root, usedBytes, 2*usedBytes)
|
|
|
|
require.NoError(t, root.Traverse(func(s ResourcePoolState) error {
|
|
if s.Name == child.Name() {
|
|
require.Equal(t, usedBytes, s.Used)
|
|
}
|
|
return nil
|
|
}))
|
|
}
|
|
|
|
func TestResourcePoolNoDeadlocks(t *testing.T) {
|
|
root := genPool("root", nil)
|
|
defer root.Stop()
|
|
|
|
var wg sync.WaitGroup
|
|
const numGoroutines = 10
|
|
done := make(chan struct{})
|
|
for i := range numGoroutines {
|
|
wg.Add(1)
|
|
go func(i int) {
|
|
defer wg.Done()
|
|
rng := rand.New(rand.NewSource(1))
|
|
for {
|
|
select {
|
|
case <-done:
|
|
return
|
|
default:
|
|
func() {
|
|
m := genPool(fmt.Sprintf("m%d", i), root)
|
|
defer m.Stop()
|
|
numOps := rng.Intn(10 + 1)
|
|
var reserved int64
|
|
defer func() {
|
|
m.release(reserved)
|
|
}()
|
|
for op := 0; op < numOps; op++ {
|
|
if reserved > 0 && rng.Float64() < 0.5 {
|
|
toRelease := int64(rng.Intn(int(reserved))) + 1
|
|
m.release(toRelease)
|
|
reserved -= toRelease
|
|
} else {
|
|
toReserve := int64(rng.Intn(1000) + 1)
|
|
_ = m.allocate(toReserve)
|
|
reserved += toReserve
|
|
}
|
|
time.Sleep(time.Duration(rng.Intn(1000)) * time.Microsecond)
|
|
}
|
|
}()
|
|
time.Sleep(time.Duration(rng.Intn(2000)) * time.Microsecond)
|
|
}
|
|
}
|
|
}(i)
|
|
}
|
|
|
|
rng := rand.New(rand.NewSource(1))
|
|
for range 1000 {
|
|
var pools []ResourcePoolState
|
|
_ = root.Traverse(func(pool ResourcePoolState) error {
|
|
pools = append(pools, pool)
|
|
return nil
|
|
})
|
|
for _, m := range pools {
|
|
require.NotEqual(t, ResourcePoolState{}, m)
|
|
}
|
|
time.Sleep(time.Duration(rng.Intn(3000)) * time.Microsecond)
|
|
}
|
|
close(done)
|
|
wg.Wait()
|
|
}
|
|
|
|
func BenchmarkBudgetGrow(b *testing.B) {
|
|
m := NewResourcePoolDefault("test",
|
|
1e9,
|
|
)
|
|
m.Start(nil, (1e9))
|
|
|
|
a := m.CreateBudget()
|
|
for range b.N {
|
|
_ = a.Grow(1)
|
|
}
|
|
}
|
|
|
|
func BenchmarkTraverseTree(b *testing.B) {
|
|
makePoolTree := func(numLevels int, numChildrenPerPool int) (root *ResourcePool, cleanup func()) {
|
|
allPools := make([][]*ResourcePool, numLevels)
|
|
allPools[0] = []*ResourcePool{genPool("root", nil)}
|
|
for level := 1; level < numLevels; level++ {
|
|
allPools[level] = make([]*ResourcePool, 0, len(allPools[level-1])*numChildrenPerPool)
|
|
for parent, parentMon := range allPools[level-1] {
|
|
for child := range numChildrenPerPool {
|
|
name := fmt.Sprintf("child%d_parent%d", child, parent)
|
|
allPools[level] = append(allPools[level], genPool(name, parentMon))
|
|
}
|
|
}
|
|
}
|
|
cleanup = func() {
|
|
for i := len(allPools) - 1; i >= 0; i-- {
|
|
for _, m := range allPools[i] {
|
|
m.Stop()
|
|
}
|
|
}
|
|
}
|
|
return allPools[0][0], cleanup
|
|
}
|
|
for _, numLevels := range []int{2, 4, 8} {
|
|
for _, numChildrenPerPool := range []int{2, 4, 8} {
|
|
b.Run(fmt.Sprintf("levels=%d/children=%d", numLevels, numChildrenPerPool), func(b *testing.B) {
|
|
root, cleanup := makePoolTree(numLevels, numChildrenPerPool)
|
|
defer cleanup()
|
|
b.ResetTimer()
|
|
for range b.N {
|
|
var numPools int
|
|
_ = root.Traverse(func(ResourcePoolState) error {
|
|
numPools++
|
|
return nil
|
|
})
|
|
}
|
|
})
|
|
}
|
|
}
|
|
}
|
|
|
|
func randomSize(rnd *rand.Rand, mag int64) int64 {
|
|
return int64(rnd.ExpFloat64() * float64(mag) * 0.3679)
|
|
}
|