Add feature flag to control the encoded token format

3 files changed, 95 insertions, 66 deletions

diff --git a/internal/gitaly/service/commit/tree_entries.go b/internal/gitaly/service/commit/tree_entries.go
index b3f8d25f2..a91c632d9 100644
--- a/internal/gitaly/service/commit/tree_entries.go
+++ b/internal/gitaly/service/commit/tree_entries.go
@@ -15,6 +15,7 @@ import (
 	"gitlab.com/gitlab-org/gitaly/v15/internal/git/localrepo"
 	"gitlab.com/gitlab-org/gitaly/v15/internal/git/lstree"
 	"gitlab.com/gitlab-org/gitaly/v15/internal/helper/chunk"
+	"gitlab.com/gitlab-org/gitaly/v15/internal/metadata/featureflag"
 	"gitlab.com/gitlab-org/gitaly/v15/proto/go/gitalypb"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/status"
@@ -178,7 +179,7 @@ func (s *server) sendTreeEntries(
 
 	cursor := ""
 	if p != nil {
-		entries, cursor, err = paginateTreeEntries(entries, p)
+		entries, cursor, err = paginateTreeEntries(ctx, entries, p)
 		if err != nil {
 			return err
 		}
@@ -302,7 +303,7 @@ func (s *server) GetTreeEntries(in *gitalypb.GetTreeEntriesRequest, stream gital
 	return s.sendTreeEntries(stream, repo, revision, path, in.Recursive, in.SkipFlatPaths, in.GetSort(), in.GetPaginationParams())
 }
 
-func paginateTreeEntries(entries []*gitalypb.TreeEntry, p *gitalypb.PaginationParameter) ([]*gitalypb.TreeEntry, string, error) {
+func paginateTreeEntries(ctx context.Context, entries []*gitalypb.TreeEntry, p *gitalypb.PaginationParameter) ([]*gitalypb.TreeEntry, string, error) {
 	limit := int(p.GetLimit())
 	start, tokenType := decodePageToken(p.GetPageToken())
 	index := -1
@@ -333,12 +334,16 @@ func paginateTreeEntries(entries []*gitalypb.TreeEntry, p *gitalypb.PaginationPa
 
 	paginated := entries[index : index+limit]
 
-	newPageToken, err := encodePageToken(paginated[len(paginated)-1])
-	if err != nil {
-		return nil, "", fmt.Errorf("encode page token: %w", err)
+	if featureflag.TreeEntriesNewPageTokenFormat.IsEnabled(ctx) {
+		newPageToken, err := encodePageToken(paginated[len(paginated)-1])
+		if err != nil {
+			return nil, "", fmt.Errorf("encode page token: %w", err)
+		}
+
+		return paginated, newPageToken, nil
 	}
 
-	return paginated, newPageToken, nil
+	return paginated, paginated[len(paginated)-1].GetOid(), nil
 }
 
 func buildEntryToken(entry *gitalypb.TreeEntry, tokenType pageTokenType) string {
@@ -364,7 +369,6 @@ const (
 	pageTokenTypeFilename
 )
 
-
 // decodePageToken decodes the given Base64-encoded page token. It returns the
 // continuation point of the token and its type.
 func decodePageToken(token string) (string, pageTokenType) {
diff --git a/internal/gitaly/service/commit/tree_entries_test.go b/internal/gitaly/service/commit/tree_entries_test.go
index 436836bcd..14dd4fe28 100644
--- a/internal/gitaly/service/commit/tree_entries_test.go
+++ b/internal/gitaly/service/commit/tree_entries_test.go
@@ -3,6 +3,7 @@
 package commit
 
 import (
+	"context"
 	"errors"
 	"io"
 	"strconv"
@@ -12,6 +13,7 @@ import (
 	"gitlab.com/gitlab-org/gitaly/v15/internal/git"
 	"gitlab.com/gitlab-org/gitaly/v15/internal/git/gittest"
 	"gitlab.com/gitlab-org/gitaly/v15/internal/helper"
+	"gitlab.com/gitlab-org/gitaly/v15/internal/metadata/featureflag"
 	"gitlab.com/gitlab-org/gitaly/v15/internal/testhelper"
 	"gitlab.com/gitlab-org/gitaly/v15/proto/go/gitalypb"
 	"google.golang.org/grpc/codes"
@@ -342,11 +344,11 @@ func TestGetTreeEntries_successful(t *testing.T) {
 		return strippedEntries
 	}
 
-	getPageToken := func(t testing.TB, entry *gitalypb.TreeEntry) string {
-		t.Helper()
+	getPageToken := func(tb testing.TB, entry *gitalypb.TreeEntry) string {
+		tb.Helper()
 
 		pageToken, err := encodePageToken(entry)
-		require.NoError(t, err)
+		require.NoError(tb, err)
 
 		return pageToken
 	}
@@ -362,6 +364,7 @@ func TestGetTreeEntries_successful(t *testing.T) {
 		pageToken       string
 		pageLimit       int32
 		cursor          string
+		legacyCursor    string
 		skipFlatPaths   bool
 	}{
 		{
@@ -437,13 +440,14 @@ func TestGetTreeEntries_successful(t *testing.T) {
 			sortBy:      gitalypb.GetTreeEntriesRequest_TREES_FIRST,
 		},
 		{
-			description: "with root path and sorted by trees with pagination",
-			revision:    []byte(commitID),
-			path:        []byte("."),
-			entries:     sortedAndPaginated,
-			pageLimit:   4,
-			sortBy:      gitalypb.GetTreeEntriesRequest_TREES_FIRST,
-			cursor:      "eyJmaWxlX25hbWUiOiIuRFNfU3RvcmUifQ==",
+			description:  "with root path and sorted by trees with pagination",
+			revision:     []byte(commitID),
+			path:         []byte("."),
+			entries:      sortedAndPaginated,
+			pageLimit:    4,
+			sortBy:       gitalypb.GetTreeEntriesRequest_TREES_FIRST,
+			cursor:       "eyJmaWxlX25hbWUiOiIuRFNfU3RvcmUifQ==",
+			legacyCursor: "fd90a3d2d21d6b4f9bec2c33fb7f49780c55f0d2",
 		},
 		{
 			description:     "with pagination parameters",
@@ -454,6 +458,7 @@ func TestGetTreeEntries_successful(t *testing.T) {
 			pageToken:       getPageToken(t, rootEntries[2]),
 			pageLimit:       3,
 			cursor:          "eyJmaWxlX25hbWUiOiJMSUNFTlNFIn0=",
+			legacyCursor:    "50b27c6518be44c42c4d87966ae2481ce895624c",
 		},
 		{
 			description:     "with pagination parameters larger than length",
@@ -482,61 +487,72 @@ func TestGetTreeEntries_successful(t *testing.T) {
 			pageLimit:       0,
 		},
 		{
-			description: "with a blank pagination token",
-			revision:    []byte(commitID),
-			path:        []byte("."),
-			pageToken:   "",
-			entries:     rootEntries[0:2],
-			pageLimit:   2,
-			cursor:      "eyJmaWxlX25hbWUiOiIuZ2l0aWdub3JlIn0=",
+			description:  "with a blank pagination token",
+			revision:     []byte(commitID),
+			path:         []byte("."),
+			pageToken:    "",
+			entries:      rootEntries[0:2],
+			pageLimit:    2,
+			cursor:       "eyJmaWxlX25hbWUiOiIuZ2l0aWdub3JlIn0=",
+			legacyCursor: "470ad2fcf1e33798f1afc5781d08e60c40f51e7a",
 		},
 	}
 
 	for _, testCase := range testCases {
-		t.Run(testCase.description, func(t *testing.T) {
-			for _, tc := range []struct {
-				desc      string
-				pageToken string
-			}{
-				{
-					desc:      "legacy token",
-					pageToken: testCase.legacyPageToken,
-				},
-				{
-					desc:      "structured page token",
-					pageToken: testCase.pageToken,
-				},
-			} {
-				t.Run(tc.desc, func(t *testing.T) {
-					request := &gitalypb.GetTreeEntriesRequest{
-						Repository:    repo,
-						Revision:      testCase.revision,
-						Path:          testCase.path,
-						Recursive:     testCase.recursive,
-						Sort:          testCase.sortBy,
-						SkipFlatPaths: testCase.skipFlatPaths,
-					}
-
-					if tc.pageToken != "" || testCase.pageLimit > 0 {
-						request.PaginationParams = &gitalypb.PaginationParameter{
-							PageToken: tc.pageToken,
-							Limit:     testCase.pageLimit,
-						}
-					}
-
-					c, err := client.GetTreeEntries(ctx, request)
-
-					require.NoError(t, err)
-					fetchedEntries, cursor := getTreeEntriesFromTreeEntryClient(t, c, nil)
-					testhelper.ProtoEqual(t, testCase.entries, fetchedEntries)
-
-					if testCase.pageLimit > 0 && len(testCase.entries) < len(rootEntries) {
-						require.NotNil(t, cursor)
-						require.Equal(t, testCase.cursor, cursor.NextCursor)
+		testhelper.NewFeatureSets(featureflag.TreeEntriesNewPageTokenFormat).Run(
+			t,
+			func(t *testing.T, ctx context.Context) {
+				t.Run(testCase.description, func(t *testing.T) {
+					for _, tc := range []struct {
+						desc      string
+						pageToken string
+					}{
+						{
+							desc:      "legacy token",
+							pageToken: testCase.legacyPageToken,
+						},
+						{
+							desc:      "structured page token",
+							pageToken: testCase.pageToken,
+						},
+					} {
+						t.Run(tc.desc, func(t *testing.T) {
+							request := &gitalypb.GetTreeEntriesRequest{
+								Repository:    repo,
+								Revision:      testCase.revision,
+								Path:          testCase.path,
+								Recursive:     testCase.recursive,
+								Sort:          testCase.sortBy,
+								SkipFlatPaths: testCase.skipFlatPaths,
+							}
+
+							if tc.pageToken != "" || testCase.pageLimit > 0 {
+								request.PaginationParams = &gitalypb.PaginationParameter{
+									PageToken: tc.pageToken,
+									Limit:     testCase.pageLimit,
+								}
+							}
+
+							c, err := client.GetTreeEntries(ctx, request)
+
+							require.NoError(t, err)
+							fetchedEntries, cursor := getTreeEntriesFromTreeEntryClient(t, c, nil)
+							testhelper.ProtoEqual(t, testCase.entries, fetchedEntries)
+
+							if testCase.pageLimit > 0 && len(testCase.entries) < len(rootEntries) {
+								require.NotNil(t, cursor)
+
+								if featureflag.TreeEntriesNewPageTokenFormat.IsEnabled(ctx) {
+									require.Equal(t, testCase.cursor, cursor.NextCursor)
+								} else {
+									require.Equal(t, testCase.legacyCursor, cursor.NextCursor)
+								}
+							}
+						})
 					}
 				})
-			}
-		})
+			},
+		)
 	}
 }
 
diff --git a/internal/metadata/featureflag/ff_tree_entries_new_page_token_format.go b/internal/metadata/featureflag/ff_tree_entries_new_page_token_format.go
new file mode 100644
index 000000000..c26194cfe
--- /dev/null
+++ b/internal/metadata/featureflag/ff_tree_entries_new_page_token_format.go
@@ -0,0 +1,9 @@
+package featureflag
+
+// TreeEntriesNewPageTokenFormat enables support for new page token format in TreeEntries RPC
+var TreeEntriesNewPageTokenFormat = NewFeatureFlag(
+	"tree_entries_new_page_token_format",
+	"v15.5.0",
+	"https://gitlab.com/gitlab-org/gitaly/-/issues/4542",
+	false,
+)