depager/v2/depager.go

/*
 * This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at https://mozilla.org/MPL/2.0/.
 */

package depager

import (
	"fmt"
)

/*
The `Page` interface must wrap server responses. This
allows pagers to calculate page sizes and iterate over
response aggregates.

If the underlying value of this interface is `nil` (e.g. a
nil pointer to a struct or a nil slice), `Elems()` will
panic.
*/
type Page[T any] interface {
	// Elems must return the items from the current page
	Elems() []T
}

// Exposes the part of the client that depager understands.
type Client[T any] interface {
	// NextPage returns the next page or it returns an error
	NextPage(offset uint64) (
		page Page[T],
		count uint64, // total count of all items being paged
		err error,
	)
}

type Pager[T any] interface {
	// Iter is intended to be used in a for-range loop
	Iter() <-chan T

	// LastErr must return the first error encountered, if any
	LastErr() error
}

func NewPager[T any](
	c Client[T],
	pageSize uint64,
) Pager[T] {
	return &pager[T]{
		client: c,
		n:      pageSize,
		p:      4,
	}
}

/*
Retrieve n items in the range [m*n, m*n + n - 1], inclusive.
Keep p pages buffered.
*/
type pager[T any] struct {
	client Client[T]
	m      uint64
	n      uint64
	err    error
	p      int
	cnt    uint64
}

func (p *pager[T]) iteratePages() <-chan Page[T] {
	ch := make(chan Page[T], p.p)
	go func() {
		defer close(ch)
		for {
			page, cnt, err :=
				p.client.NextPage(p.m * p.n)
			if err != nil {
				p.err = err
				return
			}
			if p.cnt == 0 {
				p.cnt = cnt
			}
			ch <- page

			if (p.m*p.n + p.n) >= p.cnt {
				return
			}
			p.m++
		}
	}()
	return ch
}

func (p *pager[T]) Iter() <-chan T {
	ch := make(chan T, p.n)
	go func() {
		defer close(ch)
		for page := range p.iteratePages() {
			for _, i := range page.Elems() {
				ch <- i
			}
			if p.err != nil {
				p.err = fmt.Errorf("pager: iterate items: %s", p.err)
				return
			}
		}
	}()
	return ch
}

func (p *pager[T]) LastErr() error {
	return p.err
}