Merge #82555 #82569 #82652 #82683

82555: sql: fix CREATE TABLE LIKE with implicit pk r=jasonmchan a=jasonmchan

Previously, `CREATE TABLE LIKE` copied implicitly created columns (e.g.
for the rowid default primary key and hash sharded index). Defaults for
some of these columns were not properly copied over in some cases,
causing unexpected constraint violations to surface.

This commit fixes this by skipping copying such columns; instead, they
will be freshly created. Followup work is needed for REGIONAL BY ROW.

Fixes #82401

Release note: None

82569: sql/schemachanger/rel,scplan/rules: add support for rules, _; adopt r=ajwerner a=ajwerner

The first commit extends the `rel` language with support for rules and `_` and adopts it for the dep rules. The second commit contains further cleanup and adopts in op rules. 

Release note: None

82652: ccl/sqlproxyccl: fix inaccurate CurConnCount metric due to goroutine leak r=JeffSwenson a=jaylim-crl

Previously, there was a possibility where a processor can return from
resuming because the client's connection was closed _before_ waitResumed even
has the chance to wake up to check on the resumed field. When that happens,
the connection goroutine will be blocked forever, and the CurConnCount metric
will never be decremented, even if the connection has already been terminated.

When the client's connection was closed, the forwarder's context will be
cancelled as well. The ideal behavior would be to terminate all waiters when
that happens, but the current code does not do that. This commit fixes that
issue by adding a new closed state to the processors, and ensuring that the
processor is closed whenever resume returns with an error. waitResumed can
then check on this state before going back to wait.

Release note: None

82683: server: don't re-run node decommissioning callback r=aayushshah15 a=aayushshah15

This commit fixes a bug from #80993. Without this commit, nodes
might re-run the callback to enqueue a decommissioning node's ranges
into their replicate queues if they received a gossip update from that
decommissioning node that was perceived to be newer. Re-running this
callback on every newer gossip update from a decommissioning node will
be too expensive for nodes with a lot of replicas.

Release note: None

Co-authored-by: Jason Chan <[email protected]>
Co-authored-by: Andrew Werner <[email protected]>
Co-authored-by: Jay <[email protected]>
Co-authored-by: Aayush Shah <[email protected]>

pkg/ccl/sqlproxyccl/forwarder.go

@@ -202,6 +202,8 @@ func (f *forwarder) Context() context.Context {
 //
 // Close implements the balancer.ConnectionHandle interface.
 func (f *forwarder) Close() {
+	// Cancelling the forwarder's context and connections will automatically
+	// cause the processors to exit, and close themselves.
 	f.ctxCancel()
 
 	// Whenever Close is called while both of the processors are suspended, the
@@ -389,7 +391,10 @@ func makeLogicalClockFn() func() uint64 {
 // cancellation of dials.
 var aLongTimeAgo = timeutil.Unix(1, 0)
 
-var errProcessorResumed = errors.New("processor has already been resumed")
+var (
+	errProcessorResumed = errors.New("processor has already been resumed")
+	errProcessorClosed  = errors.New("processor has been closed")
+)
 
 // processor must always be constructed through newProcessor.
 type processor struct {
@@ -402,6 +407,7 @@ type processor struct {
 	mu struct {
 		syncutil.Mutex
 		cond       *sync.Cond
+		closed     bool
 		resumed    bool
 		inPeek     bool
 		suspendReq bool // Indicates that a suspend has been requested.
@@ -424,13 +430,15 @@ func newProcessor(logicalClockFn func() uint64, src, dst *interceptor.PGConn) *p
 
 // resume starts the processor and blocks during the processing. When the
 // processing has been terminated, this returns nil if the processor can be
-// resumed again in the future. If an error (except errProcessorResumed) was
-// returned, the processor should not be resumed again, and the forwarder should
-// be closed.
-func (p *processor) resume(ctx context.Context) error {
+// resumed again in the future. If an error was returned, the processor should
+// not be resumed again, and the forwarder must be closed.
+func (p *processor) resume(ctx context.Context) (retErr error) {
 	enterResume := func() error {
 		p.mu.Lock()
 		defer p.mu.Unlock()
+		if p.mu.closed {
+			return errProcessorClosed
+		}
 		if p.mu.resumed {
 			return errProcessorResumed
 		}
@@ -441,6 +449,10 @@ func (p *processor) resume(ctx context.Context) error {
 	exitResume := func() {
 		p.mu.Lock()
 		defer p.mu.Unlock()
+		// If there's an error, close the processor.
+		if retErr != nil {
+			p.mu.closed = true
+		}
 		p.mu.resumed = false
 		p.mu.cond.Broadcast()
 	}
@@ -495,6 +507,9 @@ func (p *processor) resume(ctx context.Context) error {
 	}
 
 	if err := enterResume(); err != nil {
+		if errors.Is(err, errProcessorResumed) {
+			return nil
+		}
 		return err
 	}
 	defer exitResume()
@@ -524,6 +539,9 @@ func (p *processor) waitResumed(ctx context.Context) error {
 		if ctx.Err() != nil {
 			return ctx.Err()
 		}
+		if p.mu.closed {
+			return errProcessorClosed
+		}
 		p.mu.cond.Wait()
 	}
 	return nil
@@ -536,6 +554,11 @@ func (p *processor) suspend(ctx context.Context) error {
 	p.mu.Lock()
 	defer p.mu.Unlock()
 
+	// If the processor has been closed, it cannot be suspended at all.
+	if p.mu.closed {
+		return errProcessorClosed
+	}
+
 	defer func() {
 		if p.mu.suspendReq {
 			p.mu.suspendReq = false

pkg/ccl/sqlproxyccl/forwarder_test.go

@@ -521,12 +521,15 @@ func TestSuspendResumeProcessor(t *testing.T) {
 			interceptor.NewPGConn(serverProxy),
 		)
 		require.EqualError(t, p.resume(ctx), context.Canceled.Error())
+		p.mu.Lock()
+		require.True(t, p.mu.closed)
+		p.mu.Unlock()
 
 		// Set resumed to true to simulate suspend loop.
 		p.mu.Lock()
 		p.mu.resumed = true
 		p.mu.Unlock()
-		require.EqualError(t, p.suspend(ctx), context.Canceled.Error())
+		require.EqualError(t, p.suspend(ctx), errProcessorClosed.Error())
 	})
 
 	t.Run("wait_for_resumed", func(t *testing.T) {
@@ -586,15 +589,15 @@ func TestSuspendResumeProcessor(t *testing.T) {
 			interceptor.NewPGConn(serverProxy),
 		)
 
-		// Ensure that everything will return a resumed error except 1.
+		// Ensure that two resume calls will return right away.
 		errCh := make(chan error, 2)
 		go func() { errCh <- p.resume(ctx) }()
 		go func() { errCh <- p.resume(ctx) }()
 		go func() { errCh <- p.resume(ctx) }()
 		err := <-errCh
-		require.EqualError(t, err, errProcessorResumed.Error())
+		require.NoError(t, err)
 		err = <-errCh
-		require.EqualError(t, err, errProcessorResumed.Error())
+		require.NoError(t, err)
 
 		// Suspend the last goroutine.
 		err = p.waitResumed(ctx)
@@ -604,7 +607,7 @@ func TestSuspendResumeProcessor(t *testing.T) {
 
 		// Validate suspension.
 		err = <-errCh
-		require.Nil(t, err)
+		require.NoError(t, err)
 		p.mu.Lock()
 		require.False(t, p.mu.resumed)
 		require.False(t, p.mu.inPeek)
@@ -694,10 +697,7 @@ func TestSuspendResumeProcessor(t *testing.T) {
 		// Wait until all resume calls except 1 have returned.
 		for i := 0; i < concurrency-1; i++ {
 			err := <-errResumeCh
-			// If error is not nil, it has to be an already resumed error.
-			if err != nil {
-				require.EqualError(t, err, errProcessorResumed.Error())
-			}
+			require.NoError(t, err)
 		}
 
 		// Wait until the last one returns. We can guarantee that this is for

pkg/ccl/sqlproxyccl/proxy_handler.go

@@ -407,7 +407,7 @@ func (handler *proxyHandler) handle(ctx context.Context, incomingConn *proxyConn
 	if err := f.run(fe.Conn, crdbConn); err != nil {
 		// Don't send to the client here for the same reason below.
 		handler.metrics.updateForError(err)
-		return err
+		return errors.Wrap(err, "running forwarder")
 	}
 
 	// Block until an error is received, or when the stopper starts quiescing,

pkg/ccl/sqlproxyccl/proxy_handler_test.go

@@ -1325,6 +1325,76 @@ func TestConnectionMigration(t *testing.T) {
 	}, 10*time.Second, 100*time.Millisecond)
 }
 
+// TestCurConnCountMetric ensures that the CurConnCount metric is accurate.
+// Previously, there was a regression where the CurConnCount metric wasn't
+// decremented whenever the connections were closed due to a goroutine leak.
+func TestCurConnCountMetric(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+	defer log.Scope(t).Close(t)
+
+	ctx := context.Background()
+
+	// Start KV server.
+	params, _ := tests.CreateTestServerParams()
+	s, _, _ := serverutils.StartServer(t, params)
+	defer s.Stopper().Stop(ctx)
+
+	// Start a single SQL pod.
+	tenantID := serverutils.TestTenantID()
+	tenants := startTestTenantPods(ctx, t, s, tenantID, 1)
+	defer func() {
+		for _, tenant := range tenants {
+			tenant.Stopper().Stop(ctx)
+		}
+	}()
+
+	// Register the SQL pod in the directory server.
+	tds := tenantdirsvr.NewTestStaticDirectoryServer(s.Stopper(), nil /* timeSource */)
+	tds.CreateTenant(tenantID, "tenant-cluster")
+	tds.AddPod(tenantID, &tenant.Pod{
+		TenantID:       tenantID.ToUint64(),
+		Addr:           tenants[0].SQLAddr(),
+		State:          tenant.RUNNING,
+		StateTimestamp: timeutil.Now(),
+	})
+	require.NoError(t, tds.Start(ctx))
+
+	opts := &ProxyOptions{SkipVerify: true, DisableConnectionRebalancing: true}
+	opts.testingKnobs.directoryServer = tds
+	proxy, addr := newSecureProxyServer(ctx, t, s.Stopper(), opts)
+	connectionString := fmt.Sprintf("postgres://testuser:hunter2@%s/?sslmode=require&options=--cluster=tenant-cluster-%s", addr, tenantID)
+
+	// Open 500 connections to the SQL pod.
+	const numConns = 500
+	var wg sync.WaitGroup
+	wg.Add(numConns)
+	for i := 0; i < numConns; i++ {
+		go func() {
+			defer wg.Done()
+
+			// Opens a new connection, runs SELECT 1, and closes it right away.
+			// Ignore all connection errors.
+			conn, err := pgx.Connect(ctx, connectionString)
+			if err != nil {
+				return
+			}
+			_ = conn.Ping(ctx)
+			_ = conn.Close(ctx)
+		}()
+	}
+	wg.Wait()
+
+	// Ensure that the CurConnCount metric gets decremented to 0 whenever all
+	// the connections are closed.
+	testutils.SucceedsSoon(t, func() error {
+		val := proxy.metrics.CurConnCount.Value()
+		if val == 0 {
+			return nil
+		}
+		return errors.Newf("expected CurConnCount=0, but got %d", val)
+	})
+}
+
 func TestClusterNameAndTenantFromParams(t *testing.T) {
 	defer leaktest.AfterTest(t)()
 	defer log.Scope(t).Close(t)

pkg/server/admin_test.go

@@ -2444,6 +2444,19 @@ func TestDecommissionEnqueueReplicas(t *testing.T) {
 
 		// Ensure that the scratch range's replica was proactively enqueued.
 		require.Equal(t, <-enqueuedRangeIDs, tc.LookupRangeOrFatal(t, scratchKey).RangeID)
+
+		// Check that the node was marked as decommissioning in each of the nodes'
+		// decommissioningNodeMap. This needs to be wrapped in a SucceedsSoon to
+		// deal with gossip propagation delays.
+		testutils.SucceedsSoon(t, func() error {
+			for i := 0; i < tc.NumServers(); i++ {
+				srv := tc.Server(i)
+				if _, exists := srv.DecommissioningNodeMap()[decommissioningSrv.NodeID()]; !exists {
+					return errors.Newf("node %d not detected to be decommissioning", decommissioningSrv.NodeID())
+				}
+			}
+			return nil
+		})
 	}
 
 	decommissionAndCheck(2 /* decommissioningSrvIdx */)

pkg/server/decommission.go

@@ -54,6 +54,7 @@ func (t *decommissioningNodeMap) makeOnNodeDecommissioningCallback(
 			// Nothing more to do.
 			return
 		}
+		t.nodes[decommissioningNodeID] = struct{}{}
 
 		logLimiter := log.Every(5 * time.Second) // avoid log spam
 		if err := stores.VisitStores(func(store *kvserver.Store) error {
@@ -216,3 +217,15 @@ func (s *Server) Decommission(
 	}
 	return nil
 }
+
+// DecommissioningNodeMap returns the set of node IDs that are decommissioning
+// from the perspective of the server.
+func (s *Server) DecommissioningNodeMap() map[roachpb.NodeID]interface{} {
+	s.decomNodeMap.RLock()
+	defer s.decomNodeMap.RUnlock()
+	nodes := make(map[roachpb.NodeID]interface{})
+	for key, val := range s.decomNodeMap.nodes {
+		nodes[key] = val
+	}
+	return nodes
+}

pkg/server/server.go

@@ -123,6 +123,7 @@ type Server struct {
 	admin           *adminServer
 	status          *statusServer
 	drain           *drainServer
+	decomNodeMap    *decommissioningNodeMap
 	authentication  *authenticationServer
 	migrationServer *migrationServer
 	tsDB            *ts.DB
@@ -844,6 +845,7 @@ func NewServer(cfg Config, stopper *stop.Stopper) (*Server, error) {
 		admin:                  sAdmin,
 		status:                 sStatus,
 		drain:                  drain,
+		decomNodeMap:           decomNodeMap,
 		authentication:         sAuth,
 		tsDB:                   tsDB,
 		tsServer:               &sTS,

pkg/sql/create_table.go

@@ -2509,18 +2509,11 @@ func replaceLikeTableOpts(n *tree.CreateTable, params runParams) (tree.TableDefs
 		// This is required to ensure the newly created table still works as expected
 		// as these columns are required for certain features to work when used
 		// as an index.
+		// TODO(#82672): We shouldn't need this. This is only still required for
+		// the REGIONAL BY ROW column.
 		shouldCopyColumnDefaultSet := make(map[string]struct{})
 		if opts.Has(tree.LikeTableOptIndexes) {
 			for _, idx := range td.NonDropIndexes() {
-				// Copy the rowid default if it was created implicitly by not specifying
-				// PRIMARY KEY.
-				if idx.Primary() && td.IsPrimaryIndexDefaultRowID() {
-					for i := 0; i < idx.NumKeyColumns(); i++ {
-						shouldCopyColumnDefaultSet[idx.GetKeyColumnName(i)] = struct{}{}
-					}
-				}
-				// Copy any implicitly created columns (e.g. hash-sharded indexes,
-				// REGIONAL BY ROW).
 				for i := 0; i < idx.ExplicitColumnStartIdx(); i++ {
 					for i := 0; i < idx.NumKeyColumns(); i++ {
 						shouldCopyColumnDefaultSet[idx.GetKeyColumnName(i)] = struct{}{}
@@ -2530,12 +2523,15 @@ func replaceLikeTableOpts(n *tree.CreateTable, params runParams) (tree.TableDefs
 		}
 
 		defs := make(tree.TableDefs, 0)
-		// Add all columns. Columns are always added.
+		// Add user-defined columns.
 		for i := range td.Columns {
 			c := &td.Columns[i]
-			if c.Inaccessible {
-				// Inaccessible columns automatically get added by
-				// the system; we don't need to add them ourselves here.
+			implicit, err := isImplicitlyCreatedBySystem(td, c)
+			if err != nil {
+				return nil, err
+			}
+			if implicit {
+				// Don't add system-created implicit columns.
 				continue
 			}
 			def := tree.ColumnTableDef{
@@ -2615,6 +2611,11 @@ func replaceLikeTableOpts(n *tree.CreateTable, params runParams) (tree.TableDefs
 		}
 		if opts.Has(tree.LikeTableOptIndexes) {
 			for _, idx := range td.NonDropIndexes() {
+				if idx.Primary() && td.IsPrimaryIndexDefaultRowID() {
+					// We won't copy over the default rowid primary index; instead
+					// we'll just generate a new one.
+					continue
+				}
 				indexDef := tree.IndexTableDef{
 					Name:     tree.Name(idx.GetName()),
 					Inverted: idx.GetType() == descpb.IndexDescriptor_INVERTED,
@@ -2883,3 +2884,23 @@ func validateUniqueConstraintParamsForCreateTableAs(n *tree.CreateTable) error {
 	}
 	return nil
 }
+
+// Checks if the column was automatically added by the system (e.g. for a rowid
+// primary key or hash sharded index).
+func isImplicitlyCreatedBySystem(td *tabledesc.Mutable, c *descpb.ColumnDescriptor) (bool, error) {
+	// TODO(#82672): add check for REGIONAL BY ROW column
+	if td.IsPrimaryIndexDefaultRowID() && c.ID == td.GetPrimaryIndex().GetKeyColumnID(0) {
+		return true, nil
+	}
+	col, err := td.FindColumnWithID(c.ID)
+	if err != nil {
+		return false, err
+	}
+	if td.IsShardColumn(col) {
+		return true, nil
+	}
+	if c.Inaccessible {
+		return true, nil
+	}
+	return false, nil
+}