inference: fix embedding runtime lifecycle race

Author: ChangRui-Ryan

.agents/skills/tidb-test-guidelines/references/domain-case-map.md

@@ -10,7 +10,7 @@
 
 ### Tests
 - `pkg/domain/db_test.go` - Tests domain session.
-- `pkg/domain/domain_test.go` - Tests domain info.
+- `pkg/domain/domain_test.go` - Tests domain info and concurrent embedding runtime lifecycle.
 - `pkg/domain/domain_utils_test.go` - Tests error code.
 - `pkg/domain/domainctx_test.go` - Tests domain ctx.
 - `pkg/domain/extract_test.go` - Tests extract plan without history view.

pkg/domain/BUILD.bazel

@@ -156,7 +156,7 @@ go_test(
     ],
     embed = [":domain"],
     flaky = True,
-    shard_count = 30,
+    shard_count = 31,
     deps = [
         "//pkg/config",
         "//pkg/ddl",

pkg/domain/domain.go

@@ -229,7 +229,7 @@ type Domain struct {
 	minJobIDRefresher *systable.MinJobIDRefresher
 
 	instancePlanCache sessionctx.InstancePlanCache // the instance level plan cache
-	embedFn           *inference.EmbedFn
+	embedFn           atomic.Pointer[inference.EmbedFn]
 
 	statsOwner owner.Manager
 

pkg/domain/domain_test.go

@@ -47,6 +47,29 @@ import (
 	"go.etcd.io/etcd/tests/v3/integration"
 )
 
+func TestEmbeddingRuntimeConcurrentLifecycle(t *testing.T) {
+	do := &Domain{}
+	for range 32 {
+		do.initInferenceProviders()
+		start := make(chan struct{})
+		readDone := make(chan struct{})
+		closeDone := make(chan struct{})
+		go func() {
+			<-start
+			_ = do.GetEmbedFn()
+			close(readDone)
+		}()
+		go func() {
+			<-start
+			do.closeInferenceProviders()
+			close(closeDone)
+		}()
+		close(start)
+		<-readDone
+		<-closeDone
+	}
+}
+
 func TestInfo(t *testing.T) {
 	t.Skip("TestInfo will hang currently, it should be fixed later")
 

pkg/domain/inference.go

@@ -17,22 +17,23 @@ package domain
 import "github.com/pingcap/tidb/pkg/inference"
 
 func (do *Domain) initInferenceProviders() {
-	do.closeInferenceProviders()
-	do.embedFn = inference.NewEmbedFn()
+	oldEmbedFn := do.embedFn.Swap(inference.NewEmbedFn())
+	if oldEmbedFn != nil {
+		oldEmbedFn.Close()
+	}
 }
 
 func (do *Domain) closeInferenceProviders() {
-	if do.embedFn == nil {
-		return
+	embedFn := do.embedFn.Swap(nil)
+	if embedFn != nil {
+		embedFn.Close()
 	}
-	do.embedFn.Close()
-	do.embedFn = nil
 }
 
 // GetEmbedFn returns the embedding function managed by this Domain.
 func (do *Domain) GetEmbedFn() *inference.EmbedFn {
 	if do == nil {
 		return nil
 	}
-	return do.embedFn
+	return do.embedFn.Load()
 }

pkg/expression/integration_test/integration_test.go

@@ -4530,7 +4530,7 @@ func TestEmbedTextFunction(t *testing.T) {
 	err := tk.QueryToErr(`select embed_text('mock/json', '[1, 3, 4]')`)
 	require.ErrorContains(t, err, "EMBED_TEXT is only supported in starter deployment mode")
 	if !enableStarterDeployModeForEmbeddingTest(t) {
-		return
+		t.Skip("EMBED_TEXT functional tests require nextgen kernel in starter deployment mode")
 	}
 
 	err = tk.QueryToErr("select embed_text('text-embedding-3', 'hello world')")

pkg/inference/sqlembed_test.go

@@ -27,6 +27,8 @@ import (
 	"github.com/stretchr/testify/require"
 )
 
+const embeddingTestTimeout = 10 * time.Second
+
 type staticEmbedder struct {
 	embeddings [][]float32
 	err        error
@@ -197,13 +199,8 @@ func TestEmbedFnSharedCallCancellation(t *testing.T) {
 		result2 <- result{embedding: embedding, err: err}
 	}()
 	require.Eventually(t, func() bool {
-		embedFn.mu.Lock()
-		defer embedFn.mu.Unlock()
-		for _, call := range embedFn.inFlight {
-			return call.waiters == 2
-		}
-		return false
-	}, time.Second, 10*time.Millisecond)
+		return hasSingleInFlightCallWithWaiters(embedFn, 2)
+	}, embeddingTestTimeout, 10*time.Millisecond)
 
 	firstCause := errors.New("first caller canceled")
 	cancel1(firstCause)
@@ -245,21 +242,16 @@ func TestEmbedFnCancelsProviderAfterAllCallersCancel(t *testing.T) {
 		err2 <- err
 	}()
 	require.Eventually(t, func() bool {
-		embedFn.mu.Lock()
-		defer embedFn.mu.Unlock()
-		for _, call := range embedFn.inFlight {
-			return call.waiters == 2
-		}
-		return false
-	}, time.Second, 10*time.Millisecond)
+		return hasSingleInFlightCallWithWaiters(embedFn, 2)
+	}, embeddingTestTimeout, 10*time.Millisecond)
 
 	cancel1()
 	cancel2()
 	require.ErrorIs(t, receiveFromChannel(t, err1, "first caller cancellation"), context.Canceled)
 	require.ErrorIs(t, receiveFromChannel(t, err2, "second caller cancellation"), context.Canceled)
 	select {
 	case <-provider.canceled:
-	case <-time.After(time.Second):
+	case <-time.After(embeddingTestTimeout):
 		t.Fatal("provider request was not canceled after all callers canceled")
 	}
 }
@@ -278,12 +270,24 @@ func waitForChannel(t *testing.T, ch <-chan struct{}, description string) {
 	receiveFromChannel(t, ch, description)
 }
 
+func hasSingleInFlightCallWithWaiters(embedFn *EmbedFn, waiters int) bool {
+	embedFn.mu.Lock()
+	defer embedFn.mu.Unlock()
+	if len(embedFn.inFlight) != 1 {
+		return false
+	}
+	for _, call := range embedFn.inFlight {
+		return call.waiters == waiters
+	}
+	return false
+}
+
 func receiveFromChannel[T any](t *testing.T, ch <-chan T, description string) T {
 	t.Helper()
 	select {
 	case value := <-ch:
 		return value
-	case <-time.After(time.Second):
+	case <-time.After(embeddingTestTimeout):
 		t.Fatalf("timed out waiting for %s", description)
 		var zero T
 		return zero