Merge #144520 #144570

144520: optbuilder: disable mutation of vector indexes while backfilling r=mw5h a=mw5h

Merging concurrent writes on vector indexes while backfill is in progress is somewhat more complicated problem than can be solved in the 25.2 timeframe. In order to ensure data integrity, we will disable mutation on vector indexes while backfill is in progress.

Informs: #144443
Release note: None

144570: sqlproxyccl/acl: use newer btree r=RaduBerinde a=RaduBerinde

Use `BtreeG` in the acl watcher code.
Informs: #144504
Release note: None

Co-authored-by: Matt White <[email protected]>
Co-authored-by: Radu Berinde <[email protected]>

Author: 3 people

go.mod

@@ -34,7 +34,7 @@ require (
 	github.com/golang/mock v1.6.0
 	github.com/golang/protobuf v1.5.4
 	github.com/golang/snappy v0.0.5-0.20231225225746-43d5d4cd4e0e
-	github.com/google/btree v1.1.3 // indirect
+	github.com/google/btree v1.1.3
 	github.com/google/pprof v0.0.0-20240227163752-401108e1b7e7
 	github.com/google/uuid v1.6.0 // indirect
 	google.golang.org/api v0.114.0

pkg/ccl/sqlproxyccl/acl/BUILD.bazel

@@ -21,9 +21,9 @@ go_library(
         "//pkg/util/syncutil",
         "//pkg/util/timeutil",
         "@com_github_cockroachdb_errors//:errors",
+        "@com_github_google_btree//:btree",
         "@com_github_pires_go_proxyproto//:go-proxyproto",
         "@com_github_pires_go_proxyproto//tlvparse",
-        "@com_github_raduberinde_btree//:btree",
         "@in_gopkg_yaml_v2//:yaml_v2",
     ],
 )
@@ -50,7 +50,6 @@ go_test(
         "@com_github_cockroachdb_errors//:errors",
         "@com_github_pires_go_proxyproto//:go-proxyproto",
         "@com_github_pires_go_proxyproto//tlvparse",
-        "@com_github_raduberinde_btree//:btree",
         "@com_github_stretchr_testify//require",
         "@in_gopkg_yaml_v2//:yaml_v2",
     ],

pkg/ccl/sqlproxyccl/acl/watcher.go

@@ -9,11 +9,11 @@ import (
 	"context"
 	"time"
 
-	"github.com/RaduBerinde/btree" // TODO(#144504): switch to the newer btree
 	"github.com/cockroachdb/cockroach/pkg/util/log"
 	"github.com/cockroachdb/cockroach/pkg/util/metric"
 	"github.com/cockroachdb/cockroach/pkg/util/syncutil"
 	"github.com/cockroachdb/cockroach/pkg/util/timeutil"
+	"github.com/google/btree"
 )
 
 // Watcher maintains a list of connections waiting for changes to the
@@ -33,7 +33,7 @@ type Watcher struct {
 	options *aclOptions
 
 	// All of the listeners waiting for changes to the access control list.
-	listeners *btree.BTree
+	listeners *btree.BTreeG[*listener]
 
 	// These control whether or not a connection is allowd based on it's
 	// ConnectionTags.
@@ -128,7 +128,7 @@ func NewWatcher(ctx context.Context, opts ...Option) (*Watcher, error) {
 		opt(options)
 	}
 	w := &Watcher{
-		listeners:   btree.New(8),
+		listeners:   btree.NewG[*listener](8, func(a, b *listener) bool { return a.id < b.id }),
 		options:     options,
 		controllers: make([]AccessController, 0),
 	}
@@ -235,7 +235,7 @@ func (w *Watcher) addAccessController(
 func (w *Watcher) updateAccessController(
 	ctx context.Context, index int, controller AccessController,
 ) {
-	var copy *btree.BTree
+	var copy *btree.BTreeG[*listener]
 	var controllers []AccessController
 	func() {
 		w.mu.Lock()
@@ -304,14 +304,10 @@ func (w *Watcher) removeListener(l *listener) {
 	w.listeners.Delete(l)
 }
 
-// Less implements the btree.Item interface for listener.
-func (l *listener) Less(than btree.Item) bool {
-	return l.id < than.(*listener).id
-}
-
-func checkListeners(ctx context.Context, listeners *btree.BTree, controllers []AccessController) {
-	listeners.Ascend(func(i btree.Item) bool {
-		lst := i.(*listener)
+func checkListeners(
+	ctx context.Context, listeners *btree.BTreeG[*listener], controllers []AccessController,
+) {
+	listeners.Ascend(func(lst *listener) bool {
 		if err := checkConnection(ctx, lst.connection, controllers); err != nil {
 			lst.mu.Lock()
 			defer lst.mu.Unlock()

pkg/ccl/sqlproxyccl/acl/watcher_test.go

@@ -11,7 +11,6 @@ import (
 	"testing"
 	"time"
 
-	"github.com/RaduBerinde/btree" // TODO(#144504): switch to the newer btree
 	"github.com/cockroachdb/cockroach/pkg/ccl/sqlproxyccl/tenant"
 	"github.com/cockroachdb/cockroach/pkg/ccl/sqlproxyccl/tenantdirsvr"
 	"github.com/cockroachdb/cockroach/pkg/ccl/testutilsccl"
@@ -403,8 +402,8 @@ func TestACLWatcher(t *testing.T) {
 		require.Nil(t, err)
 
 		var l *listener
-		watcher.listeners.Ascend(func(i btree.Item) bool {
-			l = i.(*listener)
+		watcher.listeners.Ascend(func(lst *listener) bool {
+			l = lst
 			return false
 		})
 		require.NotNil(t, l)

pkg/sql/alter_primary_key.go

@@ -54,6 +54,15 @@ func (p *planner) AlterPrimaryKey(
 	alterPKNode tree.AlterTableAlterPrimaryKey,
 	alterPrimaryKeyLocalitySwap *alterPrimaryKeyLocalitySwap,
 ) error {
+	// Check if sql_safe_updates is enabled and the table has vector indexes
+	if p.EvalContext().SessionData().SafeUpdates {
+		for _, idx := range tableDesc.AllIndexes() {
+			if idx.GetType() == idxtype.VECTOR {
+				return pgerror.DangerousStatementf("ALTER PRIMARY KEY on a table with vector indexes will disable writes to the table while the index is being rebuilt")
+			}
+		}
+	}
+
 	if err := paramparse.ValidateUniqueConstraintParams(
 		alterPKNode.StorageParams,
 		paramparse.UniqueConstraintParamContext{

pkg/sql/backfill/backfill_test.go

@@ -7,6 +7,7 @@ package backfill_test
 
 import (
 	"context"
+	"sync"
 	"testing"
 
 	"github.com/cockroachdb/cockroach/pkg/base"
@@ -54,16 +55,18 @@ func TestVectorColumnAndIndexBackfill(t *testing.T) {
 	sqlDB.Exec(t, `
 		CREATE TABLE vectors (
 			id INT PRIMARY KEY,
-			vec VECTOR(3)
+			vec VECTOR(3),
+			data INT
 		)
 	`)
 
 	// Insert 200 rows with random vector data
 	sqlDB.Exec(t, `
-		INSERT INTO vectors (id, vec)
+		INSERT INTO vectors (id, vec, data)
 		SELECT 
 			generate_series(1, 200) as id,
-			ARRAY[random(), random(), random()]::vector(3) as vec
+			ARRAY[random(), random(), random()]::vector(3) as vec,
+			generate_series(1, 200) as data
 	`)
 
 	// Create a vector index on the vector column
@@ -88,3 +91,87 @@ func TestVectorColumnAndIndexBackfill(t *testing.T) {
 	// be 200 in most cases.
 	require.Greater(t, matchCount, 190, "Expected to find at least 190 similar vectors")
 }
+
+func TestConcurrentOperationsDuringVectorIndexCreation(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+	defer log.Scope(t).Close(t)
+
+	// Channel to block the backfill process
+	blockBackfill := make(chan struct{})
+	backfillBlocked := make(chan struct{})
+
+	ctx := context.Background()
+	srv, db, _ := serverutils.StartServer(t, base.TestServerArgs{
+		Knobs: base.TestingKnobs{
+			DistSQL: &execinfra.TestingKnobs{
+				// Block the backfill process after the first batch
+				RunAfterBackfillChunk: func() {
+					backfillBlocked <- struct{}{}
+					<-blockBackfill
+				},
+			},
+		},
+	})
+	defer srv.Stopper().Stop(ctx)
+	sqlDB := sqlutils.MakeSQLRunner(db)
+
+	// Create a table with a vector column
+	sqlDB.Exec(t, `
+		CREATE TABLE vectors (
+			id INT PRIMARY KEY,
+			vec VECTOR(3),
+			data INT
+		)
+	`)
+
+	// Insert some initial data
+	sqlDB.Exec(t, `
+		INSERT INTO vectors (id, vec, data) VALUES
+		(1, '[1, 2, 3]', 100),
+		(2, '[4, 5, 6]', 200),
+		(3, '[7, 8, 9]', 300)
+	`)
+
+	wg := sync.WaitGroup{}
+	wg.Add(1)
+
+	// Start creating the vector index in a goroutine
+	var createIndexErr error
+	go func() {
+		defer wg.Done()
+		_, createIndexErr = db.ExecContext(ctx, `CREATE VECTOR INDEX vec_idx ON vectors (vec)`)
+	}()
+
+	// Wait for the backfill to be blocked
+	<-backfillBlocked
+
+	// Attempt concurrent operations while the index is being created
+	// These should fail with the appropriate error
+	_, err := db.ExecContext(ctx, `UPDATE vectors SET vec = '[10, 11, 12]', data = 150 WHERE id = 1`)
+	require.Error(t, err)
+	require.Contains(t, err.Error(), "Cannot write to a vector index while it is being built")
+
+	_, err = db.ExecContext(ctx, `INSERT INTO vectors (id, vec, data) VALUES (4, '[13, 14, 15]', 400)`)
+	require.Error(t, err)
+	require.Contains(t, err.Error(), "Cannot write to a vector index while it is being built")
+
+	_, err = db.ExecContext(ctx, `DELETE FROM vectors WHERE id = 2`)
+	require.Error(t, err)
+	require.Contains(t, err.Error(), "Cannot write to a vector index while it is being built")
+
+	// This update should succeed since it only modifies the data column
+	_, err = db.ExecContext(ctx, `UPDATE vectors SET data = 250 WHERE id = 2`)
+	require.NoError(t, err, "Updating only the data column should succeed during index creation")
+
+	// Unblock the backfill process
+	close(blockBackfill)
+
+	wg.Wait()
+	// Wait for the index creation to complete
+	require.NoError(t, createIndexErr)
+
+	// Verify the index was created successfully
+	var id int
+	sqlDB.QueryRow(t, `SELECT id FROM vectors@vec_idx ORDER BY vec <-> '[1, 2, 3]' LIMIT 1`).Scan(&id)
+	require.Equal(t, 1, id)
+}

pkg/sql/catalog/tabledesc/index.go

@@ -464,12 +464,7 @@ func (w index) CreatedAt() time.Time {
 	return timeutil.Unix(0, w.desc.CreatedAtNanos)
 }
 
-// IsTemporaryIndexForBackfill returns true iff the index is
-// an index being used as the temporary index being used by an
-// in-progress index backfill.
-//
-// TODO(ssd): This could be its own boolean or we could store the ID
-// of the index it is a temporary index for.
+// IsTemporaryIndexForBackfill implements the cat.Index interface.
 func (w index) IsTemporaryIndexForBackfill() bool {
 	return w.desc.UseDeletePreservingEncoding
 }

pkg/sql/create_index.go

@@ -57,6 +57,12 @@ func (p *planner) CreateIndex(ctx context.Context, n *tree.CreateIndex) (planNod
 	); err != nil {
 		return nil, err
 	}
+
+	// Check if sql_safe_updates is enabled and this is a vector index
+	if n.Type == idxtype.VECTOR && p.EvalContext().SessionData().SafeUpdates {
+		return nil, pgerror.DangerousStatementf("CREATE VECTOR INDEX will disable writes to the table while the index is being built")
+	}
+
 	_, tableDesc, err := p.ResolveMutableTableDescriptor(
 		ctx, &n.Table, true /*required*/, tree.ResolveRequireTableOrViewDesc,
 	)

pkg/sql/logictest/testdata/logic_test/vector_index

@@ -2,15 +2,30 @@
 # CREATE TABLE/INDEX tests.
 # ------------------------------------------------------------------------------
 
+# Test guardrail for vector index creation.
+# TODO(mw5h): remove these two statements once online modifications are supported.
+statement ok
+SET sql_safe_updates = true
+
 # Simple vector index.
 statement ok
 CREATE TABLE simple (
   a INT PRIMARY KEY,
+  b INT NOT NULL,
   vec1 VECTOR(3),
   VECTOR INDEX (vec1),
   FAMILY (a, vec1)
 )
 
+statement error pgcode 01000 pq: rejected \(sql_safe_updates = true\): CREATE VECTOR INDEX will disable writes to the table while the index is being built
+CREATE VECTOR INDEX ON simple (vec1)
+
+statement error pgcode 01000 pq: rejected \(sql_safe_updates = true\): ALTER PRIMARY KEY on a table with vector indexes will disable writes to the table while the index is being rebuilt
+ALTER TABLE simple ALTER PRIMARY KEY USING COLUMNS (b)
+
+statement ok
+SET sql_safe_updates = false
+
 statement ok
 CREATE VECTOR INDEX ON simple (vec1)
 
@@ -23,12 +38,13 @@ SHOW CREATE TABLE simple
 ----
 simple  CREATE TABLE public.simple (
           a INT8 NOT NULL,
+          b INT8 NOT NULL,
           vec1 VECTOR(3) NULL,
           CONSTRAINT simple_pkey PRIMARY KEY (a ASC),
           VECTOR INDEX simple_vec1_idx (vec1),
           VECTOR INDEX simple_vec1_idx1 (vec1),
           VECTOR INDEX simple_vec1_idx2 (vec1),
-          FAMILY fam_0_a_vec1 (a, vec1)
+          FAMILY fam_0_a_vec1_b (a, vec1, b)
         )
 
 statement ok

pkg/sql/opt/cat/index.go

@@ -204,6 +204,10 @@ type Index interface {
 	// Partition returns the ith PARTITION BY LIST partition within the index
 	// definition, where i < PartitionCount.
 	Partition(i int) Partition
+
+	// IsTemporaryIndexForBackfill returns true iff the index is an index being
+	// used as the temporary index being used by an in-progress index backfill.
+	IsTemporaryIndexForBackfill() bool
 }
 
 // IndexColumn describes a single column that is part of an index definition.