frostfs-sdk-go/client/container_list_stream.go

package client

import (
	"context"
	"errors"
	"fmt"
	"io"

	"git.frostfs.info/TrueCloudLab/frostfs-sdk-go/api/container"
	"git.frostfs.info/TrueCloudLab/frostfs-sdk-go/api/refs"
	rpcapi "git.frostfs.info/TrueCloudLab/frostfs-sdk-go/api/rpc"
	"git.frostfs.info/TrueCloudLab/frostfs-sdk-go/api/rpc/client"
	v2session "git.frostfs.info/TrueCloudLab/frostfs-sdk-go/api/session"
	"git.frostfs.info/TrueCloudLab/frostfs-sdk-go/api/signature"
	apistatus "git.frostfs.info/TrueCloudLab/frostfs-sdk-go/client/status"
	cid "git.frostfs.info/TrueCloudLab/frostfs-sdk-go/container/id"
	"git.frostfs.info/TrueCloudLab/frostfs-sdk-go/session"
	"git.frostfs.info/TrueCloudLab/frostfs-sdk-go/user"
)

// PrmContainerListStream groups parameters of ContainerListStream operation.
type PrmContainerListStream struct {
	XHeaders []string

	OwnerID user.ID

	Session *session.Container
}

func (x *PrmContainerListStream) buildRequest(c *Client) (*container.ListStreamRequest, error) {
	if x.OwnerID.IsEmpty() {
		return nil, errorAccountNotSet
	}

	var ownerV2 refs.OwnerID
	x.OwnerID.WriteToV2(&ownerV2)

	reqBody := new(container.ListStreamRequestBody)
	reqBody.SetOwnerID(&ownerV2)

	var meta v2session.RequestMetaHeader
	writeXHeadersToMeta(x.XHeaders, &meta)

	if x.Session != nil {
		var tokv2 v2session.Token
		x.Session.WriteToV2(&tokv2)

		meta.SetSessionToken(&tokv2)
	}

	var req container.ListStreamRequest
	req.SetBody(reqBody)
	c.prepareRequest(&req, &meta)
	return &req, nil
}

// ResContainerListStream groups the final result values of ContainerListStream operation.
type ResContainerListStream struct {
	statusRes
}

// ContainerListReader is designed to read list of container identifiers from FrostFS system.
//
// Must be initialized using Client.ContainerListInit, any other usage is unsafe.
type ContainerListReader struct {
	client          *Client
	cancelCtxStream context.CancelFunc
	err             error
	res             ResContainerListStream
	stream          interface {
		Read(resp *container.ListStreamResponse) error
	}
	tail []refs.ContainerID
}

// Read reads another list of the container identifiers. Works similar to
// io.Reader.Read but copies cid.ID and returns success flag instead of error.
//
// Failure reason can be received via Close.
//
// Panics if buf has zero length.
func (x *ContainerListReader) Read(buf []cid.ID) (int, bool) {
	// if len(buf) == 0 {
	// 	panic("empty buffer in ContainerListReader.ReadList")
	// }

	read := copyCnrIDBuffers(buf, x.tail)
	x.tail = x.tail[read:]

	if len(buf) == read {
		return read, true
	}

	for {
		var resp container.ListStreamResponse
		x.err = x.stream.Read(&resp)
		if x.err != nil {
			return read, false
		}

		x.res.st, x.err = x.client.processResponse(&resp)
		if x.err != nil || !apistatus.IsSuccessful(x.res.st) {
			return read, false
		}

		// read new chunk of containers
		ids := resp.GetBody().GetContainerIDs()
		if len(ids) == 0 {
			// just skip empty lists since they are not prohibited by protocol
			continue
		}

		ln := copyCnrIDBuffers(buf[read:], ids)
		read += ln

		if read == len(buf) {
			// save the tail
			x.tail = append(x.tail, ids[ln:]...)

			return read, true
		}
	}
}

func copyCnrIDBuffers(dst []cid.ID, src []refs.ContainerID) int {
	var i int
	for ; i < len(dst) && i < len(src); i++ {
		_ = dst[i].ReadFromV2(src[i])
	}
	return i
}

// Iterate iterates over the list of found container identifiers.
// f can return true to stop iteration earlier.
//
// Returns an error if container can't be read.
func (x *ContainerListReader) Iterate(f func(cid.ID) bool) error {
	buf := make([]cid.ID, 1)

	for {
		// Do not check first return value because `len(buf) == 1`,
		// so false means nothing was read.
		_, ok := x.Read(buf)
		if !ok {
			res, err := x.Close()
			if err != nil {
				return err
			}
			return apistatus.ErrFromStatus(res.Status())
		}
		if f(buf[0]) {
			return nil
		}
	}
}

// Close ends reading list of the matched containers and returns the result of the operation
// along with the final results. Must be called after using the ContainerListReader.
//
// Exactly one return value is non-nil. By default, server status is returned in res structure.
// Any client's internal or transport errors are returned as Go built-in error.
// If Client is tuned to resolve FrostFS API statuses, then FrostFS failures
// codes are returned as error.
func (x *ContainerListReader) Close() (*ResContainerListStream, error) {
	defer x.cancelCtxStream()

	if x.err != nil && !errors.Is(x.err, io.EOF) {
		return nil, x.err
	}

	return &x.res, nil
}

// ContainerListInit initiates container selection through a remote server using FrostFS API protocol.
//
// The call only opens the transmission channel, explicit fetching of identifiers of the account-owned
// containers is done using the ContainerListReader. Exactly one return value is non-nil.
// Resulting reader must be finally closed.
//
// Returns an error if parameters are set incorrectly (see PrmContainerListStream docs).
// Context is required and must not be nil. It is used for network communication.
func (c *Client) ContainerListInit(ctx context.Context, prm PrmContainerListStream) (*ContainerListReader, error) {
	req, err := prm.buildRequest(c)
	if err != nil {
		return nil, err
	}

	err = signature.SignServiceMessage(&c.prm.Key, req)
	if err != nil {
		return nil, fmt.Errorf("sign request: %w", err)
	}

	var r ContainerListReader
	ctx, r.cancelCtxStream = context.WithCancel(ctx)

	r.stream, err = rpcapi.ListContainersStream(&c.c, req, client.WithContext(ctx))
	if err != nil {
		return nil, fmt.Errorf("open stream: %w", err)
	}
	r.client = c

	return &r, nil
}