frostfs-node/pkg/local_object_storage/engine/container_stat.go

package engine

import (
	"bytes"
	"context"
	"errors"
	"sort"
	"sync"

	"git.frostfs.info/TrueCloudLab/frostfs-node/pkg/local_object_storage/shard"
	"git.frostfs.info/TrueCloudLab/frostfs-observability/tracing"
	cid "git.frostfs.info/TrueCloudLab/frostfs-sdk-go/container/id"
	"go.opentelemetry.io/otel/attribute"
	"go.opentelemetry.io/otel/trace"
	"golang.org/x/sync/errgroup"
)

type ContainerStatPrm struct {
	ContainerID []cid.ID

	Limit                uint32
	StartFromContainerID *cid.ID
}

type ContainerStatRes struct {
	ContainerStats []ContainerStat
	Partial        bool
}

type ContainerStat struct {
	ContainerID                     cid.ID
	SizeLogic                       uint64
	CountPhy, CountLogic, CountUser uint64
}

var errInvalidLimit = errors.New("limit must be greater than zero")

func (e *StorageEngine) ContainerStat(ctx context.Context, prm ContainerStatPrm) (*ContainerStatRes, error) {
	if e.metrics != nil {
		defer elapsed("ContainerStat", e.metrics.AddMethodDuration)()
	}

	ctx, span := tracing.StartSpanFromContext(ctx, "StorageEngine.ContainerStat",
		trace.WithAttributes(
			attribute.Int("container_ids", len(prm.ContainerID)),
			attribute.Int64("limit", int64(prm.Limit)),
			attribute.Bool("start_from_container_id", prm.StartFromContainerID != nil),
		))
	defer span.End()

	if len(prm.ContainerID) == 0 && prm.Limit == 0 {
		return nil, errInvalidLimit
	}

	var result *ContainerStatRes
	err := e.execIfNotBlocked(func() error {
		var sErr error
		result, sErr = e.containerStat(ctx, prm)
		return sErr
	})
	return result, err
}

func (e *StorageEngine) containerStat(ctx context.Context, prm ContainerStatPrm) (*ContainerStatRes, error) {
	e.mtx.RLock()
	defer e.mtx.RUnlock()

	if len(prm.ContainerID) > 0 {
		sort.Slice(prm.ContainerID, func(i, j int) bool {
			return bytes.Compare(prm.ContainerID[i][:], prm.ContainerID[j][:]) < 0
		})
	}

	shardResults, partial, err := e.collectShardContainerStats(ctx, prm)
	if err != nil {
		return nil, err
	}

	return &ContainerStatRes{
		ContainerStats: e.mergeShardContainerStats(shardResults, int(prm.Limit)),
		Partial:        partial,
	}, nil
}

func (e *StorageEngine) collectShardContainerStats(ctx context.Context, prm ContainerStatPrm) ([][]shard.ContainerStat, bool, error) {
	if len(prm.ContainerID) > 0 {
		sort.Slice(prm.ContainerID, func(i, j int) bool {
			return bytes.Compare(prm.ContainerID[i][:], prm.ContainerID[j][:]) < 0
		})
	}

	var shardResults [][]shard.ContainerStat
	var shardErrors []error
	var resultsGuard sync.Mutex

	shPrm := shard.ContainerStatPrm{
		ContainerID:          prm.ContainerID,
		Limit:                prm.Limit,
		StartFromContainerID: prm.StartFromContainerID,
	}

	eg, egCtx := errgroup.WithContext(ctx)
	var shardsCount int
	e.iterateOverUnsortedShards(func(hs hashedShard) (stop bool) {
		shardsCount++
		eg.Go(func() error {
			s, err := hs.ContainerStat(egCtx, shPrm)
			resultsGuard.Lock()
			defer resultsGuard.Unlock()

			if err != nil {
				shardErrors = append(shardErrors, err)
				return nil
			}

			if len(s) > 0 {
				shardResults = append(shardResults, s)
			}

			return nil
		})
		return false
	})
	if err := eg.Wait(); err != nil {
		return nil, false, err
	}
	if shardsCount == len(shardErrors) {
		return nil, false, errors.Join(shardErrors...)
	}

	return shardResults, len(shardErrors) > 0, nil
}

func (e *StorageEngine) mergeShardContainerStats(shardResults [][]shard.ContainerStat, limit int) []ContainerStat {
	var stats []ContainerStat
	for len(stats) <= limit && len(shardResults) > 0 {
		// shard results are sorted by container ID
		sort.Slice(shardResults, func(i, j int) bool {
			return bytes.Compare(shardResults[i][0].ContainerID[:], shardResults[j][0].ContainerID[:]) < 0
		})

		if len(stats) > 0 && stats[len(stats)-1].ContainerID == shardResults[0][0].ContainerID {
			stats[len(stats)-1].SizeLogic += shardResults[0][0].SizeLogic
			stats[len(stats)-1].CountPhy += shardResults[0][0].CountPhy
			stats[len(stats)-1].CountLogic += shardResults[0][0].CountLogic
			stats[len(stats)-1].CountUser += shardResults[0][0].CountUser
		} else {
			stats = append(stats, ContainerStat{
				ContainerID: shardResults[0][0].ContainerID,
				SizeLogic:   shardResults[0][0].SizeLogic,
				CountPhy:    shardResults[0][0].CountPhy,
				CountLogic:  shardResults[0][0].CountLogic,
				CountUser:   shardResults[0][0].CountUser,
			})
		}

		if len(shardResults[0]) == 1 { // last item for shard
			shardResults = shardResults[1:]
		} else {
			shardResults[0] = shardResults[0][1:]
		}
	}
	if len(stats) > limit {
		stats = stats[:limit]
	}
	return stats
}