planner: add delete-by-key support for instance plan cache

Author: winoros

pkg/planner/core/plan_cache_instance.go

@@ -43,8 +43,11 @@ type instancePCNode struct {
 	value    *PlanCacheValue
 	lastUsed atomic.Time
 	next     atomic.Pointer[instancePCNode]
+	deleted  atomic.Bool
 }
 
+var deletedInstancePCNode = &instancePCNode{}
+
 // instancePlanCache is a lock-free implementation of InstancePlanCache interface.
 // [key1] --> [headNode1] --> [node1] --> [node2] --> [node3]
 // [key2] --> [headNode2] --> [node4] --> [node5]
@@ -62,32 +65,45 @@ type instancePlanCache struct {
 }
 
 func (pc *instancePlanCache) getHead(key string, create bool) *instancePCNode {
-	headNode, ok := pc.heads.Load(key)
-	if ok { // cache hit
-		return headNode.(*instancePCNode)
-	}
-	if !create { // cache miss
-		return nil
-	}
-	newHeadNode := pc.createNode(nil)
-	actual, _ := pc.heads.LoadOrStore(key, newHeadNode)
-	if headNode, ok := actual.(*instancePCNode); ok { // for safety
-		return headNode
+	for {
+		headNode, ok := pc.heads.Load(key)
+		if ok { // cache hit
+			head := headNode.(*instancePCNode)
+			if !head.deleted.Load() {
+				return head
+			}
+			if !create {
+				return nil
+			}
+			pc.heads.CompareAndDelete(key, headNode)
+			continue
+		}
+		if !create { // cache miss
+			return nil
+		}
+		newHeadNode := pc.createNode(nil)
+		actual, loaded := pc.heads.LoadOrStore(key, newHeadNode)
+		if !loaded {
+			return newHeadNode
+		}
+		if headNode, ok := actual.(*instancePCNode); ok && !headNode.deleted.Load() { // for safety
+			return headNode
+		}
+		pc.heads.CompareAndDelete(key, actual)
 	}
-	return nil
 }
 
 // Get gets the cached value according to key and paramTypes.
 func (pc *instancePlanCache) Get(key string, paramTypes any) (value any, ok bool) {
 	headNode := pc.getHead(key, false)
-	if headNode == nil { // cache miss
+	if headNode == nil || headNode.deleted.Load() { // cache miss
 		return nil, false
 	}
 	return pc.getPlanFromList(headNode, paramTypes)
 }
 
 func (pc *instancePlanCache) getPlanFromList(headNode *instancePCNode, paramTypes any) (any, bool) {
-	for node := headNode.next.Load(); node != nil; node = node.next.Load() {
+	for node := headNode.next.Load(); node != nil && node != deletedInstancePCNode; node = node.next.Load() {
 		if checkTypesCompatibility4PC(node.value.ParamTypes, paramTypes) { // v.Plan is read-only, no need to lock
 			if !pc.inEvict.Load() {
 				node.lastUsed.Store(time.Now()) // atomically update the lastUsed field
@@ -109,17 +125,20 @@ func (pc *instancePlanCache) Put(key string, value, paramTypes any) (succ bool)
 		return // do nothing if it exceeds the hard limit
 	}
 	headNode := pc.getHead(key, true)
-	if headNode == nil {
+	if headNode == nil || headNode.deleted.Load() {
 		return false // for safety
 	}
 	if _, ok := pc.getPlanFromList(headNode, paramTypes); ok {
 		return // some other thread has inserted the same plan before
 	}
-	if pc.inEvict.Load() {
+	if pc.inEvict.Load() || headNode.deleted.Load() {
 		return // do nothing if eviction is in progress
 	}
 
 	firstNode := headNode.next.Load()
+	if firstNode == deletedInstancePCNode {
+		return
+	}
 	currNode := pc.createNode(value)
 	currNode.next.Store(firstNode)
 	if headNode.next.CompareAndSwap(firstNode, currNode) { // if failed, some other thread has updated this node,
@@ -130,6 +149,28 @@ func (pc *instancePlanCache) Put(key string, value, paramTypes any) (succ bool)
 	return
 }
 
+// Delete removes all cached values under the exact cache key.
+func (pc *instancePlanCache) Delete(key string) (numDeleted int) {
+	pc.evictMutex.Lock() // serialize against Evict and key deletion for safety
+	defer pc.evictMutex.Unlock()
+	pc.inEvict.Store(true)
+	defer pc.inEvict.Store(false)
+
+	headNode := pc.getHead(key, false)
+	if headNode == nil {
+		return 0
+	}
+	headNode.deleted.Store(true)
+	firstNode := headNode.next.Swap(deletedInstancePCNode)
+	pc.heads.Delete(key)
+	for node := firstNode; node != nil && node != deletedInstancePCNode; node = node.next.Load() {
+		pc.totCost.Sub(node.value.MemoryUsage())
+		pc.totPlan.Sub(1)
+		numDeleted++
+	}
+	return
+}
+
 // All returns all cached values.
 // All returned values are read-only, don't modify them.
 func (pc *instancePlanCache) All() (values []any) {
@@ -226,7 +267,7 @@ func (pc *instancePlanCache) calcEvictionThreshold(lastUsedTimes []time.Time) (t
 func (pc *instancePlanCache) foreach(callback func(prev, this *instancePCNode) (thisRemoved bool)) {
 	_, headNodes := pc.headNodes()
 	for _, headNode := range headNodes {
-		for prev, this := headNode, headNode.next.Load(); this != nil; {
+		for prev, this := headNode, headNode.next.Load(); this != nil && this != deletedInstancePCNode; {
 			thisRemoved := callback(prev, this)
 			if !thisRemoved { // this node is removed, no need to update the prev node in this case
 				prev, this = this, this.next.Load()

pkg/planner/core/plan_cache_instance_test.go

@@ -23,8 +23,10 @@ import (
 	"time"
 
 	"github.com/pingcap/tidb/pkg/domain"
+	"github.com/pingcap/tidb/pkg/parser/mysql"
 	"github.com/pingcap/tidb/pkg/planner/util/coretestsdk"
 	"github.com/pingcap/tidb/pkg/sessionctx"
+	"github.com/pingcap/tidb/pkg/types"
 	"github.com/stretchr/testify/require"
 )
 
@@ -44,6 +46,11 @@ func _miss(t *testing.T, pc sessionctx.InstancePlanCache, testKey, statsHash int
 	require.False(t, ok)
 }
 
+func _putWithParamTypes(pc sessionctx.InstancePlanCache, key string, memUsage int64, paramTypes []*types.FieldType) (succ bool) {
+	v := &PlanCacheValue{Memory: memUsage, ParamTypes: paramTypes}
+	return pc.Put(key, v, paramTypes)
+}
+
 func TestInstancePlanCacheBasic(t *testing.T) {
 	sctx := coretestsdk.MockContext()
 	defer func() {
@@ -75,6 +82,20 @@ func TestInstancePlanCacheBasic(t *testing.T) {
 	_put(pc, 1, 101, 0)
 	require.Equal(t, pc.MemUsage(), int64(100)) // the second one will be ignored
 
+	// delete an exact key
+	pc = NewInstancePlanCache(1000, 1000)
+	_put(pc, 1, 100, 0)
+	_put(pc, 2, 100, 0)
+	_put(pc, 3, 100, 0)
+	numDeleted := pc.Delete("2-0")
+	require.Equal(t, 1, numDeleted)
+	require.Equal(t, int64(200), pc.MemUsage())
+	require.Equal(t, int64(2), pc.Size())
+	_hit(t, pc, 1, 0)
+	_miss(t, pc, 2, 0)
+	_hit(t, pc, 3, 0)
+	require.Equal(t, 0, pc.Delete("not-exist"))
+
 	// eviction
 	pc = NewInstancePlanCache(320, 500)
 	_put(pc, 1, 100, 0)
@@ -158,6 +179,28 @@ func TestInstancePlanCacheWithMatchOpts(t *testing.T) {
 	_miss(t, pc, 3, 2)
 	_miss(t, pc, 3, 3)
 
+	// same exact key with different param types should be deleted together
+	pc = NewInstancePlanCache(1000, 1000)
+	key := "shared-key"
+	paramTypes1 := []*types.FieldType{types.NewFieldType(mysql.TypeLonglong)}
+	paramTypes2 := []*types.FieldType{types.NewFieldType(mysql.TypeDouble)}
+	require.True(t, _putWithParamTypes(pc, key, 100, paramTypes1))
+	require.True(t, _putWithParamTypes(pc, key, 100, paramTypes2))
+	_put(pc, 2, 100, 1)
+	_, ok := pc.Get(key, paramTypes1)
+	require.True(t, ok)
+	_, ok = pc.Get(key, paramTypes2)
+	require.True(t, ok)
+	numDeleted := pc.Delete(key)
+	require.Equal(t, 2, numDeleted)
+	require.Equal(t, int64(100), pc.MemUsage())
+	require.Equal(t, int64(1), pc.Size())
+	_, ok = pc.Get(key, paramTypes1)
+	require.False(t, ok)
+	_, ok = pc.Get(key, paramTypes2)
+	require.False(t, ok)
+	_hit(t, pc, 2, 1)
+
 	// hard limit can take effect in this case
 	pc = NewInstancePlanCache(200, 200)
 	_put(pc, 1, 100, 1)

pkg/sessionctx/context.go

@@ -61,6 +61,8 @@ type InstancePlanCache interface {
 	Get(key string, paramTypes any) (value any, ok bool)
 	// Put puts the key and value into the cache.
 	Put(key string, value, paramTypes any) (succ bool)
+	// Delete removes all cached values under the exact cache key.
+	Delete(key string) (numDeleted int)
 	// All returns all cached values.
 	// Returned values are read-only, don't modify them.
 	All() (values []any)