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

package engine

import (
	"context"
	"sync"

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

type RebuildPrm struct {
	ShardIDs          []*shard.ID
	ConcurrencyLimit  uint32
	TargetFillPercent uint32
}

type ShardRebuildResult struct {
	ShardID  *shard.ID
	Success  bool
	ErrorMsg string
}

type RebuildRes struct {
	ShardResults []ShardRebuildResult
}

func (e *StorageEngine) Rebuild(ctx context.Context, prm RebuildPrm) (RebuildRes, error) {
	ctx, span := tracing.StartSpanFromContext(ctx, "StorageEngine.Rebuild",
		trace.WithAttributes(
			attribute.Int("shard_id_count", len(prm.ShardIDs)),
			attribute.Int64("target_fill_percent", int64(prm.TargetFillPercent)),
			attribute.Int64("concurrency_limit", int64(prm.ConcurrencyLimit)),
		))
	defer span.End()

	res := RebuildRes{
		ShardResults: make([]ShardRebuildResult, 0, len(prm.ShardIDs)),
	}
	resGuard := &sync.Mutex{}

	limiter := shard.NewRebuildLimiter(prm.ConcurrencyLimit)

	eg, egCtx := errgroup.WithContext(ctx)
	for _, shardID := range prm.ShardIDs {
		eg.Go(func() error {
			e.mtx.RLock()
			sh, ok := e.shards[shardID.String()]
			e.mtx.RUnlock()

			if !ok {
				resGuard.Lock()
				defer resGuard.Unlock()
				res.ShardResults = append(res.ShardResults, ShardRebuildResult{
					ShardID:  shardID,
					ErrorMsg: errShardNotFound.Error(),
				})
				return nil
			}

			err := sh.ScheduleRebuild(egCtx, shard.RebuildPrm{
				ConcurrencyLimiter: limiter,
				TargetFillPercent:  prm.TargetFillPercent,
			})

			resGuard.Lock()
			defer resGuard.Unlock()

			if err != nil {
				res.ShardResults = append(res.ShardResults, ShardRebuildResult{
					ShardID:  shardID,
					ErrorMsg: err.Error(),
				})
			} else {
				res.ShardResults = append(res.ShardResults, ShardRebuildResult{
					ShardID: shardID,
					Success: true,
				})
			}
			return nil
		})
	}

	if err := eg.Wait(); err != nil {
		return RebuildRes{}, err
	}
	return res, nil
}