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)) }