ITWorx-Pulse-Public/internal/forecast/storage.go

package forecast

import (
	"context"
	"errors"
	"sort"
	"strings"
	"time"

	"github.com/itworx/pulse/internal/pool"
	"github.com/itworx/pulse/internal/share"
	"github.com/jackc/pgx/v5/pgxpool"
)

type HistoryReader interface {
	Points(context.Context, string, string, time.Time, int) ([]Point, error)
}

type PostgresHistory struct{ Pool *pgxpool.Pool }

func (r PostgresHistory) Points(ctx context.Context, kind, entityID string, since time.Time, limit int) ([]Point, error) {
	if r.Pool == nil {
		return nil, errors.New("capacity history store is unavailable")
	}
	if (kind != "share" && kind != "pool" && kind != "disk") || strings.TrimSpace(entityID) == "" || limit < 1 || limit > 512 {
		return nil, errors.New("capacity history query is invalid")
	}
	rows, err := r.Pool.Query(ctx, `SELECT observed_at,used_bytes FROM (
    SELECT DISTINCT ON (sampled_at) sampled_at,observed_at,used_bytes,source_id
    FROM capacity_samples
    WHERE entity_kind=$1 AND entity_id=$2 AND sampled_at >= $3
    ORDER BY sampled_at DESC,source_id ASC
    LIMIT $4
) history ORDER BY sampled_at ASC`, kind, entityID, since.UTC(), limit)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	points := make([]Point, 0, limit)
	for rows.Next() {
		var point Point
		if err := rows.Scan(&point.ObservedAt, &point.UsedBytes); err != nil {
			return nil, err
		}
		point.ObservedAt = point.ObservedAt.UTC()
		points = append(points, point)
	}
	if err := rows.Err(); err != nil {
		return nil, err
	}
	return points, nil
}

type StorageProvider struct {
	Shares  share.Provider
	Pools   pool.Provider
	History HistoryReader
	Policy  Policy
	Now     func() time.Time
}

func (p StorageProvider) Snapshot(ctx context.Context) (Snapshot, error) {
	if ctx == nil {
		return Snapshot{}, errors.New("forecast context is nil")
	}
	if err := ctx.Err(); err != nil {
		return Snapshot{}, err
	}
	now := time.Now().UTC()
	if p.Now != nil {
		now = p.Now().UTC()
	}
	policy := p.Policy.withDefaults()
	if err := policy.Validate(); err != nil {
		return Snapshot{}, err
	}
	view := PolicyView{Enabled: policy.Enabled, WindowSeconds: int64(policy.Window / time.Second), MinPoints: policy.MinPoints, Method: MethodLinearMedian}
	if !policy.Enabled {
		view.Method = MethodDisabled
	}
	if p.Shares == nil || p.Pools == nil {
		unknown := UnknownSnapshot(now, "source_unavailable")
		unknown.Policy = view
		return unknown, nil
	}
	shares, err := p.Shares.Snapshot(ctx)
	if err != nil {
		return Snapshot{}, err
	}
	pools, err := p.Pools.Snapshot(ctx)
	if err != nil {
		return Snapshot{}, err
	}
	if len(shares.Shares) == 0 {
		reason := shares.Source.Reason
		if reason == "" {
			reason = "no_capacity_entities"
		}
		return Snapshot{ContractVersion: ContractVersion, GeneratedAt: now, Policy: view, Items: []Forecast{}, Reason: reason}, nil
	}
	capacities := make(map[string]uint64, len(pools.Pools))
	for _, item := range pools.Pools {
		if item.UsableBytes > 0 {
			capacities[canonical(item.ID)] = item.UsableBytes
		}
	}
	items := make([]Forecast, 0, len(shares.Shares))
	qualified := 0
	for _, item := range shares.Shares {
		points := pointsFromShare(item)
		if p.History != nil {
			// Reserve one slot for a current observation that may not have reached
			// its six-hour history bucket yet.
			points, err = p.History.Points(ctx, "share", item.ID, now.Add(-policy.Window), policy.MaxPoints-1)
			if err != nil {
				return Snapshot{}, err
			}
			points = mergeCurrentPoint(points, item)
		}
		capacity := capacityForShare(item, capacities)
		forecast, predictErr := Predict(item.ID, item.Name, "share", capacity, points, now, policy)
		if predictErr != nil {
			return Snapshot{}, predictErr
		}
		if shares.Source.Freshness == share.Stale || item.SizeState == share.SizeStale {
			forecast.Method = MethodInsufficient
			forecast.Confidence = ConfidenceNone
			forecast.DaysToCapacity = nil
			forecast.ProjectedAt = nil
			forecast.Reason = "history_stale"
		} else if item.SizeState == share.SizeUnknown {
			forecast.Method = MethodInsufficient
			forecast.Confidence = ConfidenceNone
			forecast.DaysToCapacity = nil
			forecast.ProjectedAt = nil
			forecast.Reason = "history_unavailable"
		}
		if forecast.Confidence == ConfidenceHigh || forecast.Confidence == ConfidenceMedium {
			qualified++
		}
		items = append(items, forecast)
	}
	sort.SliceStable(items, func(i, j int) bool {
		if items[i].Name != items[j].Name {
			return items[i].Name < items[j].Name
		}
		return items[i].EntityID < items[j].EntityID
	})
	return Snapshot{ContractVersion: ContractVersion, GeneratedAt: now, Policy: view, Items: items, QualifiedCount: qualified}, nil
}

func mergeCurrentPoint(points []Point, item share.Share) []Point {
	if item.SizeObservedAt.IsZero() {
		return points
	}
	result := append([]Point(nil), points...)
	for index := range result {
		if result[index].ObservedAt.Equal(item.SizeObservedAt) {
			result[index].UsedBytes = item.UsedBytes
			return result
		}
	}
	return append(result, Point{ObservedAt: item.SizeObservedAt, UsedBytes: item.UsedBytes})
}

func pointsFromShare(item share.Share) []Point {
	points := make([]Point, 0, len(item.GrowthHistory)+1)
	for _, point := range item.GrowthHistory {
		points = append(points, Point{ObservedAt: point.ObservedAt, UsedBytes: point.UsedBytes})
	}
	if !item.SizeObservedAt.IsZero() {
		found := false
		for index := range points {
			if points[index].ObservedAt.Equal(item.SizeObservedAt) {
				points[index].UsedBytes = item.UsedBytes
				found = true
				break
			}
		}
		if !found {
			points = append(points, Point{ObservedAt: item.SizeObservedAt, UsedBytes: item.UsedBytes})
		}
	}
	return points
}

func capacityForShare(item share.Share, capacities map[string]uint64) uint64 {
	seen := make(map[string]struct{})
	var total uint64
	for _, placement := range item.Placements {
		id := canonical(placement.PoolID)
		capacity, exists := capacities[id]
		if !exists {
			continue
		}
		if _, duplicate := seen[id]; duplicate {
			continue
		}
		seen[id] = struct{}{}
		if ^uint64(0)-total < capacity {
			return ^uint64(0)
		}
		total += capacity
	}
	if total == 0 {
		total = capacities[canonical(item.StoragePolicy.PrimaryPool)]
	}
	return total
}

func canonical(value string) string { return strings.ToLower(strings.TrimSpace(value)) }