sql/fingerprint: process range-sized spans

Subdividing each span from partition spans into n range-sized spans
ensures both that no single span passed to an ExportRequest (or some
SQL-based span fingerprinter) will run for an excessive time, allowing
the use of reasonable timeouts to catch underlying failures, and makes
the unit of progress we can complete and checkpoint more granular.
It should have no effect on the result (as tested here).

Release note: none.
Epic: none.

Author: dt

pkg/sql/fingerprint/BUILD.bazel

@@ -55,15 +55,14 @@ go_test(
         "//pkg/security/username",
         "//pkg/server",
         "//pkg/sql",
-        "//pkg/testutils",
+        "//pkg/testutils/jobutils",
         "//pkg/testutils/serverutils",
         "//pkg/testutils/sqlutils",
         "//pkg/testutils/testcluster",
         "//pkg/util/hlc",
         "//pkg/util/leaktest",
         "//pkg/util/log",
         "//pkg/util/span",
-        "@com_github_cockroachdb_errors//:errors",
         "@com_github_stretchr_testify//require",
     ],
 )

pkg/sql/fingerprint/fingerprint_job.go

@@ -78,7 +78,7 @@ func (r *resumer) Resume(ctx context.Context, execCtx interface{}) error {
 			id:      jobspb.JobID(*r),
 			execCtx: jobExecCtx,
 		},
-		partitioner: partitionSpans{jobExecCtx},
+		partitioner: partitionSpans{jobExecCtx, false},
 		fn: kvFingerprinter{
 			sender:   jobExecCtx.ExecCfg().DB.NonTransactionalSender(),
 			asOf:     details.AsOf,
@@ -225,6 +225,7 @@ func (p *persist) store(ctx context.Context, state checkpointState, frac float64
 
 type partitionSpans struct {
 	sql.JobExecContext
+	rangeSized bool
 }
 
 var _ partitioner = partitionSpans{}
@@ -249,6 +250,36 @@ func (p partitionSpans) partition(
 	if err != nil {
 		return nil, err
 	}
+	if p.rangeSized {
+		// Sub-divide each partition into indivially range-aligned sub-spans.
+		for part := range spanPartitions {
+			subdivided := make([]roachpb.Span, 0, len(spanPartitions[part].Spans))
+			for _, sp := range spanPartitions[part].Spans {
+				rdi, err := p.JobExecContext.ExecCfg().RangeDescIteratorFactory.NewLazyIterator(ctx, sp, 64)
+				if err != nil {
+					return nil, err
+				}
+				remaining := sp
+				for ; rdi.Valid(); rdi.Next() {
+					rangeDesc := rdi.CurRangeDescriptor()
+					rangeSpan := roachpb.Span{Key: rangeDesc.StartKey.AsRawKey(), EndKey: rangeDesc.EndKey.AsRawKey()}
+					subspan := remaining.Intersect(rangeSpan)
+					if !subspan.Valid() {
+						return nil, errors.AssertionFailedf("%s not in %s of %s", rangeSpan, remaining, sp)
+					}
+					subdivided = append(subdivided, subspan)
+					remaining.Key = subspan.EndKey
+				}
+				if err := rdi.Error(); err != nil {
+					return nil, err
+				}
+				if remaining.Valid() {
+					subdivided = append(subdivided, remaining)
+				}
+			}
+			spanPartitions[part].Spans = subdivided
+		}
+	}
 	return spanPartitions, nil
 }
 

pkg/sql/fingerprint/fingerprint_job_test.go

@@ -8,21 +8,21 @@ package fingerprint
 import (
 	"context"
 	"testing"
+	"time"
 
 	"github.com/cockroachdb/cockroach/pkg/base"
 	"github.com/cockroachdb/cockroach/pkg/jobs"
 	"github.com/cockroachdb/cockroach/pkg/jobs/jobspb"
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
 	"github.com/cockroachdb/cockroach/pkg/security/username"
 	"github.com/cockroachdb/cockroach/pkg/sql"
-	"github.com/cockroachdb/cockroach/pkg/testutils"
+	"github.com/cockroachdb/cockroach/pkg/testutils/jobutils"
 	"github.com/cockroachdb/cockroach/pkg/testutils/serverutils"
 	"github.com/cockroachdb/cockroach/pkg/testutils/sqlutils"
 	"github.com/cockroachdb/cockroach/pkg/util/hlc"
 	"github.com/cockroachdb/cockroach/pkg/util/leaktest"
 	"github.com/cockroachdb/cockroach/pkg/util/log"
 	"github.com/cockroachdb/cockroach/pkg/util/span"
-	"github.com/cockroachdb/errors"
 	"github.com/stretchr/testify/require"
 )
 
@@ -141,6 +141,77 @@ func TestKVFingerprinter(t *testing.T) {
 	require.Equal(t, uint64(0x2249712bb6b0388a), result.fingerprint)
 }
 
+func TestRangeSizedSpans(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+	defer log.Scope(t).Close(t)
+
+	ctx := context.Background()
+	s, db, _ := serverutils.StartServer(t, base.TestServerArgs{})
+	defer s.Stopper().Stop(ctx)
+
+	runner := sqlutils.MakeSQLRunner(db)
+
+	// Create a table with enough data to span multiple ranges.
+	runner.Exec(t, `CREATE TABLE r (i PRIMARY KEY, j) AS SELECT i, i::string FROM generate_series(1, 1000) i`)
+	runner.Exec(t, `ALTER TABLE r SPLIT AT SELECT i*10 FROM generate_series(1, 100) i`)
+	var tableID uint32
+	runner.QueryRow(t, `SELECT id FROM system.namespace WHERE name = 'r'`).Scan(&tableID)
+
+	// Use actual key encoding for the table.
+	tableSpan := roachpb.Span{
+		Key:    s.Codec().TablePrefix(tableID),
+		EndKey: s.Codec().TablePrefix(tableID).PrefixEnd(),
+	}
+
+	execCfg := s.ApplicationLayer().ExecutorConfig().(sql.ExecutorConfig)
+	jobExecCtx, cleanup := sql.MakeJobExecContext(ctx, "test", username.RootUserName(), &sql.MemoryMetrics{}, &execCfg)
+	defer cleanup()
+
+	parts := partitionSpans{JobExecContext: jobExecCtx}
+
+	g := gatherer{
+		spans: []roachpb.Span{tableSpan},
+		fn: kvFingerprinter{
+			sender:   s.DB().NonTransactionalSender(),
+			asOf:     s.Clock().Now(),
+			stripped: true,
+		}.fingerprintSpan,
+		partitioner: parts,
+		chkptFreq:   func() time.Duration { return time.Hour },
+		persister:   &testPersister{},
+	}
+
+	r1, err := g.run(ctx)
+	require.NoError(t, err)
+
+	spCount := func() int {
+		total := 0
+		for _, p := range g.mu.prog {
+			total += p.total
+		}
+		return total
+	}
+
+	// Reset g and run again to verify determinism.
+	g.persister, g.mu.done, g.mu.fingerprint, g.mu.prog = &testPersister{}, nil, 0, nil
+	r2, err := g.run(ctx)
+	require.NoError(t, err)
+	require.Equal(t, r1, r2)
+	// Single-node => partitionSpans just returned one span, so total spans = 1.
+	require.Equal(t, 1, spCount())
+
+	// Now reset g and run again, but this time with range-sized spans.
+	g.persister, g.mu.done, g.mu.fingerprint, g.mu.prog = &testPersister{}, nil, 0, nil
+	parts.rangeSized = true
+	g.partitioner = parts
+	r3, err := g.run(ctx)
+	require.NoError(t, err)
+	require.Equal(t, r2, r3)
+
+	// Verify that using range-sized spans resulted in an extra span per split.
+	require.Equal(t, 101, spCount())
+}
+
 // TestJobE2E tests that a fingerprint job can be created as adoptable and
 // successfully runs through the job system.
 func TestJobE2E(t *testing.T) {
@@ -183,21 +254,8 @@ func TestJobE2E(t *testing.T) {
 	job, err := registry.CreateAdoptableJobWithTxn(ctx, record, registry.MakeJobID(), nil /* txn */)
 	require.NoError(t, err)
 
-	// Wait for the job to be adopted and succeed.
-	testutils.SucceedsSoon(t, func() error {
-		// Nudge the adoption queue to pick up the job immediately.
-		registry.TestingNudgeAdoptionQueue()
-
-		var status string
-		runner.QueryRow(t, `SELECT status FROM system.jobs WHERE id = $1`, job.ID()).Scan(&status)
-		if status == string(jobs.StateSucceeded) {
-			return nil
-		}
-		if status == string(jobs.StateFailed) || status == string(jobs.StateCanceled) {
-			return errors.Errorf("job failed with status %s", status)
-		}
-		return errors.Errorf("job not yet succeeded, current status: %s", status)
-	})
+	registry.TestingNudgeAdoptionQueue()
+	jobutils.WaitForJobToSucceed(t, runner, job.ID())
 
 	// Verify that the job completed and has a fingerprint stored.
 	execCfg := s.ApplicationLayer().ExecutorConfig().(sql.ExecutorConfig)