migrations: Add IndexStatus to store (sourcegraph#30270)

Author: efritz

internal/database/migration/store/observability.go

@@ -10,6 +10,7 @@ import (
 type Operations struct {
 	down              *observation.Operation
 	ensureSchemaTable *observation.Operation
+	indexStatus       *observation.Operation
 	lock              *observation.Operation
 	tryLock           *observation.Operation
 	up                *observation.Operation
@@ -35,6 +36,7 @@ func NewOperations(observationContext *observation.Context) *Operations {
 	return &Operations{
 		down:              op("Down"),
 		ensureSchemaTable: op("EnsureSchemaTable"),
+		indexStatus:       op("IndexStatus"),
 		lock:              op("Lock"),
 		tryLock:           op("TryLock"),
 		up:                op("Up"),

internal/database/migration/store/store.go

@@ -23,6 +23,45 @@ type Store struct {
 	operations      *Operations
 }
 
+// IndexStatus describes the state of an index. Is{Valid,Ready,Live} is taken
+// from the `pg_index` system table. If the index is currently being created,
+// then the remaining reference fields will be populated describing the index
+// creation progress.
+type IndexStatus struct {
+	IsValid      bool
+	IsReady      bool
+	IsLive       bool
+	Phase        *string
+	LockersDone  *int
+	LockersTotal *int
+	BlocksDone   *int
+	BlocksTotal  *int
+	TuplesDone   *int
+	TuplesTotal  *int
+}
+
+// CreateIndexConcurrentlyPhases is an ordered list of phases that occur during
+// a CREATE INDEX CONCURRENTLY operation. The phase of an ongoing operation can
+// found in the system view `view pg_stat_progress_create_index` (since PG 12).
+//
+// If the phase value found in the system view may not match these values exactly
+// and may only indicate a prefix. The phase may have more specific information
+// following the initial phase description. Do not compare phase values exactly.
+//
+// See https://www.postgresql.org/docs/12/progress-reporting.html#CREATE-INDEX-PROGRESS-REPORTING.
+var CreateIndexConcurrentlyPhases = []string{
+	"initializing",
+	"waiting for writers before build",
+	"building index",
+	"waiting for writers before validation",
+	"index validation: scanning index",
+	"index validation: sorting tuples",
+	"index validation: scanning table",
+	"waiting for old snapshots",
+	"waiting for readers before marking dead",
+	"waiting for readers before dropping",
+}
+
 func NewWithDB(db dbutil.DB, migrationsTable string, operations *Operations) *Store {
 	return &Store{
 		Store:           basestore.NewWithDB(db, sql.TxOptions{}),
@@ -198,6 +237,36 @@ func (s *Store) Down(ctx context.Context, definition definition.Definition) (err
 	return nil
 }
 
+// IndexStatus returns an object describing the current validity status and creation progress of the
+// index with the given name. If the index does not exist, a false-valued flag is returned.
+func (s *Store) IndexStatus(ctx context.Context, tableName, indexName string) (_ IndexStatus, _ bool, err error) {
+	ctx, endObservation := s.operations.indexStatus.With(ctx, &err, observation.Args{})
+	defer endObservation(1, observation.Args{})
+
+	return scanFirstIndexStatus(s.Query(ctx, sqlf.Sprintf(indexStatusQuery, tableName, indexName)))
+}
+
+const indexStatusQuery = `
+-- source: internal/database/migration/store/store.go:IndexStatus
+SELECT
+	pi.indisvalid,
+	pi.indisready,
+	pi.indislive,
+	p.phase,
+	p.lockers_total,
+	p.lockers_done,
+	p.blocks_total,
+	p.blocks_done,
+	p.tuples_total,
+	p.tuples_done
+FROM pg_stat_all_indexes ai
+JOIN pg_index pi ON pi.indexrelid = ai.indexrelid
+LEFT JOIN pg_stat_progress_create_index p ON p.relid = ai.relid AND p.index_relid = ai.indexrelid
+WHERE
+	ai.relname = %s AND
+	ai.indexrelname = %s
+`
+
 func (s *Store) runMigrationQuery(ctx context.Context, definitionVersion int, up bool, query *sqlf.Query) (err error) {
 	targetVersion := definitionVersion
 	expectedCurrentVersion := definitionVersion - 1
@@ -329,3 +398,32 @@ func scanMigrationLogs(rows *sql.Rows, queryErr error) (_ []migrationLog, err er
 
 	return logs, nil
 }
+
+// scanFirstIndexStatus scans a slice of index status objects from the return value of `*Store.query`.
+func scanFirstIndexStatus(rows *sql.Rows, queryErr error) (status IndexStatus, _ bool, err error) {
+	if queryErr != nil {
+		return IndexStatus{}, false, queryErr
+	}
+	defer func() { err = basestore.CloseRows(rows, err) }()
+
+	if rows.Next() {
+		if err := rows.Scan(
+			&status.IsValid,
+			&status.IsReady,
+			&status.IsLive,
+			&status.Phase,
+			&status.LockersDone,
+			&status.LockersTotal,
+			&status.BlocksDone,
+			&status.BlocksTotal,
+			&status.TuplesDone,
+			&status.TuplesTotal,
+		); err != nil {
+			return IndexStatus{}, false, err
+		}
+
+		return status, true, nil
+	}
+
+	return IndexStatus{}, false, nil
+}

internal/database/migration/store/store_test.go

@@ -2,11 +2,14 @@ package store
 
 import (
 	"context"
+	"database/sql"
 	"strings"
 	"testing"
+	"time"
 
 	"github.com/google/go-cmp/cmp"
 	"github.com/keegancsmith/sqlf"
+	"golang.org/x/sync/errgroup"
 
 	"github.com/sourcegraph/sourcegraph/internal/database/basestore"
 	"github.com/sourcegraph/sourcegraph/internal/database/dbtest"
@@ -425,6 +428,159 @@ func TestDown(t *testing.T) {
 	})
 }
 
+func TestIndexStatus(t *testing.T) {
+	db := dbtest.NewDB(t)
+	store := testStore(db)
+	ctx := context.Background()
+
+	if _, err := db.ExecContext(ctx, "CREATE TABLE tbl (id text, name text);"); err != nil {
+		t.Fatalf("unexpected error: %s", err)
+	}
+
+	// Index does not (yet) exist
+	if _, ok, err := store.IndexStatus(ctx, "tbl", "idx"); err != nil {
+		t.Fatalf("unexpected error: %s", err)
+	} else if ok {
+		t.Fatalf("unexpected index status")
+	}
+
+	// Wrap context in a small timeout; we do tight for-loops here to determine
+	// when we can continue on to/unblock the next operation, but none of the
+	// steps should take any significant time.
+	ctx, cancel := context.WithTimeout(ctx, time.Second*10)
+	group, groupCtx := errgroup.WithContext(ctx)
+	defer cancel()
+
+	whileEmpty := func(ctx context.Context, conn dbutil.DB, query string) error {
+		for {
+			rows, err := conn.QueryContext(ctx, query)
+			if err != nil {
+				return err
+			}
+
+			lockVisible := rows.Next()
+
+			if err := basestore.CloseRows(rows, nil); err != nil {
+				return err
+			}
+
+			if lockVisible {
+				return nil
+			}
+		}
+	}
+
+	// Create separate connections to precise control contention of resources
+	// so we can examine what this method returns while an index is being created.
+
+	conns := make([]*sql.Conn, 3)
+	for i := 0; i < 3; i++ {
+		conn, err := db.Conn(ctx)
+		if err != nil {
+			t.Fatalf("failed to open new connection: %s", err)
+		}
+		t.Cleanup(func() { conn.Close() })
+
+		conns[i] = conn
+	}
+	connA, connB, connC := conns[0], conns[1], conns[2]
+
+	lockQuery := `SELECT pg_advisory_lock(10, 10)`
+	unlockQuery := `SELECT pg_advisory_unlock(10, 10)`
+	createIndexQuery := `CREATE INDEX CONCURRENTLY idx ON tbl(id)`
+
+	// Session A
+	// Successfully take and hold advisory lock
+	if _, err := connA.ExecContext(ctx, lockQuery); err != nil {
+		t.Fatalf("unexpected error: %s", err)
+	}
+
+	// Session B
+	// Try to take advisory lock; blocked by Session A
+	group.Go(func() error {
+		_, err := connB.ExecContext(groupCtx, lockQuery)
+		return err
+	})
+
+	// Session C
+	// try to create index concurrently; blocked by session B waiting on session A
+	group.Go(func() error {
+		// Wait until we can see Session B's lock before attempting to create index
+		if err := whileEmpty(groupCtx, connC, "SELECT 1 FROM pg_locks WHERE locktype = 'advisory' AND NOT granted"); err != nil {
+			t.Fatalf("unexpected error: %s", err)
+		}
+
+		_, err := connC.ExecContext(groupCtx, createIndexQuery)
+		return err
+	})
+
+	// Wait until we can see Session C's lock before querying index status
+	if err := whileEmpty(ctx, db, "SELECT 1 FROM pg_locks WHERE locktype = 'relation' AND mode = 'ShareUpdateExclusiveLock'"); err != nil {
+		t.Fatalf("unexpected error: %s", err)
+	}
+
+	// "waiting for old snapshots" will be the phase that is blocked by the concurrent
+	// sessions holding advisory locks. We may happen to hit one of the earlier phases
+	// if we're quick enough, so we'll keep polling progress until we hit the target.
+	blockingPhase := "waiting for old snapshots"
+	nonblockingPhasePrefixes := make([]string, 0, len(CreateIndexConcurrentlyPhases))
+	for _, prefix := range CreateIndexConcurrentlyPhases {
+		if prefix == blockingPhase {
+			break
+		}
+
+		nonblockingPhasePrefixes = append(nonblockingPhasePrefixes, prefix)
+	}
+	compareWithPrefix := func(value, prefix string) bool {
+		return value == prefix || strings.HasPrefix(value, prefix+":")
+	}
+
+retryLoop:
+	for {
+		if status, ok, err := store.IndexStatus(ctx, "tbl", "idx"); err != nil {
+			t.Fatalf("unexpected error: %s", err)
+		} else if !ok {
+			t.Fatalf("expected index status")
+		} else if status.Phase == nil {
+			t.Fatalf("unexpected phase. want=%q have=nil", blockingPhase)
+		} else if *status.Phase == blockingPhase {
+			break
+		} else {
+			for _, prefix := range nonblockingPhasePrefixes {
+				if compareWithPrefix(*status.Phase, prefix) {
+					continue retryLoop
+				}
+			}
+
+			t.Fatalf("unexpected phase. want=%q have=%q", blockingPhase, *status.Phase)
+		}
+	}
+
+	// Session A
+	// Unlock, unblocking both Session B and Session C
+	if _, err := connA.ExecContext(ctx, unlockQuery); err != nil {
+		t.Fatalf("unexpected error: %s", err)
+	}
+
+	// Wait for index creation to complete
+	if err := group.Wait(); err != nil {
+		t.Fatalf("unexpected error: %s", err)
+	}
+
+	if status, ok, err := store.IndexStatus(ctx, "tbl", "idx"); err != nil {
+		t.Fatalf("unexpected error: %s", err)
+	} else if !ok {
+		t.Fatalf("expected index status")
+	} else {
+		if !status.IsValid {
+			t.Fatalf("unexpected isvalid. want=%v have=%v", true, status.IsValid)
+		}
+		if status.Phase != nil {
+			t.Fatalf("unexpected phase. want=%v have=%v", nil, status.Phase)
+		}
+	}
+}
+
 func testStore(db dbutil.DB) *Store {
 	return NewWithDB(db, "test_migrations_table", NewOperations(&observation.TestContext))
 }