[usage] Implement CollectUsage

Author: easyCZ

components/usage/pkg/apiv1/usage.go

@@ -6,7 +6,11 @@ package apiv1
 
 import (
 	context "context"
+	"database/sql"
+	"fmt"
 	"github.com/gitpod-io/gitpod/common-go/log"
+	"github.com/gitpod-io/gitpod/usage/pkg/contentservice"
+	"github.com/gitpod-io/gitpod/usage/pkg/controller"
 	"time"
 
 	v1 "github.com/gitpod-io/gitpod/usage-api/v1"
@@ -21,12 +25,17 @@ var _ v1.UsageServiceServer = (*UsageService)(nil)
 
 type UsageService struct {
 	conn *gorm.DB
+
+	contentService contentservice.Interface
+
+	reportGenerator *ReportGenerator
+
 	v1.UnimplementedUsageServiceServer
 }
 
 const maxQuerySize = 31 * 24 * time.Hour
 
-func (us *UsageService) ListBilledUsage(ctx context.Context, in *v1.ListBilledUsageRequest) (*v1.ListBilledUsageResponse, error) {
+func (s *UsageService) ListBilledUsage(ctx context.Context, in *v1.ListBilledUsageRequest) (*v1.ListBilledUsageResponse, error) {
 	to := time.Now()
 	if in.To != nil {
 		to = in.To.AsTime()
@@ -52,7 +61,7 @@ func (us *UsageService) ListBilledUsage(ctx context.Context, in *v1.ListBilledUs
 		order = db.DescendingOrder
 	}
 
-	usageRecords, err := db.ListUsage(ctx, us.conn, db.AttributionID(in.GetAttributionId()), from, to, order)
+	usageRecords, err := db.ListUsage(ctx, s.conn, db.AttributionID(in.GetAttributionId()), from, to, order)
 	if err != nil {
 		log.Log.
 			WithField("attribution_id", in.AttributionId).
@@ -88,6 +97,212 @@ func (us *UsageService) ListBilledUsage(ctx context.Context, in *v1.ListBilledUs
 	}, nil
 }
 
-func NewUsageService(conn *gorm.DB) *UsageService {
-	return &UsageService{conn: conn}
+func (s *UsageService) ReconcileUsage(ctx context.Context, req *v1.ReconcileUsageRequest) (*v1.ReconcileUsageResponse, error) {
+	from := req.GetStartTime().AsTime()
+	to := req.GetEndTime().AsTime()
+
+	if to.Before(from) {
+		return nil, status.Errorf(codes.InvalidArgument, "End time must be after start time")
+	}
+
+	report, err := s.reportGenerator.GenerateUsageReport(ctx, from, to)
+	if err != nil {
+		log.Log.WithError(err).Error("Failed to reconcile time range.")
+		return nil, status.Error(codes.Internal, "failed to reconcile time range")
+	}
+
+	err = db.CreateUsageRecords(ctx, s.conn, report.UsageRecords)
+	if err != nil {
+		log.Log.WithError(err).Error("Failed to persist usage records.")
+		return nil, status.Error(codes.Internal, "failed to persist usage records")
+	}
+
+	filename := fmt.Sprintf("%s.gz", time.Now().Format(time.RFC3339))
+	err = s.contentService.UploadUsageReport(ctx, filename, report.UsageRecords)
+	if err != nil {
+		log.Log.WithError(err).Error("Failed to persist usage report to content service.")
+		return nil, status.Error(codes.Internal, "failed to persist usage report to content service")
+	}
+
+	var sessions []*v1.BilledSession
+	for _, instance := range report.UsageRecords {
+		sessions = append(sessions, usageRecordToBilledUsageProto(instance))
+	}
+
+	return &v1.ReconcileUsageResponse{
+		Sessions: sessions,
+	}, nil
+
+}
+
+func NewUsageService(conn *gorm.DB, reportGenerator *ReportGenerator, contentSvc contentservice.Interface) *UsageService {
+	return &UsageService{
+		conn:            conn,
+		reportGenerator: reportGenerator,
+		contentService:  contentSvc,
+	}
+}
+
+func usageRecordToBilledUsageProto(usageRecord db.WorkspaceInstanceUsage) *v1.BilledSession {
+	var endTime *timestamppb.Timestamp
+	if usageRecord.StoppedAt.Valid {
+		endTime = timestamppb.New(usageRecord.StoppedAt.Time)
+	}
+	return &v1.BilledSession{
+		AttributionId:  string(usageRecord.AttributionID),
+		UserId:         usageRecord.UserID.String(),
+		WorkspaceId:    usageRecord.WorkspaceID,
+		TeamId:         "",
+		WorkspaceType:  string(usageRecord.WorkspaceType),
+		ProjectId:      usageRecord.ProjectID,
+		InstanceId:     usageRecord.InstanceID.String(),
+		WorkspaceClass: usageRecord.WorkspaceClass,
+		StartTime:      timestamppb.New(usageRecord.StartedAt),
+		EndTime:        endTime,
+		Credits:        usageRecord.CreditsUsed,
+	}
+}
+
+type InvalidSession struct {
+	Reason  string
+	Session db.WorkspaceInstanceForUsage
+}
+
+type UsageReport struct {
+	GenerationTime time.Time
+
+	From time.Time
+	To   time.Time
+
+	RawSessions     []db.WorkspaceInstanceForUsage
+	InvalidSessions []InvalidSession
+
+	UsageRecords []db.WorkspaceInstanceUsage
+}
+
+func NewReportGenerator(conn *gorm.DB, pricer *controller.WorkspacePricer) *ReportGenerator {
+	return &ReportGenerator{
+		conn:    conn,
+		pricer:  pricer,
+		nowFunc: time.Now,
+	}
+}
+
+type ReportGenerator struct {
+	conn    *gorm.DB
+	pricer  *controller.WorkspacePricer
+	nowFunc func() time.Time
+}
+
+func (g *ReportGenerator) GenerateUsageReport(ctx context.Context, from, to time.Time) (UsageReport, error) {
+	now := g.nowFunc().UTC()
+	log.Infof("Gathering usage data from %s to %s", from, to)
+
+	report := UsageReport{
+		GenerationTime: now,
+		From:           from,
+		To:             to,
+	}
+
+	instances, err := db.ListWorkspaceInstancesInRange(ctx, g.conn, from, to)
+	if err != nil {
+		return report, fmt.Errorf("failed to list instances from db: %w", err)
+	}
+	report.RawSessions = instances
+
+	valid, invalid := validateInstances(instances)
+	report.InvalidSessions = invalid
+
+	if len(invalid) > 0 {
+		log.WithField("invalid_workspace_instances", invalid).Errorf("Detected %d invalid instances. These will be skipped in the current run.", len(invalid))
+	}
+	log.WithField("workspace_instances", instances).Debug("Successfully loaded workspace instances.")
+
+	trimmed := trimStartStopTime(valid, from, to)
+
+	report.UsageRecords = instancesToUsageRecords(trimmed, g.pricer, now)
+	return report, nil
+}
+
+func validateInstances(instances []db.WorkspaceInstanceForUsage) (valid []db.WorkspaceInstanceForUsage, invalid []InvalidSession) {
+	for _, i := range instances {
+		// i is a pointer to the current element, we need to assign it to ensure we're copying the value, not the current pointer.
+		instance := i
+
+		// Each instance must have a start time, without it, we do not have a baseline for usage computation.
+		if !instance.CreationTime.IsSet() {
+			invalid = append(invalid, InvalidSession{
+				Reason:  "missing creation time",
+				Session: instance,
+			})
+			continue
+		}
+
+		start := instance.CreationTime.Time()
+
+		// Currently running instances do not have a stopped time set, so we ignore these.
+		if instance.StoppedTime.IsSet() {
+			stop := instance.StoppedTime.Time()
+			if stop.Before(start) {
+				invalid = append(invalid, InvalidSession{
+					Reason:  "stop time is before start time",
+					Session: instance,
+				})
+				continue
+			}
+		}
+
+		valid = append(valid, instance)
+	}
+	return valid, invalid
+}
+
+// trimStartStopTime ensures that start time or stop time of an instance is never outside of specified start or stop time range.
+func trimStartStopTime(instances []db.WorkspaceInstanceForUsage, maximumStart, minimumStop time.Time) []db.WorkspaceInstanceForUsage {
+	var updated []db.WorkspaceInstanceForUsage
+
+	for _, instance := range instances {
+		if instance.CreationTime.Time().Before(maximumStart) {
+			instance.CreationTime = db.NewVarcharTime(maximumStart)
+		}
+
+		if instance.StoppedTime.Time().After(minimumStop) {
+			instance.StoppedTime = db.NewVarcharTime(minimumStop)
+		}
+
+		updated = append(updated, instance)
+	}
+	return updated
+}
+
+func instancesToUsageRecords(instances []db.WorkspaceInstanceForUsage, pricer *controller.WorkspacePricer, now time.Time) []db.WorkspaceInstanceUsage {
+	var usageRecords []db.WorkspaceInstanceUsage
+
+	for _, instance := range instances {
+		var stoppedAt sql.NullTime
+		if instance.StoppedTime.IsSet() {
+			stoppedAt = sql.NullTime{Time: instance.StoppedTime.Time(), Valid: true}
+		}
+
+		projectID := ""
+		if instance.ProjectID.Valid {
+			projectID = instance.ProjectID.String
+		}
+
+		usageRecords = append(usageRecords, db.WorkspaceInstanceUsage{
+			InstanceID:     instance.ID,
+			AttributionID:  instance.UsageAttributionID,
+			WorkspaceID:    instance.WorkspaceID,
+			ProjectID:      projectID,
+			UserID:         instance.OwnerID,
+			WorkspaceType:  instance.Type,
+			WorkspaceClass: instance.WorkspaceClass,
+			StartedAt:      instance.CreationTime.Time(),
+			StoppedAt:      stoppedAt,
+			CreditsUsed:    pricer.CreditsUsedByInstance(&instance, now),
+			GenerationID:   0,
+		})
+	}
+
+	return usageRecords
 }