fix(unified-storage): use continue token containing both formats for dualwriter (#106525)

pkg/storage/unified/resource/server.go

@@ -198,6 +198,8 @@ type ResourceServerOptions struct {
 	storageMetrics *StorageMetrics
 
 	IndexMetrics *BleveIndexMetrics
+
+	MaxPageSizeBytes int
 }
 
 func NewResourceServer(opts ResourceServerOptions) (ResourceServer, error) {
@@ -222,6 +224,11 @@ func NewResourceServer(opts ResourceServerOptions) (ResourceServer, error) {
 		}
 	}
 
+	if opts.MaxPageSizeBytes <= 0 {
+		// By default, we use 2MB for the page size.
+		opts.MaxPageSizeBytes = 1024 * 1024 * 2
+	}
+
 	// Initialize the blob storage
 	blobstore := opts.Blob.Backend
 	if blobstore == nil {
@@ -250,19 +257,20 @@ func NewResourceServer(opts ResourceServerOptions) (ResourceServer, error) {
 	// Make this cancelable
 	ctx, cancel := context.WithCancel(context.Background())
 	s := &server{
-		tracer:         opts.Tracer,
-		log:            logger,
-		backend:        opts.Backend,
-		blob:           blobstore,
-		diagnostics:    opts.Diagnostics,
-		access:         opts.AccessClient,
-		writeHooks:     opts.WriteHooks,
-		lifecycle:      opts.Lifecycle,
-		now:            opts.Now,
-		ctx:            ctx,
-		cancel:         cancel,
-		storageMetrics: opts.storageMetrics,
-		indexMetrics:   opts.IndexMetrics,
+		tracer:           opts.Tracer,
+		log:              logger,
+		backend:          opts.Backend,
+		blob:             blobstore,
+		diagnostics:      opts.Diagnostics,
+		access:           opts.AccessClient,
+		writeHooks:       opts.WriteHooks,
+		lifecycle:        opts.Lifecycle,
+		now:              opts.Now,
+		ctx:              ctx,
+		cancel:           cancel,
+		storageMetrics:   opts.storageMetrics,
+		indexMetrics:     opts.IndexMetrics,
+		maxPageSizeBytes: opts.MaxPageSizeBytes,
 	}
 
 	if opts.Search.Resources != nil {
@@ -307,6 +315,8 @@ type server struct {
 	// init checking
 	once    sync.Once
 	initErr error
+
+	maxPageSizeBytes int
 }
 
 // Init implements ResourceServer.
@@ -791,7 +801,7 @@ func (s *server) List(ctx context.Context, req *resourcepb.ListRequest) (*resour
 	if req.Limit < 1 {
 		req.Limit = 50 // default max 50 items in a page
 	}
-	maxPageBytes := 1024 * 1024 * 2 // 2mb/page
+	maxPageBytes := s.maxPageSizeBytes
 	pageBytes := 0
 	rsp := &resourcepb.ListResponse{}
 

pkg/storage/unified/sql/server.go

@@ -43,6 +43,12 @@ func NewResourceServer(db infraDB.DB, cfg *setting.Cfg,
 		opts.Blob.URL = "file:///" + dir
 	}
 
+	// This is mostly for testing, being able to influence when we paginate
+	// based on the page size during tests.
+	unifiedStorageCfg := cfg.SectionWithEnvOverrides("unified_storage")
+	maxPageSizeBytes := unifiedStorageCfg.Key("max_page_size_bytes")
+	opts.MaxPageSizeBytes = maxPageSizeBytes.MustInt(0)
+
 	eDB, err := dbimpl.ProvideResourceDB(db, cfg, tracer)
 	if err != nil {
 		return nil, err