[usage] Implement CollectUsage

Author: easyCZ

components/usage/pkg/apiv1/usage.go

@@ -7,6 +7,7 @@ package apiv1
 import (
 	context "context"
 	"github.com/gitpod-io/gitpod/common-go/log"
+	"github.com/gitpod-io/gitpod/usage/pkg/controller"
 	"time"
 
 	v1 "github.com/gitpod-io/gitpod/usage-api/v1"
@@ -21,12 +22,15 @@ var _ v1.UsageServiceServer = (*UsageService)(nil)
 
 type UsageService struct {
 	conn *gorm.DB
+
+	usageCtrl *controller.UsageReconciler
+
 	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 +56,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 +92,54 @@ func (us *UsageService) ListBilledUsage(ctx context.Context, in *v1.ListBilledUs
 	}, nil
 }
 
-func NewUsageService(conn *gorm.DB) *UsageService {
-	return &UsageService{conn: conn}
+func (s *UsageService) CollectUsage(ctx context.Context, req *v1.CollectUsageRequest) (*v1.CollectUsageResponse, error) {
+	from := req.GetStartTime().AsTime()
+	to := req.GetEndTime().AsTime()
+
+	if from.Before(to) {
+		return nil, status.Errorf(codes.InvalidArgument, "End time must be after start time")
+	}
+
+	_, report, err := s.usageCtrl.ReconcileTimeRange(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")
+	}
+
+	var sessions []*v1.BilledSession
+	for _, instance := range report {
+		sessions = append(sessions, usageRecordToBilledUsageProto(instance))
+	}
+
+	return &v1.CollectUsageResponse{
+		Sessions: sessions,
+	}, nil
+
+}
+
+func NewUsageService(conn *gorm.DB, usageCtrl *controller.UsageReconciler) *UsageService {
+	return &UsageService{
+		conn:      conn,
+		usageCtrl: usageCtrl,
+	}
+}
+
+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,
+	}
 }

components/usage/pkg/apiv1/usage_test.go

@@ -142,7 +142,7 @@ func TestUsageService_ListBilledUsage(t *testing.T) {
 				baseserver.WithGRPC(baseserver.MustUseRandomLocalAddress(t)),
 			)
 
-			v1.RegisterUsageServiceServer(srv.GRPC(), NewUsageService(dbconn))
+			v1.RegisterUsageServiceServer(srv.GRPC(), NewUsageService(dbconn, nil))
 			baseserver.StartServerForTests(t, srv)
 
 			conn, err := grpc.Dial(srv.GRPCAddress(), grpc.WithTransportCredentials(insecure.NewCredentials()))

components/usage/pkg/server/server.go

@@ -93,7 +93,8 @@ func Start(cfg Config) error {
 		contentService = contentservice.New(cfg.ContentServiceAddress)
 	}
 
-	ctrl, err := controller.New(schedule, controller.NewUsageReconciler(conn, pricer, v1.NewBillingServiceClient(selfConnection), contentService))
+	usageCtrl := controller.NewUsageReconciler(conn, pricer, v1.NewBillingServiceClient(selfConnection), contentService)
+	ctrl, err := controller.New(schedule, usageCtrl)
 	if err != nil {
 		return fmt.Errorf("failed to initialize usage controller: %w", err)
 	}
@@ -104,7 +105,7 @@ func Start(cfg Config) error {
 	}
 	defer ctrl.Stop()
 
-	err = registerGRPCServices(srv, conn, stripeClient)
+	err = registerGRPCServices(srv, conn, stripeClient, usageCtrl)
 	if err != nil {
 		return fmt.Errorf("failed to register gRPC services: %w", err)
 	}
@@ -122,8 +123,8 @@ func Start(cfg Config) error {
 	return nil
 }
 
-func registerGRPCServices(srv *baseserver.Server, conn *gorm.DB, stripeClient *stripe.Client) error {
-	v1.RegisterUsageServiceServer(srv.GRPC(), apiv1.NewUsageService(conn))
+func registerGRPCServices(srv *baseserver.Server, conn *gorm.DB, stripeClient *stripe.Client, usageCtrl *controller.UsageReconciler) error {
+	v1.RegisterUsageServiceServer(srv.GRPC(), apiv1.NewUsageService(conn, usageCtrl))
 	if stripeClient == nil {
 		v1.RegisterBillingServiceServer(srv.GRPC(), &apiv1.BillingServiceNoop{})
 	} else {