jobfrontier: add GetResolvedSpans helper function

This patch adds the helper function `GetResolvedSpans`, which is like
`Get` except it returns the resolved spans directly so that callers
can decide what to do with them.

Release note: None

Author: andyyang890

pkg/jobs/jobfrontier/BUILD.bazel

@@ -8,9 +8,7 @@ go_library(
     deps = [
         "//pkg/jobs",
         "//pkg/jobs/jobspb",
-        "//pkg/roachpb",
         "//pkg/sql/isql",
-        "//pkg/util/hlc",
         "//pkg/util/protoutil",
         "//pkg/util/span",
         "@com_github_cockroachdb_errors//:errors",

pkg/jobs/jobfrontier/frontier.go

@@ -11,9 +11,7 @@ import (
 
 	"github.com/cockroachdb/cockroach/pkg/jobs"
 	"github.com/cockroachdb/cockroach/pkg/jobs/jobspb"
-	"github.com/cockroachdb/cockroach/pkg/roachpb"
 	"github.com/cockroachdb/cockroach/pkg/sql/isql"
-	"github.com/cockroachdb/cockroach/pkg/util/hlc"
 	"github.com/cockroachdb/cockroach/pkg/util/protoutil"
 	"github.com/cockroachdb/cockroach/pkg/util/span"
 	"github.com/cockroachdb/errors"
@@ -32,15 +30,15 @@ func Get(
 ) (span.Frontier, bool, error) {
 	infoStorage := jobs.InfoStorageForJob(txn, jobID)
 
-	// Read all persisted entries, both as entries and as plain spans; we need the
-	// latter form to construct the frontier and the former to advance it.
-	// TODO(dt): we could avoid duplicate allocation here if we added an API to
-	// construct frontier directly from entries.
-	var entries []frontierEntry
-	var spans []roachpb.Span
-
 	keyPrefix := frontierPrefix + name + shardSep
 
+	// Construct frontier to track the set of spans found and advance it to their
+	// persisted timestamps. This implies we persist zero-timestamp spans to keep
+	// the set of tracked spans even if they do not have progress.
+	frontier, err := span.MakeFrontier()
+	if err != nil {
+		return nil, false, err
+	}
 	var found bool
 	if err := infoStorage.Iterate(ctx, keyPrefix, func(_ string, value []byte) error {
 		found = true
@@ -49,28 +47,42 @@ func Get(
 			return err
 		}
 		for _, sp := range r.ResolvedSpans {
-			entries = append(entries, frontierEntry{Span: sp.Span, Timestamp: sp.Timestamp})
-			spans = append(spans, sp.Span)
+			if err := frontier.AddSpansAt(sp.Timestamp, sp.Span); err != nil {
+				return err
+			}
 		}
 		return nil
 	}); err != nil || !found {
 		return nil, false, err
 	}
 
-	// Construct frontier to track the set of spans found and advance it to their
-	// persisted timestamps. This implies we perist zero-timestamp spans to keep
-	// the set of tracked spans even if they do not have progress.
-	frontier, err := span.MakeFrontier(spans...)
-	if err != nil {
-		return nil, false, err
-	}
-	for _, entry := range entries {
-		if _, err := frontier.Forward(entry.Span, entry.Timestamp); err != nil {
-			return nil, false, err
+	return frontier, found, nil
+}
+
+// GetResolvedSpans is like Get except it returns the resolved spans directly
+// so that callers can decide what to do with them.
+func GetResolvedSpans(
+	ctx context.Context, txn isql.Txn, jobID jobspb.JobID, name string,
+) ([]jobspb.ResolvedSpan, bool, error) {
+	infoStorage := jobs.InfoStorageForJob(txn, jobID)
+
+	keyPrefix := frontierPrefix + name + shardSep
+
+	var found bool
+	var spans []jobspb.ResolvedSpan
+	if err := infoStorage.Iterate(ctx, keyPrefix, 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 frontier, true, nil
+	return spans, found, nil
 }
 
 // Store persists a frontier's current state to storage.
@@ -168,10 +180,3 @@ func deleteEntries(ctx context.Context, infoStorage jobs.InfoStorage, name strin
 	endKey := frontierPrefix + name + string(rune(shardSep[0])+1)
 	return infoStorage.DeleteRange(ctx, startKey, endKey, 0 /* no limit */)
 }
-
-// frontierEntry represents a single persisted frontier entry.
-// This is used internally for serialization but may be useful for testing.
-type frontierEntry struct {
-	Span      roachpb.Span
-	Timestamp hlc.Timestamp
-}

pkg/jobs/jobfrontier/frontier_test.go

@@ -54,7 +54,10 @@ func TestFrontier(t *testing.T) {
 
 	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))}
+		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)
 
@@ -77,6 +80,17 @@ func TestFrontier(t *testing.T) {
 				require.Equal(t, 3, countEntries(loadedFrontier))
 				checkFrontierContainsSpan(t, loadedFrontier, sp1, ts1)
 				checkFrontierContainsSpan(t, loadedFrontier, sp3, ts1)
+
+				spans, found, err := GetResolvedSpans(ctx, txn, jobID, "basic")
+				require.NoError(t, err)
+				require.True(t, found)
+				require.ElementsMatch(t,
+					[]jobspb.ResolvedSpan{
+						{Span: sp1, Timestamp: ts1},
+						{Span: sp2},
+						{Span: sp3, Timestamp: ts1},
+					},
+					spans)
 			},
 		},
 		{
@@ -86,6 +100,10 @@ func TestFrontier(t *testing.T) {
 				_, found, err := Get(ctx, txn, jobID, "nonexistent")
 				require.NoError(t, err)
 				require.False(t, found)
+
+				_, found, err = GetResolvedSpans(ctx, txn, jobID, "nonexistent")
+				require.NoError(t, err)
+				require.False(t, found)
 			},
 		},
 		{
@@ -106,6 +124,17 @@ func TestFrontier(t *testing.T) {
 				checkFrontierContainsSpan(t, loadedFrontier, sp1, ts1)
 				checkFrontierContainsSpan(t, loadedFrontier, sp2, ts2)
 				checkFrontierContainsSpan(t, loadedFrontier, sp3, ts1)
+
+				spans, found, err := GetResolvedSpans(ctx, txn, jobID, "multi")
+				require.NoError(t, err)
+				require.True(t, found)
+				require.ElementsMatch(t,
+					[]jobspb.ResolvedSpan{
+						{Span: sp1, Timestamp: ts1},
+						{Span: sp2, Timestamp: ts2},
+						{Span: sp3, Timestamp: ts1},
+					},
+					spans)
 			},
 		},
 		{
@@ -118,6 +147,10 @@ func TestFrontier(t *testing.T) {
 				_, found, err := Get(ctx, txn, jobID, "deleteme")
 				require.NoError(t, err)
 				require.False(t, found)
+
+				_, found, err = GetResolvedSpans(ctx, txn, jobID, "deleteme")
+				require.NoError(t, err)
+				require.False(t, found)
 			},
 		},
 		{
@@ -135,6 +168,15 @@ func TestFrontier(t *testing.T) {
 				require.True(t, found)
 				require.Equal(t, 1, countEntries(loadedFrontier))
 				checkFrontierContainsSpan(t, loadedFrontier, sp2, ts2)
+
+				spans, found, err := GetResolvedSpans(ctx, txn, jobID, "overwrite")
+				require.NoError(t, err)
+				require.True(t, found)
+				require.ElementsMatch(t,
+					[]jobspb.ResolvedSpan{
+						{Span: sp2, Timestamp: ts2},
+					},
+					spans)
 			},
 		},
 	}