wip queuefeed skel

pkg/BUILD.bazel

@@ -586,6 +586,7 @@ ALL_TESTS = [
     "//pkg/sql/privilege:privilege_test",
     "//pkg/sql/protoreflect:protoreflect_test",
     "//pkg/sql/querycache:querycache_test",
+    "//pkg/sql/queuefeed:queuefeed_test",
     "//pkg/sql/randgen:randgen_test",
     "//pkg/sql/regions:regions_test",
     "//pkg/sql/row:row_disallowed_imports_test",
@@ -2228,6 +2229,9 @@ GO_TARGETS = [
     "//pkg/sql/protoreflect:protoreflect_test",
     "//pkg/sql/querycache:querycache",
     "//pkg/sql/querycache:querycache_test",
+    "//pkg/sql/queuefeed/queuebase:queuebase",
+    "//pkg/sql/queuefeed:queuefeed",
+    "//pkg/sql/queuefeed:queuefeed_test",
     "//pkg/sql/randgen:randgen",
     "//pkg/sql/randgen:randgen_test",
     "//pkg/sql/rangeprober:range_prober",

pkg/cmd/roachtest/tests/BUILD.bazel

@@ -167,6 +167,7 @@ go_library(
         "ptp.go",
         "query_comparison_util.go",
         "queue.go",
+        "queuefeed.go",
         "quit.go",
         "rapid_restart.go",
         "rebalance_load.go",
@@ -365,6 +366,7 @@ go_library(
         "@org_golang_google_protobuf//proto",
         "@org_golang_x_exp//maps",
         "@org_golang_x_oauth2//clientcredentials",
+        "@org_golang_x_sync//errgroup",
         "@org_golang_x_text//cases",
         "@org_golang_x_text//language",
     ],

pkg/cmd/roachtest/tests/queuefeed.go

@@ -0,0 +1,145 @@
+// Copyright 2025 The Cockroach Authors.
+//
+// Use of this software is governed by the CockroachDB Software License
+// included in the /LICENSE file.
+
+package tests
+
+import (
+	"context"
+	"database/sql"
+	"fmt"
+	"math/rand"
+	"strings"
+	"sync/atomic"
+	"time"
+
+	"github.com/cockroachdb/cockroach/pkg/cmd/roachtest/cluster"
+	"github.com/cockroachdb/cockroach/pkg/cmd/roachtest/option"
+	"github.com/cockroachdb/cockroach/pkg/cmd/roachtest/registry"
+	"github.com/cockroachdb/cockroach/pkg/cmd/roachtest/spec"
+	"github.com/cockroachdb/cockroach/pkg/cmd/roachtest/test"
+	"github.com/cockroachdb/cockroach/pkg/roachprod/install"
+	"github.com/cockroachdb/errors"
+	"github.com/stretchr/testify/require"
+	"golang.org/x/sync/errgroup"
+)
+
+func registerQueuefeed(r registry.Registry) {
+	r.Add(registry.TestSpec{
+		Name:    "queuefeed",
+		Owner:   registry.OwnerCDC,
+		Cluster: r.MakeClusterSpec(4, spec.WorkloadNode()),
+		Run: func(ctx context.Context, t test.Test, c cluster.Cluster) {
+			runQueuefeed(ctx, t, c)
+		},
+		CompatibleClouds: registry.AllClouds,
+		Suites:           registry.Suites(registry.Nightly),
+	})
+}
+
+func runQueuefeed(ctx context.Context, t test.Test, c cluster.Cluster) {
+	c.Start(ctx, t.L(), option.DefaultStartOpts(), install.MakeClusterSettings(), c.CRDBNodes())
+
+	db := c.Conn(ctx, t.L(), 1)
+	defer db.Close()
+
+	_, err := db.ExecContext(ctx, "SET CLUSTER SETTING kv.rangefeed.enabled = true")
+	require.NoError(t, errors.Wrap(err, "enabling rangefeeds"))
+
+	t.Status("initializing kv workload")
+	c.Run(ctx, option.WithNodes(c.WorkloadNode()),
+		"./cockroach workload init kv --splits=100 {pgurl:1}")
+
+	var tableID int64
+	err = db.QueryRowContext(ctx, "SELECT id FROM system.namespace WHERE name = 'kv' and \"parentSchemaID\" <> 0;").Scan(&tableID)
+	require.NoError(t, err)
+
+	t.Status("creating kv_queue")
+	_, err = db.ExecContext(ctx, "SELECT crdb_internal.create_queue_feed('kv_queue', $1)", tableID)
+	require.NoError(t, err)
+
+	t.Status("running queue feed queries")
+
+	ctx, cancel := context.WithTimeout(ctx, 10*time.Minute)
+	defer cancel()
+
+	g, ctx := errgroup.WithContext(ctx)
+
+	const numReaders = 10
+	counters := make([]*atomic.Int64, numReaders)
+	for i := range counters {
+		counters[i] = &atomic.Int64{}
+	}
+
+	g.Go(func() error {
+		return c.RunE(ctx, option.WithNodes(c.WorkloadNode()),
+			"./cockroach workload run kv --duration=10m {pgurl:1}")
+	})
+
+	g.Go(func() error {
+		ticker := time.NewTicker(1 * time.Second)
+		defer ticker.Stop()
+		lastCounts := make([]int64, numReaders)
+
+		for {
+			select {
+			case <-ctx.Done():
+				return ctx.Err()
+			case <-ticker.C:
+				qps := make([]string, numReaders)
+				for i := 0; i < numReaders; i++ {
+					currentCount := counters[i].Load()
+					ratePerSec := currentCount - lastCounts[i]
+					qps[i] = fmt.Sprintf("%d", ratePerSec)
+					lastCounts[i] = currentCount
+				}
+				t.L().Printf("qps: %s", strings.Join(qps, ","))
+			}
+		}
+	})
+
+	dbNodes := 1 // TODO fix bug that occurs with 3
+	nodePool := make([]*sql.DB, numReaders)
+	for i := range dbNodes {
+		nodePool[i] = c.Conn(ctx, t.L(), i+1)
+	}
+	defer func() {
+		for i := range dbNodes {
+			_ = nodePool[i].Close()
+		}
+	}()
+
+	for i := 0; i < numReaders; i++ {
+		readerIndex := i
+		g.Go(func() error {
+			// Stagger the readers a bit. This helps test re-distribution of
+			// partitions.
+			// TODO fix bug that occurs with jitter
+			// time.Sleep(time.Duration(rand.Intn(int(time.Minute))))
+
+			// Connect to a random node to simulate a tcp load balancer.
+			conn, err := nodePool[rand.Intn(dbNodes)].Conn(ctx)
+			if err != nil {
+				return errors.Wrap(err, "getting connection for the queuefeed reader")
+			}
+			defer func() { _ = conn.Close() }()
+
+			for ctx.Err() == nil {
+				var count int
+				err := conn.QueryRowContext(ctx,
+					"SELECT count(*) FROM crdb_internal.select_from_queue_feed('kv_queue', 10000)").Scan(&count)
+				if err != nil {
+					return err
+				}
+				counters[readerIndex].Add(int64(count))
+			}
+			return ctx.Err()
+		})
+	}
+
+	err = g.Wait()
+	if err != nil && ctx.Err() == nil {
+		t.Fatal(err)
+	}
+}

pkg/cmd/roachtest/tests/registry.go

@@ -131,6 +131,7 @@ func RegisterTests(r registry.Registry) {
 	registerPruneDanglingSnapshotsAndDisks(r)
 	registerPTP(r)
 	registerQueue(r)
+	registerQueuefeed(r)
 	registerQuitTransfersLeases(r)
 	registerRebalanceLoad(r)
 	registerReplicaGC(r)

pkg/kv/txn.go

@@ -90,6 +90,9 @@ type Txn struct {
 	// commitTriggers are run upon successful commit.
 	commitTriggers []func(ctx context.Context)
 
+	// rollbackTriggers are run upon rollback/abort.
+	rollbackTriggers []func(ctx context.Context)
+
 	// mu holds fields that need to be synchronized for concurrent request execution.
 	mu struct {
 		syncutil.Mutex
@@ -1093,6 +1096,16 @@ func (txn *Txn) AddCommitTrigger(trigger func(ctx context.Context)) {
 	txn.commitTriggers = append(txn.commitTriggers, trigger)
 }
 
+// AddRollbackTrigger adds a closure to be executed on rollback/abort
+// of the transaction.
+func (txn *Txn) AddRollbackTrigger(trigger func(ctx context.Context)) {
+	if txn.typ != RootTxn {
+		panic(errors.AssertionFailedf("AddRollbackTrigger() called on leaf txn"))
+	}
+
+	txn.rollbackTriggers = append(txn.rollbackTriggers, trigger)
+}
+
 // endTxnReqAlloc is used to batch the heap allocations of an EndTxn request.
 type endTxnReqAlloc struct {
 	req      kvpb.EndTxnRequest
@@ -1243,6 +1256,9 @@ func (txn *Txn) PrepareForRetry(ctx context.Context) error {
 	// Reset commit triggers. These must be reconfigured by the client during the
 	// next retry.
 	txn.commitTriggers = nil
+	// Reset rollback triggers. These must be reconfigured by the client during the
+	// next retry.
+	txn.rollbackTriggers = nil
 
 	txn.mu.Lock()
 	defer txn.mu.Unlock()
@@ -1383,9 +1399,17 @@ func (txn *Txn) Send(
 	if pErr == nil {
 		// Invoking the commit triggers here ensures they run even in the case when a
 		// commit request is issued manually (not via Commit).
-		if et, ok := ba.GetArg(kvpb.EndTxn); ok && et.(*kvpb.EndTxnRequest).Commit {
-			for _, t := range txn.commitTriggers {
-				t(ctx)
+		if et, ok := ba.GetArg(kvpb.EndTxn); ok {
+			if et.(*kvpb.EndTxnRequest).Commit {
+				for _, t := range txn.commitTriggers {
+					t(ctx)
+				}
+			} else {
+				// Invoking the rollback triggers here ensures they run even in the case when a
+				// rollback request is issued manually (not via Rollback).
+				for _, t := range txn.rollbackTriggers {
+					t(ctx)
+				}
 			}
 		}
 		return br, nil

pkg/server/BUILD.bazel

@@ -259,6 +259,7 @@ go_library(
         "//pkg/sql/physicalplan",
         "//pkg/sql/privilege",
         "//pkg/sql/querycache",
+        "//pkg/sql/queuefeed",
         "//pkg/sql/rangeprober",
         "//pkg/sql/regions",
         "//pkg/sql/rolemembershipcache",

pkg/server/server_sql.go

@@ -89,6 +89,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/sql/optionalnodeliveness"
 	"github.com/cockroachdb/cockroach/pkg/sql/pgwire"
 	"github.com/cockroachdb/cockroach/pkg/sql/querycache"
+	"github.com/cockroachdb/cockroach/pkg/sql/queuefeed"
 	"github.com/cockroachdb/cockroach/pkg/sql/rangeprober"
 	"github.com/cockroachdb/cockroach/pkg/sql/regions"
 	"github.com/cockroachdb/cockroach/pkg/sql/rolemembershipcache"
@@ -1062,6 +1063,7 @@ func newSQLServer(ctx context.Context, cfg sqlServerArgs) (*SQLServer, error) {
 		TenantReadOnly:             cfg.SQLConfig.TenantReadOnly,
 		CidrLookup:                 cfg.BaseConfig.CidrLookup,
 		LicenseEnforcer:            cfg.SQLConfig.LicenseEnforcer,
+		QueueManager:               queuefeed.NewManager(ctx, cfg.internalDB, cfg.rangeFeedFactory, cfg.rangeDescIteratorFactory, codec, leaseMgr, cfg.sqlLivenessProvider.CachedReader()),
 	}
 
 	if codec.ForSystemTenant() {
@@ -1791,6 +1793,11 @@ func (s *SQLServer) preStart(
 
 	s.startLicenseEnforcer(ctx, knobs)
 
+	// Close queue manager when the stopper stops.
+	stopper.AddCloser(stop.CloserFn(func() {
+		s.execCfg.QueueManager.Close()
+	}))
+
 	// Report a warning if the server is being shut down via the stopper
 	// before it was gracefully drained. This warning may be innocuous
 	// in tests where there is no use of the test server/cluster after

pkg/sql/BUILD.bazel

@@ -478,6 +478,8 @@ go_library(
         "//pkg/sql/privilege",
         "//pkg/sql/protoreflect",
         "//pkg/sql/querycache",
+        "//pkg/sql/queuefeed",
+        "//pkg/sql/queuefeed/queuebase",
         "//pkg/sql/regionliveness",
         "//pkg/sql/regions",
         "//pkg/sql/rolemembershipcache",