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

package meta

import (
	"bytes"
	"context"
	"crypto/sha256"
	"time"

	"git.frostfs.info/TrueCloudLab/frostfs-node/pkg/local_object_storage/internal/metaerr"
	"git.frostfs.info/TrueCloudLab/frostfs-observability/tracing"
	cid "git.frostfs.info/TrueCloudLab/frostfs-sdk-go/container/id"
	"go.etcd.io/bbolt"
	"go.opentelemetry.io/otel/attribute"
	"go.opentelemetry.io/otel/trace"
)

type ContainerStatPrm struct {
	ContainerID []cid.ID

	Limit                uint32
	StartFromContainerID *cid.ID
}

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

// ContainerStat returns object count and size for containers.
// If len(prm.ContainerID) > 0, then result slice contains records in the same order as prm.ContainerID.
// Otherwise result slice sorted by ContainerID.
func (db *DB) ContainerStat(ctx context.Context, prm ContainerStatPrm) ([]ContainerStat, error) {
	var (
		startedAt = time.Now()
		success   = false
	)
	defer func() {
		db.metrics.AddMethodDuration("ContainerStat", time.Since(startedAt), success)
	}()

	_, span := tracing.StartSpanFromContext(ctx, "metabase.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()

	db.modeMtx.RLock()
	defer db.modeMtx.RUnlock()

	if db.mode.NoMetabase() {
		return nil, ErrDegradedMode
	}

	if len(prm.ContainerID) > 0 {
		return db.containerStatByContainerID(prm.ContainerID)
	}
	return db.containerStatByLimit(prm.StartFromContainerID, prm.Limit)
}

func (db *DB) containerStatByContainerID(containerID []cid.ID) ([]ContainerStat, error) {
	var result []ContainerStat
	err := db.boltDB.View(func(tx *bbolt.Tx) error {
		for _, contID := range containerID {
			var stat ContainerStat
			stat.ContainerID = contID
			stat.SizeLogic = db.containerSize(tx, contID)

			counters, err := db.containerCounters(tx, contID)
			if err != nil {
				return err
			}
			stat.CountPhy = counters.Phy
			stat.CountLogic = counters.Logic
			stat.CountUser = counters.User
			result = append(result, stat)
		}
		return nil
	})
	if err != nil {
		return nil, metaerr.Wrap(err)
	}
	return result, nil
}

func (db *DB) containerStatByLimit(startFrom *cid.ID, limit uint32) ([]ContainerStat, error) {
	var result []ContainerStat

	var lastKey []byte
	if startFrom != nil {
		lastKey = make([]byte, sha256.Size)
		startFrom.Encode(lastKey)
	}

	var counts []containerIDObjectCounters
	var sizes []containerIDSize

	err := db.boltDB.View(func(tx *bbolt.Tx) error {
		var e error
		counts, e = getContainerCountersBatch(tx, lastKey, limit)
		if e != nil {
			return e
		}

		sizes, e = getContainerSizesBatch(tx, lastKey, limit)
		return e
	})
	if err != nil {
		return nil, metaerr.Wrap(err)
	}

	result = mergeSizeAndCounts(counts, sizes)
	if len(result) > int(limit) {
		result = result[:limit]
	}
	return result, nil
}

type containerIDObjectCounters struct {
	ContainerID cid.ID
	ObjectCounters
}

func getContainerCountersBatch(tx *bbolt.Tx, lastKey []byte, limit uint32) ([]containerIDObjectCounters, error) {
	var result []containerIDObjectCounters
	b := tx.Bucket(containerCounterBucketName)
	if b == nil {
		return result, nil
	}
	c := b.Cursor()
	var key, value []byte
	for key, value = c.Seek(lastKey); key != nil && uint32(len(result)) < limit; key, value = c.Next() {
		if bytes.Equal(lastKey, key) {
			continue
		}

		cnrID, err := parseContainerCounterKey(key)
		if err != nil {
			return nil, err
		}
		ent, err := parseContainerCounterValue(value)
		if err != nil {
			return nil, err
		}
		result = append(result, containerIDObjectCounters{
			ContainerID:    cnrID,
			ObjectCounters: ent,
		})
	}
	return result, nil
}

type containerIDSize struct {
	ContainerID cid.ID
	Size        uint64
}

func getContainerSizesBatch(tx *bbolt.Tx, lastKey []byte, limit uint32) ([]containerIDSize, error) {
	var result []containerIDSize
	b := tx.Bucket(containerVolumeBucketName)
	c := b.Cursor()
	var key, value []byte
	for key, value = c.Seek(lastKey); key != nil && uint32(len(result)) < limit; key, value = c.Next() {
		if bytes.Equal(lastKey, key) {
			continue
		}

		var r containerIDSize
		r.Size = parseContainerSize(value)
		if err := r.ContainerID.Decode(key); err != nil {
			return nil, err
		}
		result = append(result, r)
	}
	return result, nil
}

// mergeSizeAndCounts merges sizes and counts.
// As records are deleted in background, it can happen that metabase contains size record for container,
// but doesn't contain record for count.
func mergeSizeAndCounts(counts []containerIDObjectCounters, sizes []containerIDSize) []ContainerStat {
	var result []ContainerStat

	for len(counts) > 0 || len(sizes) > 0 {
		if len(counts) == 0 {
			result = append(result, ContainerStat{
				ContainerID: sizes[0].ContainerID,
				SizeLogic:   sizes[0].Size,
			})
			sizes = sizes[1:]
			continue
		}

		if len(sizes) == 0 {
			result = append(result, ContainerStat{
				ContainerID: counts[0].ContainerID,
				CountPhy:    counts[0].Phy,
				CountLogic:  counts[0].Logic,
				CountUser:   counts[0].User,
			})
			counts = counts[1:]
			continue
		}

		v := bytes.Compare(sizes[0].ContainerID[:], counts[0].ContainerID[:])

		if v == 0 { // equal
			result = append(result, ContainerStat{
				ContainerID: counts[0].ContainerID,
				CountPhy:    counts[0].Phy,
				CountLogic:  counts[0].Logic,
				CountUser:   counts[0].User,
				SizeLogic:   sizes[0].Size,
			})
			counts = counts[1:]
			sizes = sizes[1:]
		} else if v < 0 { // from sizes
			result = append(result, ContainerStat{
				ContainerID: sizes[0].ContainerID,
				SizeLogic:   sizes[0].Size,
			})
			sizes = sizes[1:]
		} else { // from counts
			result = append(result, ContainerStat{
				ContainerID: counts[0].ContainerID,
				CountPhy:    counts[0].Phy,
				CountLogic:  counts[0].Logic,
				CountUser:   counts[0].User,
			})
			counts = counts[1:]
		}
	}

	return result
}