jobfrontier: add GetAllResolvedSpans helper function

This patch adds the helper function `GetResolvedSpans`, which is like
`GetResolvedSpans` except it collects and returns the resolved spans
for all frontiers associated with a job ID.

Release note: None

Author: andyyang890

pkg/jobs/jobfrontier/frontier.go

@@ -85,6 +85,31 @@ func GetResolvedSpans(
 	return spans, found, nil
 }
 
+// GetAllResolvedSpans is like GetResolvedSpans except it collects and returns
+// the resolved spans for all frontiers associated with a job ID. The returned
+// spans are not sorted or combined in any way.
+func GetAllResolvedSpans(
+	ctx context.Context, txn isql.Txn, jobID jobspb.JobID,
+) ([]jobspb.ResolvedSpan, bool, error) {
+	infoStorage := jobs.InfoStorageForJob(txn, jobID)
+
+	var found bool
+	var spans []jobspb.ResolvedSpan
+	if err := infoStorage.Iterate(ctx, frontierPrefix, func(_ string, value []byte) error {
+		found = true
+		var r jobspb.ResolvedSpans
+		if err := protoutil.Unmarshal(value, &r); err != nil {
+			return err
+		}
+		spans = append(spans, r.ResolvedSpans...)
+		return nil
+	}); err != nil || !found {
+		return nil, false, err
+	}
+
+	return spans, found, nil
+}
+
 // Store persists a frontier's current state to storage.
 //
 // All span entries in the frontier and their current timestamps will be

pkg/jobs/jobfrontier/frontier_test.go

@@ -303,3 +303,164 @@ func TestStoreChunked(t *testing.T) {
 		})
 	}
 }
+
+func TestGetAllResolvedSpans(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+	defer log.Scope(t).Close(t)
+
+	ctx := context.Background()
+	srv := serverutils.StartServerOnly(t, base.TestServerArgs{})
+	defer srv.Stopper().Stop(ctx)
+
+	s := srv.ApplicationLayer()
+	sp := func(i int) roachpb.Span {
+		return roachpb.Span{
+			Key:    s.Codec().TablePrefix(100 + uint32(i)),
+			EndKey: s.Codec().TablePrefix(101 + uint32(i)),
+		}
+	}
+	sp1, sp2, sp3 := sp(0), sp(1), sp(2)
+
+	ts1, ts2 := hlc.Timestamp{WallTime: 100, Logical: 1}, hlc.Timestamp{WallTime: 200, Logical: 2}
+
+	tests := []struct {
+		name     string
+		setup    func(t *testing.T, txn isql.Txn, jobID jobspb.JobID)
+		expected []jobspb.ResolvedSpan
+	}{
+		{
+			name:     "no frontiers",
+			setup:    func(t *testing.T, txn isql.Txn, jobID jobspb.JobID) {},
+			expected: nil,
+		},
+		{
+			name: "single frontier",
+			setup: func(t *testing.T, txn isql.Txn, jobID jobspb.JobID) {
+				f, err := span.MakeFrontier(sp1, sp2, sp3)
+				require.NoError(t, err)
+				_, err = f.Forward(sp1, ts1)
+				require.NoError(t, err)
+				require.NoError(t, Store(ctx, txn, jobID, "single", f))
+			},
+			expected: []jobspb.ResolvedSpan{
+				{Span: sp1, Timestamp: ts1},
+				{Span: sp2},
+				{Span: sp3},
+			},
+		},
+		{
+			name: "multiple frontiers",
+			setup: func(t *testing.T, txn isql.Txn, jobID jobspb.JobID) {
+				// First frontier with sp1@ts1 and sp2@zero.
+				f1, err := span.MakeFrontier(sp1, sp2)
+				require.NoError(t, err)
+				_, err = f1.Forward(sp1, ts1)
+				require.NoError(t, err)
+				require.NoError(t, Store(ctx, txn, jobID, "frontier1", f1))
+
+				// Second frontier with sp3@ts2.
+				f2, err := span.MakeFrontier(sp3)
+				require.NoError(t, err)
+				_, err = f2.Forward(sp3, ts2)
+				require.NoError(t, err)
+				require.NoError(t, Store(ctx, txn, jobID, "frontier2", f2))
+			},
+			expected: []jobspb.ResolvedSpan{
+				{Span: sp1, Timestamp: ts1},
+				{Span: sp2},
+				{Span: sp3, Timestamp: ts2},
+			},
+		},
+		{
+			name: "multiple frontiers with overlapping spans",
+			setup: func(t *testing.T, txn isql.Txn, jobID jobspb.JobID) {
+				// First frontier with sp1@ts1.
+				f1, err := span.MakeFrontier(sp1)
+				require.NoError(t, err)
+				_, err = f1.Forward(sp1, ts1)
+				require.NoError(t, err)
+				require.NoError(t, Store(ctx, txn, jobID, "overlap1", f1))
+
+				// Second frontier with sp1@ts2 (different timestamp) and sp2@zero.
+				f2, err := span.MakeFrontier(sp1, sp2)
+				require.NoError(t, err)
+				_, err = f2.Forward(sp1, ts2)
+				require.NoError(t, err)
+				require.NoError(t, Store(ctx, txn, jobID, "overlap2", f2))
+			},
+			expected: []jobspb.ResolvedSpan{
+				{Span: sp1, Timestamp: ts1},
+				{Span: sp1, Timestamp: ts2},
+				{Span: sp2},
+			},
+		},
+		{
+			name: "multiple frontiers with some deleted",
+			setup: func(t *testing.T, txn isql.Txn, jobID jobspb.JobID) {
+				// Store first frontier with sp1@ts1.
+				f1, err := span.MakeFrontier(sp1)
+				require.NoError(t, err)
+				_, err = f1.Forward(sp1, ts1)
+				require.NoError(t, err)
+				require.NoError(t, Store(ctx, txn, jobID, "keep", f1))
+
+				// Store second frontier with sp2@ts2 then delete it.
+				f2, err := span.MakeFrontier(sp2)
+				require.NoError(t, err)
+				_, err = f2.Forward(sp2, ts2)
+				require.NoError(t, err)
+				require.NoError(t, Store(ctx, txn, jobID, "delete", f2))
+				require.NoError(t, Delete(ctx, txn, jobID, "delete"))
+
+				// Store third frontier with sp3@ts1.
+				f3, err := span.MakeFrontier(sp3)
+				require.NoError(t, err)
+				_, err = f3.Forward(sp3, ts1)
+				require.NoError(t, err)
+				require.NoError(t, Store(ctx, txn, jobID, "also_keep", f3))
+			},
+			expected: []jobspb.ResolvedSpan{
+				{Span: sp1, Timestamp: ts1},
+				{Span: sp3, Timestamp: ts1},
+			},
+		},
+		{
+			name: "empty frontiers",
+			setup: func(t *testing.T, txn isql.Txn, jobID jobspb.JobID) {
+				// Store a frontier with no spans forwarded.
+				f, err := span.MakeFrontier(sp1, sp2, sp3)
+				require.NoError(t, err)
+				require.NoError(t, Store(ctx, txn, jobID, "empty", f))
+			},
+			expected: []jobspb.ResolvedSpan{
+				{Span: sp1},
+				{Span: sp2},
+				{Span: sp3},
+			},
+		},
+	}
+
+	for i, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			jobID := jobspb.JobID(3000 + i)
+
+			require.NoError(t, s.InternalDB().(isql.DB).Txn(ctx, func(ctx context.Context, txn isql.Txn) error {
+				tc.setup(t, txn, jobID)
+				return nil
+			}))
+
+			require.NoError(t, s.InternalDB().(isql.DB).Txn(ctx, func(ctx context.Context, txn isql.Txn) error {
+				spans, found, err := GetAllResolvedSpans(ctx, txn, jobID)
+				require.NoError(t, err)
+				if len(tc.expected) > 0 {
+					require.True(t, found)
+					require.ElementsMatch(t, tc.expected, spans)
+				} else {
+					require.False(t, found)
+					require.Empty(t, spans)
+				}
+				return nil
+			}))
+		})
+	}
+}