build source
This commit is contained in:
build
commit
ee1fec43ed
4171 changed files with 1351288 additions and 0 deletions
178
internal/store/usage.go
Normal file
178
internal/store/usage.go
Normal file
|
|
@ -0,0 +1,178 @@
|
|||
package store
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"time"
|
||||
|
||||
"github.com/iliaivanov/spec-kit-remote/cmd/dev-pod-api/internal/model"
|
||||
)
|
||||
|
||||
// RecordPodStart records a pod start event for usage tracking.
|
||||
func (s *Store) RecordPodStart(ctx context.Context, userID, podName string) error {
|
||||
_, err := s.db.Exec(ctx,
|
||||
`INSERT INTO usage_records (user_id, pod_name, event_type, value, recorded_at)
|
||||
VALUES ($1, $2, $3, 0, NOW())`,
|
||||
userID, podName, model.EventPodStart)
|
||||
if err != nil {
|
||||
return fmt.Errorf("record pod start: %w", err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// RecordPodStop records a pod stop event for usage tracking.
|
||||
func (s *Store) RecordPodStop(ctx context.Context, userID, podName string) error {
|
||||
_, err := s.db.Exec(ctx,
|
||||
`INSERT INTO usage_records (user_id, pod_name, event_type, value, recorded_at)
|
||||
VALUES ($1, $2, $3, 0, NOW())`,
|
||||
userID, podName, model.EventPodStop)
|
||||
if err != nil {
|
||||
return fmt.Errorf("record pod stop: %w", err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// RecordResourceSample records periodic CPU and memory usage samples.
|
||||
// cpuMillicores is CPU usage in millicores, memBytes is memory in bytes.
|
||||
func (s *Store) RecordResourceSample(ctx context.Context, userID, podName string, cpuMillicores, memBytes float64) error {
|
||||
now := time.Now().UTC()
|
||||
_, err := s.db.Exec(ctx,
|
||||
`INSERT INTO usage_records (user_id, pod_name, event_type, value, recorded_at)
|
||||
VALUES ($1, $2, $3, $4, $5), ($1, $2, $6, $7, $5)`,
|
||||
userID, podName, model.EventCPUSample, cpuMillicores, now,
|
||||
model.EventMemSample, memBytes)
|
||||
if err != nil {
|
||||
return fmt.Errorf("record resource sample: %w", err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// RecordAIRequest records an AI proxy request event.
|
||||
func (s *Store) RecordAIRequest(ctx context.Context, userID string) error {
|
||||
_, err := s.db.Exec(ctx,
|
||||
`INSERT INTO usage_records (user_id, pod_name, event_type, value, recorded_at)
|
||||
VALUES ($1, '', $2, 1, NOW())`,
|
||||
userID, model.EventAIRequest)
|
||||
if err != nil {
|
||||
return fmt.Errorf("record ai request: %w", err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// GetUsage returns aggregated usage for a user for a given month.
|
||||
func (s *Store) GetUsage(ctx context.Context, userID string, year int, month time.Month, monthlyBudget int) (*model.UsageSummary, error) {
|
||||
return s.getUsage(ctx, userID, year, month, monthlyBudget, time.Now().UTC())
|
||||
}
|
||||
|
||||
// getUsage is the internal implementation that accepts a "now" parameter for testing.
|
||||
func (s *Store) getUsage(ctx context.Context, userID string, year int, month time.Month, monthlyBudget int, now time.Time) (*model.UsageSummary, error) {
|
||||
monthStart := time.Date(year, month, 1, 0, 0, 0, 0, time.UTC)
|
||||
monthEnd := monthStart.AddDate(0, 1, 0)
|
||||
|
||||
// Calculate pod-hours from pod_start/pod_stop pairs.
|
||||
// For each pod_start, find the nearest subsequent pod_stop for the same user/pod.
|
||||
// If no stop found, use "now" (pod is still running).
|
||||
// Sessions that cross month boundaries are clamped to [monthStart, monthEnd].
|
||||
var podHours float64
|
||||
err := s.db.QueryRow(ctx, `
|
||||
WITH sessions AS (
|
||||
SELECT
|
||||
GREATEST(recorded_at, $2) AS start_time,
|
||||
(
|
||||
SELECT MIN(r2.recorded_at)
|
||||
FROM usage_records r2
|
||||
WHERE r2.user_id = r.user_id
|
||||
AND r2.pod_name = r.pod_name
|
||||
AND r2.event_type = 'pod_stop'
|
||||
AND r2.recorded_at > r.recorded_at
|
||||
) AS stop_time
|
||||
FROM usage_records r
|
||||
WHERE r.user_id = $1
|
||||
AND r.event_type = 'pod_start'
|
||||
AND r.recorded_at < $3
|
||||
)
|
||||
SELECT COALESCE(SUM(
|
||||
EXTRACT(EPOCH FROM LEAST(COALESCE(stop_time, $4), $3) - start_time) / 3600.0
|
||||
), 0)
|
||||
FROM sessions
|
||||
WHERE COALESCE(stop_time, $4) > $2`,
|
||||
userID, monthStart, monthEnd, now).Scan(&podHours)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("calculate pod hours: %w", err)
|
||||
}
|
||||
|
||||
// CPU-hours from resource samples.
|
||||
// Each sample represents 60s of CPU at the sampled rate (millicores).
|
||||
// CPU-hours = sum(millicores) * 60s / 3600s / 1000m = sum(millicores) / 60000
|
||||
var cpuHours float64
|
||||
err = s.db.QueryRow(ctx, `
|
||||
SELECT COALESCE(SUM(value) / 60000.0, 0)
|
||||
FROM usage_records
|
||||
WHERE user_id = $1 AND event_type = $2
|
||||
AND recorded_at >= $3 AND recorded_at < $4`,
|
||||
userID, model.EventCPUSample, monthStart, monthEnd).Scan(&cpuHours)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("calculate cpu hours: %w", err)
|
||||
}
|
||||
|
||||
// Count AI requests.
|
||||
var aiRequests int64
|
||||
err = s.db.QueryRow(ctx, `
|
||||
SELECT COUNT(*)
|
||||
FROM usage_records
|
||||
WHERE user_id = $1 AND event_type = $2
|
||||
AND recorded_at >= $3 AND recorded_at < $4`,
|
||||
userID, model.EventAIRequest, monthStart, monthEnd).Scan(&aiRequests)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("count ai requests: %w", err)
|
||||
}
|
||||
|
||||
var budgetUsedPct float64
|
||||
if monthlyBudget > 0 {
|
||||
budgetUsedPct = (podHours / float64(monthlyBudget)) * 100
|
||||
}
|
||||
|
||||
return &model.UsageSummary{
|
||||
PodHours: podHours,
|
||||
CPUHours: cpuHours,
|
||||
AIRequests: aiRequests,
|
||||
BudgetUsedPct: budgetUsedPct,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// GetDailyUsage returns daily usage breakdown for a user for a given month.
|
||||
// Pod-hours per day are estimated from CPU sample count (each sample = 1/60 hour of pod time).
|
||||
func (s *Store) GetDailyUsage(ctx context.Context, userID string, year int, month time.Month) ([]model.DailyUsage, error) {
|
||||
monthStart := time.Date(year, month, 1, 0, 0, 0, 0, time.UTC)
|
||||
monthEnd := monthStart.AddDate(0, 1, 0)
|
||||
|
||||
rows, err := s.db.Query(ctx, `
|
||||
SELECT
|
||||
DATE(recorded_at) AS day,
|
||||
COALESCE(SUM(CASE WHEN event_type = 'cpu_sample' THEN 1.0/60.0 ELSE 0 END), 0) AS pod_hours,
|
||||
COALESCE(SUM(CASE WHEN event_type = 'cpu_sample' THEN value / 60000.0 ELSE 0 END), 0) AS cpu_hours,
|
||||
COUNT(*) FILTER (WHERE event_type = 'ai_request') AS ai_requests
|
||||
FROM usage_records
|
||||
WHERE user_id = $1
|
||||
AND recorded_at >= $2 AND recorded_at < $3
|
||||
AND event_type IN ('cpu_sample', 'ai_request')
|
||||
GROUP BY DATE(recorded_at)
|
||||
ORDER BY DATE(recorded_at)`,
|
||||
userID, monthStart, monthEnd)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("query daily usage: %w", err)
|
||||
}
|
||||
defer rows.Close()
|
||||
|
||||
var result []model.DailyUsage
|
||||
for rows.Next() {
|
||||
var d model.DailyUsage
|
||||
var day time.Time
|
||||
if err := rows.Scan(&day, &d.PodHours, &d.CPUHours, &d.AIRequests); err != nil {
|
||||
return nil, fmt.Errorf("scan daily usage: %w", err)
|
||||
}
|
||||
d.Date = day.Format("2006-01-02")
|
||||
result = append(result, d)
|
||||
}
|
||||
return result, rows.Err()
|
||||
}
|
||||
Loading…
Add table
Add a link
Reference in a new issue