vecindex: lazily create root partition metadata

When preparing to update a K-means tree, we check its root partition metadata.
If that does not yet exist, we will create it using a CPut operation and then
acquire a SHARED lock that prevents a background split operation until the
transaction ends.

Epic: CRDB-42943

Release note: None

Author: andy-kimball

pkg/sql/vecindex/vecstore/BUILD.bazel

@@ -45,6 +45,7 @@ go_test(
     deps = [
         "//pkg/base",
         "//pkg/keys",
+        "//pkg/kv",
         "//pkg/kv/kvpb",
         "//pkg/security/securityassets",
         "//pkg/security/securitytest",

pkg/sql/vecindex/vecstore/store_test.go

@@ -8,10 +8,13 @@ package vecstore
 import (
 	"context"
 	"fmt"
+	"sync"
+	"sync/atomic"
 	"testing"
 
 	"github.com/cockroachdb/cockroach/pkg/base"
 	"github.com/cockroachdb/cockroach/pkg/keys"
+	"github.com/cockroachdb/cockroach/pkg/kv"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvpb"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog/catenumpb"
@@ -146,6 +149,73 @@ func TestStore(t *testing.T) {
 	// Re-run the tests with a prefixed index.
 	usePrefix = true
 	suite.Run(t, commontest.NewStoreTestSuite(ctx, makeStore))
+
+	// Ensure that races to create partition metadata either do not error, or
+	// they error with WriteTooOldError.
+	t.Run("race to create partition metadata", func(t *testing.T) {
+		store := makeStore(quantize.NewUnQuantizer(2)).(*testStore)
+
+		var done atomic.Int64
+		getMetadata := func(treeKey cspann.TreeKey) {
+			var err error
+			tx, err := store.BeginTransaction(ctx)
+			require.NoError(t, err)
+			defer func() {
+				if err != nil {
+					err = store.AbortTransaction(ctx, tx)
+				}
+			}()
+
+			// Enable stepping in the txn, which is what SQL does.
+			tx.(*Txn).kv.ConfigureStepping(ctx, kv.SteppingEnabled)
+
+			_, err = tx.GetPartitionMetadata(ctx, treeKey, cspann.RootKey, true /* forUpdate */)
+			if err != nil {
+				require.ErrorContains(t, err, "WriteTooOldError")
+				done.Store(1)
+				return
+			}
+
+			// Run GetPartitionMetadata again, to ensure that it succeeds, as a
+			// way of simulating multiple vectors being inserted in the same
+			// SQL statement.
+			_, err = tx.GetPartitionMetadata(ctx, treeKey, cspann.RootKey, true /* forUpdate */)
+			require.NoError(t, err)
+
+			err = store.CommitTransaction(ctx, tx)
+			require.NoError(t, err)
+		}
+
+		for i := range 100 {
+			var wait sync.WaitGroup
+			wait.Add(2)
+			treeKey := store.MakeTreeKey(t, i)
+			go func() {
+				defer func() {
+					wait.Done()
+				}()
+				getMetadata(treeKey)
+			}()
+			go func() {
+				defer func() {
+					wait.Done()
+				}()
+				getMetadata(treeKey)
+			}()
+			wait.Wait()
+
+			// Fail on foreground goroutine if a background goroutine failed.
+			if t.Failed() {
+				t.FailNow()
+			}
+
+			// End the test once we find at least one WriteTooOldError.
+			if done.Load() == 1 {
+				break
+			}
+		}
+	})
+
 }
 
 func TestQuantizeAndEncode(t *testing.T) {

pkg/sql/vecindex/vecstore/store_txn.go

@@ -322,11 +322,13 @@ func (tx *Txn) GetPartitionMetadata(
 			errors.Wrapf(err, "getting partition metadata for %d", partitionKey)
 	}
 
-	metadata, err := tx.extractMetadataFromKVResult(treeKey, partitionKey, &b.Results[0])
-	if err != nil {
-		return cspann.PartitionMetadata{}, err
+	// If we're preparing to update the root partition, then lazily create its
+	// metadata if it does not yet exist.
+	if forUpdate && partitionKey == cspann.RootKey && b.Results[0].Rows[0].Value == nil {
+		return tx.createRootPartition(ctx, metadataKey)
 	}
-	return metadata, nil
+
+	return tx.extractMetadataFromKVResult(treeKey, partitionKey, &b.Results[0])
 }
 
 // AddToPartition implements the cspann.Txn interface.
@@ -575,6 +577,42 @@ func (tx *Txn) GetFullVectors(
 	return err
 }
 
+// createRootPartition uses the KV CPut operation to create metadata for the
+// root partition, and then returns that metadata.
+//
+// NOTE: CPut always "sees" the latest write of the metadata record, even if the
+// timestamp of that write is higher than this transaction's.
+func (tx *Txn) createRootPartition(
+	ctx context.Context, metadataKey roachpb.Key,
+) (cspann.PartitionMetadata, error) {
+	b := tx.kv.NewBatch()
+	metadata := cspann.PartitionMetadata{Level: cspann.LeafLevel, Centroid: tx.store.emptyVec}
+	encoded, err := EncodePartitionMetadata(metadata.Level, metadata.Centroid)
+	if err != nil {
+		return cspann.PartitionMetadata{}, err
+	}
+
+	// Use CPutAllowingIfNotExists in order to handle the case where the same
+	// transaction inserts multiple vectors (e.g. multiple VALUES rows). In that
+	// case, the first row will trigger creation of the metadata record. However,
+	// subsequent inserts will not be able to "see" this record, since they will
+	// read at a lower sequence number than the metadata record was written.
+	// However, CPutAllowingIfNotExists will read at the higher sequence number
+	// and see that the record was already created.
+	//
+	// On the other hand, if a different transaction wrote the record, it will
+	// have a higher timestamp, and that will trigger a WriteTooOld error.
+	// Transactions which lose that race need to be refreshed.
+	var roachval roachpb.Value
+	roachval.SetBytes(encoded)
+	b.CPutAllowingIfNotExists(metadataKey, &roachval, roachval.TagAndDataBytes())
+	if err := tx.kv.Run(ctx, b); err != nil {
+		// Lost the race to a different transaction.
+		return cspann.PartitionMetadata{}, errors.Wrapf(err, "creating root partition metadata")
+	}
+	return metadata, nil
+}
+
 // QuantizeAndEncode quantizes the given vector (which has already been
 // randomized by the caller) with respect to the given centroid. It returns the
 // encoded form of that quantized vector.