frostfs-node/pkg/services/container/transport_splitter.go

package container

import (
	"context"
	"fmt"

	"git.frostfs.info/TrueCloudLab/frostfs-node/pkg/services/util"
	"git.frostfs.info/TrueCloudLab/frostfs-node/pkg/services/util/response"
	"git.frostfs.info/TrueCloudLab/frostfs-sdk-go/api/container"
)

type (
	TransportSplitter struct {
		next Server

		respSvc   *response.Service
		cnrAmount uint32
	}

	listStreamMsgSizeCtrl struct {
		util.ServerStream
		stream    ListStream
		respSvc   *response.Service
		cnrAmount uint32
	}
)

func NewSplitterService(cnrAmount uint32, respSvc *response.Service, next Server) Server {
	return &TransportSplitter{
		next:      next,
		respSvc:   respSvc,
		cnrAmount: cnrAmount,
	}
}

func (s *TransportSplitter) Put(ctx context.Context, req *container.PutRequest) (*container.PutResponse, error) {
	return s.next.Put(ctx, req)
}

func (s *TransportSplitter) Delete(ctx context.Context, req *container.DeleteRequest) (*container.DeleteResponse, error) {
	return s.next.Delete(ctx, req)
}

func (s *TransportSplitter) Get(ctx context.Context, req *container.GetRequest) (*container.GetResponse, error) {
	return s.next.Get(ctx, req)
}

func (s *TransportSplitter) List(ctx context.Context, req *container.ListRequest) (*container.ListResponse, error) {
	return s.next.List(ctx, req)
}

func (s *TransportSplitter) ListStream(req *container.ListStreamRequest, stream ListStream) error {
	return s.next.ListStream(req, &listStreamMsgSizeCtrl{
		ServerStream: stream,
		stream:       stream,
		respSvc:      s.respSvc,
		cnrAmount:    s.cnrAmount,
	})
}

func (s *listStreamMsgSizeCtrl) Send(resp *container.ListStreamResponse) error {
	s.respSvc.SetMeta(resp)
	body := resp.GetBody()
	ids := body.GetContainerIDs()

	var newResp *container.ListStreamResponse

	for {
		if newResp == nil {
			newResp = new(container.ListStreamResponse)
			newResp.SetBody(body)
		}

		cut := min(s.cnrAmount, uint32(len(ids)))

		body.SetContainerIDs(ids[:cut])
		newResp.SetMetaHeader(resp.GetMetaHeader())
		newResp.SetVerificationHeader(resp.GetVerificationHeader())

		if err := s.stream.Send(newResp); err != nil {
			return fmt.Errorf("TransportSplitter: %w", err)
		}

		ids = ids[cut:]

		if len(ids) == 0 {
			break
		}
	}

	return nil
}