vecindex: fix race condition when creating root partition

When multiple workers are racing to create the root partition, it's
possible to trigger an error when the root partition already exists
and is in a state other than Ready (e.g. Splitting). Fix this by
creating the root metadata record only if doesn't already exist.
If it does exist, just return the current metadata.

Epic: CRDB-42943

Release note: None

Author: andy-kimball

pkg/sql/vecindex/cspann/partition_metadata.go

@@ -147,6 +147,9 @@ type PartitionStateDetails struct {
 	//   be merged into the root partition.
 	Source PartitionKey
 	// Timestamp is the time of the last state transition for the partition.
+	// NOTE: We use NowNoMono to get the current time because the monotonic clock
+	// reading is not round-tripped through the encoding/decoding functions, since
+	// it's not useful to store.
 	Timestamp time.Time
 }
 

pkg/sql/vecindex/vecencoding/encoding.go

@@ -17,9 +17,14 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/util/vector"
 )
 
-/* Vector indexes are encoded as follows:
+/* Vector indexes are encoded as shown below.
 
-Metadata KV key:
+NOTE: Key formats always are always suffixed by a Family ID byte of 0. This is
+necessary to include so that the key can be parsed by the KV range split code.
+That code uses the family ID to check that it doesn't split column families for
+the same row across ranges.
+
+Metadata KV Key:
   Metadata keys always sort before vector keys, since Family ID 0 always
   sorts before Level, which is >= 1.
   ┌────────────┬──────────────┬────────────┬───────────┐
@@ -29,11 +34,11 @@ Metadata KV Value:
   ┌─────┬─────┬───────┬───────┬──────┬─────────┬────────┐
   │Level│State│Target1│Target2│Source|Timestamp│Centroid|
   └─────┴─────┴───────┴───────┴──────┴─────────┴────────┘
-Vector KV key (interior, non-leaf partition):
+Vector KV Key (interior, non-leaf partition):
   ┌────────────┬──────────────┬────────────┬─────┬──────────────────┬───────────┐
   │Index Prefix│Prefix Columns│PartitionKey│Level│Child PartitionKey│Family ID 0│
   └────────────┴──────────────┴────────────┴─────┴──────────────────┴───────────┘
-Vector KV key (leaf partition):
+Vector KV Key (leaf partition):
   ┌────────────┬──────────────┬────────────┬─────┬──────────┬───────────┐
   │Index Prefix│Prefix Columns│PartitionKey│Level│PrimaryKey│Family ID 0│
   └────────────┴──────────────┴────────────┴─────┴──────────┴───────────┘
@@ -147,18 +152,17 @@ func (vik *DecodedVectorKey) Encode(appendTo []byte) []byte {
 }
 
 // EncodeVectorValue takes a quantized vector entry and any composite key data
-// and returns the byte slice encoding the value side of the vector index entry.
-// This value will still need to be further encoded as Bytes in the
-// valueside.Value
+// and returns the byte slice encoding the value of the vector index entry. This
+// value will still need to be further encoded as Bytes in valueside.Value.
 func EncodeVectorValue(appendTo []byte, vectorData []byte, compositeData []byte) []byte {
-	// The value side is encoded as a concatenation of the vector data and the
+	// The value is encoded as a concatenation of the vector data and the
 	// composite data.
 	appendTo = append(appendTo, vectorData...)
 	return append(appendTo, compositeData...)
 }
 
 // EncodedVectorValueLen returns the number of bytes needed to encode the value
-// side of a vector index entry.
+// of a vector index entry.
 func EncodedVectorValueLen(vectorData []byte, compositeData []byte) int {
 	return len(vectorData) + len(compositeData)
 }

pkg/sql/vecindex/vecstore/store_test.go

@@ -179,9 +179,6 @@ func TestStore(t *testing.T) {
 				return
 			}
 
-			// Step so that we can see the new root partition.
-			require.NoError(t, tx.(*Txn).kv.Step(ctx, true /* allowReadTimestampStep */))
-
 			// Run GetPartitionMetadata again, to ensure that it succeeds, as a
 			// way of simulating multiple vectors being inserted in the same
 			// SQL statement.

pkg/sql/vecindex/vecstore/store_txn.go

@@ -304,22 +304,48 @@ func (tx *Txn) DeletePartition(
 func (tx *Txn) GetPartitionMetadata(
 	ctx context.Context, treeKey cspann.TreeKey, partitionKey cspann.PartitionKey, forUpdate bool,
 ) (cspann.PartitionMetadata, error) {
-	// TODO(mw5h): Add to an existing batch instead of starting a new one.
-	b := tx.kv.NewBatch()
-
 	metadataKey := vecencoding.EncodeMetadataKey(tx.store.prefix, treeKey, partitionKey)
-	if forUpdate {
-		// By acquiring a shared lock on metadata key, we prevent splits/merges of
-		// this partition from conflicting with the add operation.
-		b.GetForShare(metadataKey, tx.lockDurability)
-	} else {
-		b.Get(metadataKey)
-	}
 
-	// Run the batch and get the partition metadata from results.
-	if err := tx.kv.Run(ctx, b); err != nil {
-		return cspann.PartitionMetadata{},
-			errors.Wrapf(err, "getting partition metadata for %d", partitionKey)
+	// By acquiring a shared lock on metadata key, we prevent splits/merges of
+	// this partition from conflicting with the add operation.
+	b, err := func() (b *kv.Batch, err error) {
+		// TODO(mw5h): Add to an existing batch instead of starting a new one.
+		b = tx.kv.NewBatch()
+
+		if tx.kv.Sender().GetSteppingMode(ctx) == kv.SteppingEnabled {
+			// When there are multiple inserts within the same SQL statement, the
+			// first insert 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. Handle this issue by temporarily stepping the read sequence
+			// number so the latest metadata can be read.
+			prevSeqNum := tx.kv.GetReadSeqNum()
+			if err = tx.kv.Step(ctx, false /* allowReadTimestampStep */); err != nil {
+				return nil, err
+			}
+			defer func() {
+				// Restore the original sequence number.
+				if readErr := tx.kv.SetReadSeqNum(prevSeqNum); err != nil {
+					err = errors.CombineErrors(err, readErr)
+				}
+			}()
+		}
+
+		if forUpdate {
+			b.GetForShare(metadataKey, tx.lockDurability)
+		} else {
+			b.Get(metadataKey)
+		}
+
+		// Run the batch.
+		if err := tx.kv.Run(ctx, b); err != nil {
+			return nil, errors.Wrapf(err, "getting partition metadata for %d", partitionKey)
+		}
+
+		return b, nil
+	}()
+	if err != nil {
+		return cspann.PartitionMetadata{}, err
 	}
 
 	// If we're preparing to update the root partition, then lazily create its
@@ -591,10 +617,10 @@ func (tx *Txn) GetFullVectors(
 }
 
 // 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.
+// root partition, and then returns that metadata. If another transaction races
+// and creates the root partition at a higher timestamp, createRootPartition
+// returns a WriteTooOld error, which will trigger a refresh of this transaction
+// at the higher timestamp.
 func (tx *Txn) createRootPartition(
 	ctx context.Context, metadataKey roachpb.Key,
 ) (cspann.PartitionMetadata, error) {
@@ -606,22 +632,11 @@ func (tx *Txn) createRootPartition(
 	}
 	encoded := vecencoding.EncodeMetadataValue(metadata)
 
-	// 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())
+	// Use CPut to detect the case where another transaction is racing to create
+	// the root partition. CPut always "sees" the latest version of the metadata
+	// record.
+	b.CPut(metadataKey, encoded, nil /* expValue */)
 	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