[server] Integrate BillingService

geropl · geropl · commit bdbe61761b44 · 2022-08-08T08:27:54.000Z
diff --git a/components/server/ee/src/billing/entitlement-service-ubp.ts b/components/server/ee/src/billing/entitlement-service-ubp.ts
@@ -12,6 +12,7 @@ import {
     WORKSPACE_TIMEOUT_DEFAULT_LONG,
     WORKSPACE_TIMEOUT_DEFAULT_SHORT,
 } from "@gitpod/gitpod-protocol";
+import { AttributionId } from "@gitpod/gitpod-protocol/lib/attribution";
 import { inject, injectable } from "inversify";
 import {
     EntitlementService,
@@ -20,6 +21,7 @@ import {
 } from "../../../src/billing/entitlement-service";
 import { Config } from "../../../src/config";
 import { BillingModes } from "./billing-mode";
+import { BillingService } from "./billing-service";
 
 const MAX_PARALLEL_WORKSPACES_FREE = 4;
 const MAX_PARALLEL_WORKSPACES_PAID = 16;
@@ -32,6 +34,7 @@ export class EntitlementServiceUBP implements EntitlementService {
     @inject(Config) protected readonly config: Config;
     @inject(UserDB) protected readonly userDb: UserDB;
     @inject(BillingModes) protected readonly billingModes: BillingModes;
+    @inject(BillingService) protected readonly billingService: BillingService;
 
     async mayStartWorkspace(
         user: User,
@@ -50,19 +53,23 @@ export class EntitlementServiceUBP implements EntitlementService {
                 return undefined;
             }
         };
-        const [spendingLimitReached, hitParallelWorkspaceLimit] = await Promise.all([
-            this.checkSpendingLimitReached(user.id, date),
+        const [spendingLimitReachedOnCostCenter, hitParallelWorkspaceLimit] = await Promise.all([
+            this.checkSpendingLimitReached(user, date),
             hasHitParallelWorkspaceLimit(),
         ]);
 
         return {
-            spendingLimitReached,
+            spendingLimitReachedOnCostCenter,
             hitParallelWorkspaceLimit,
         };
     }
 
-    protected async checkSpendingLimitReached(userId: string, date: Date): Promise<boolean> {
-        return false;
+    protected async checkSpendingLimitReached(user: User, date: Date): Promise<AttributionId | undefined> {
+        const result = await this.billingService.checkSpendingLimitReached(user);
+        if (result.reached) {
+            return result.attributionId;
+        }
+        return undefined;
     }
 
     protected async getMaxParallelWorkspaces(user: User, date: Date): Promise<number> {
diff --git a/components/server/ee/src/workspace/gitpod-server-impl.ts b/components/server/ee/src/workspace/gitpod-server-impl.ts
@@ -46,7 +46,6 @@ import {
     FindPrebuildsParams,
     TeamMemberRole,
     WORKSPACE_TIMEOUT_DEFAULT_SHORT,
-    WorkspaceType,
     PrebuildEvent,
 } from "@gitpod/gitpod-protocol";
 import { ResponseError } from "vscode-jsonrpc";
@@ -72,7 +71,7 @@ import { BlockedRepository } from "@gitpod/gitpod-protocol/lib/blocked-repositor
 import { EligibilityService } from "../user/eligibility-service";
 import { AccountStatementProvider } from "../user/account-statement-provider";
 import { GithubUpgradeURL, PlanCoupon } from "@gitpod/gitpod-protocol/lib/payment-protocol";
-import { BillableSession, BillableSessionRequest, SortOrder } from "@gitpod/gitpod-protocol/lib/usage";
+import { BillableSession, BillableSessionRequest } from "@gitpod/gitpod-protocol/lib/usage";
 import {
     AssigneeIdentityIdentifier,
     TeamSubscription,
@@ -107,13 +106,13 @@ import { BitbucketAppSupport } from "../bitbucket/bitbucket-app-support";
 import { URL } from "url";
 import { UserCounter } from "../user/user-counter";
 import { AttributionId } from "@gitpod/gitpod-protocol/lib/attribution";
-import { CachingUsageServiceClientProvider } from "@gitpod/usage-api/lib/usage/v1/sugar";
-import * as usage from "@gitpod/usage-api/lib/usage/v1/usage_pb";
+import { CachingUsageServiceClientProvider, UsageService } from "@gitpod/usage-api/lib/usage/v1/sugar";
 import { Timestamp } from "google-protobuf/google/protobuf/timestamp_pb";
 import { EntitlementService } from "../../../src/billing/entitlement-service";
 import { BillingMode } from "@gitpod/gitpod-protocol/lib/billing-mode";
 import { BillingModes } from "../billing/billing-mode";
 import { getExperimentsClientForBackend } from "@gitpod/gitpod-protocol/lib/experiments/configcat-server";
+import { BillingService } from "../billing/billing-service";
 
 @injectable()
 export class GitpodServerEEImpl extends GitpodServerImpl {
@@ -160,6 +159,7 @@ export class GitpodServerEEImpl extends GitpodServerImpl {
     @inject(EntitlementService) protected readonly entitlementService: EntitlementService;
 
     @inject(BillingModes) protected readonly billingModes: BillingModes;
+    @inject(BillingService) protected readonly billingService: BillingService;
 
     initialize(
         client: GitpodClient | undefined,
@@ -256,39 +256,22 @@ export class GitpodServerEEImpl extends GitpodServerImpl {
     ): Promise<void> {
         await super.mayStartWorkspace(ctx, user, runningInstances);
 
-        // TODO(at) replace the naive implementation based on usage service
-        // with a proper call check against the upcoming invoice.
-        // For now this should just enable the work on fronend.
-        if (await this.isUsageBasedFeatureFlagEnabled(user)) {
-            // dummy implementation to test frontend bits
-            const attributionId = await this.userService.getWorkspaceUsageAttributionId(user);
-            const costCenter = !!attributionId && (await this.costCenterDB.findById(attributionId));
-            if (costCenter) {
-                const allSessions = await this.listBilledUsage(ctx, {
-                    attributionId,
-                    startedTimeOrder: SortOrder.Descending,
-                });
-                const totalUsage = allSessions.map((s) => s.credits).reduce((a, b) => a + b, 0);
-
-                if (totalUsage >= costCenter.spendingLimit) {
-                    throw new ResponseError(
-                        ErrorCodes.PAYMENT_SPENDING_LIMIT_REACHED,
-                        "Increase spending limit and try again.",
-                        {
-                            attributionId: user.usageAttributionId,
-                        },
-                    );
-                }
-            }
-        }
-
         const result = await this.entitlementService.mayStartWorkspace(user, new Date(), runningInstances);
         if (!result.enoughCredits) {
             throw new ResponseError(
                 ErrorCodes.NOT_ENOUGH_CREDIT,
                 `Not enough monthly workspace hours. Please upgrade your account to get more hours for your workspaces.`,
             );
         }
+        if (result.spendingLimitReachedOnCostCenter) {
+            throw new ResponseError(
+                ErrorCodes.PAYMENT_SPENDING_LIMIT_REACHED,
+                "Increase spending limit and try again.",
+                {
+                    attributionId: result.spendingLimitReachedOnCostCenter,
+                },
+            );
+        }
         if (!!result.hitParallelWorkspaceLimit) {
             throw new ResponseError(
                 ErrorCodes.TOO_MANY_RUNNING_WORKSPACES,
@@ -2146,24 +2129,17 @@ export class GitpodServerEEImpl extends GitpodServerImpl {
     async getNotifications(ctx: TraceContext): Promise<string[]> {
         const result = await super.getNotifications(ctx);
         const user = this.checkAndBlockUser("getNotifications");
-        if (user.usageAttributionId) {
-            // This change doesn't matter much because the listBilledUsage() call
-            // will be removed anyway in https://github.com/gitpod-io/gitpod/issues/11692
-            const request = {
-                attributionId: user.usageAttributionId,
-                startedTimeOrder: SortOrder.Descending,
-            };
-            const allSessions = await this.listBilledUsage(ctx, request);
-            const totalUsage = allSessions.map((s) => s.credits).reduce((a, b) => a + b, 0);
-            const costCenter = await this.costCenterDB.findById(user.usageAttributionId);
+
+        const billingMode = await this.billingModes.getBillingModeForUser(user, new Date());
+        if (billingMode.mode === "usage-based") {
+            const limit = await this.billingService.checkSpendingLimitReached(user);
+            const costCenter = await this.costCenterDB.findById(AttributionId.render(limit.attributionId));
             if (costCenter) {
-                if (totalUsage > costCenter.spendingLimit) {
+                if (limit.reached) {
                     result.unshift("The spending limit is reached.");
-                } else if (totalUsage > costCenter.spendingLimit * 0.8) {
+                } else if (limit.almostReached) {
                     result.unshift("The spending limit is almost reached.");
                 }
-            } else {
-                log.warn("No costcenter found.", { userId: user.id, attributionId: user.usageAttributionId });
             }
         }
         return result;
@@ -2192,7 +2168,7 @@ export class GitpodServerEEImpl extends GitpodServerImpl {
             timestampFrom,
             timestampTo,
         );
-        const sessions = response.getSessionsList().map((s) => this.mapBilledSession(s));
+        const sessions = response.getSessionsList().map((s) => UsageService.mapBilledSession(s));
 
         return sessions;
     }
@@ -2233,28 +2209,6 @@ export class GitpodServerEEImpl extends GitpodServerImpl {
         await this.guardAccess({ kind: "costCenter", /*subject: costCenter,*/ owner }, operation);
     }
 
-    protected mapBilledSession(s: usage.BilledSession): BillableSession {
-        function mandatory<T>(v: T, m: (v: T) => string = (s) => "" + s): string {
-            if (!v) {
-                throw new Error(`Empty value in usage.BilledSession for instanceId '${s.getInstanceId()}'`);
-            }
-            return m(v);
-        }
-        return {
-            attributionId: mandatory(s.getAttributionId()),
-            userId: s.getUserId() || undefined,
-            teamId: s.getTeamId() || undefined,
-            projectId: s.getProjectId() || undefined,
-            workspaceId: mandatory(s.getWorkspaceId()),
-            instanceId: mandatory(s.getInstanceId()),
-            workspaceType: mandatory(s.getWorkspaceType()) as WorkspaceType,
-            workspaceClass: s.getWorkspaceClass(),
-            startTime: mandatory(s.getStartTime(), (t) => t!.toDate().toISOString()),
-            endTime: s.getEndTime()?.toDate().toISOString(),
-            credits: s.getCredits(), // optional
-        };
-    }
-
     async getBillingModeForUser(ctx: TraceContextWithSpan): Promise<BillingMode> {
         traceAPIParams(ctx, {});
 
diff --git a/components/usage-api/typescript/src/usage/v1/sugar.ts b/components/usage-api/typescript/src/usage/v1/sugar.ts
@@ -8,11 +8,13 @@ import { UsageServiceClient } from "./usage_grpc_pb";
 import { TraceContext } from "@gitpod/gitpod-protocol/lib/util/tracing";
 import * as opentracing from "opentracing";
 import { Metadata } from "@grpc/grpc-js";
-import { ListBilledUsageRequest, ListBilledUsageResponse } from "./usage_pb";
+import { BilledSession, ListBilledUsageRequest, ListBilledUsageResponse } from "./usage_pb";
 import { injectable, inject, optional } from "inversify";
 import { createClientCallMetricsInterceptor, IClientCallMetrics } from "@gitpod/gitpod-protocol/lib/util/grpc";
 import * as grpc from "@grpc/grpc-js";
 import { Timestamp } from "google-protobuf/google/protobuf/timestamp_pb";
+import { BillableSession } from "@gitpod/gitpod-protocol/lib/usage";
+import { WorkspaceType } from "@gitpod/gitpod-protocol";
 
 export const UsageServiceClientProvider = Symbol("UsageServiceClientProvider");
 
@@ -91,7 +93,13 @@ export class PromisifiedUsageServiceClient {
         );
     }
 
-    public async listBilledUsage(_ctx: TraceContext, attributionId: string, order: ListBilledUsageRequest.Ordering, from?: Timestamp, to?: Timestamp): Promise<ListBilledUsageResponse> {
+    public async listBilledUsage(
+        _ctx: TraceContext,
+        attributionId: string,
+        order: ListBilledUsageRequest.Ordering,
+        from?: Timestamp,
+        to?: Timestamp,
+    ): Promise<ListBilledUsageResponse> {
         const ctx = TraceContext.childContext(`/usage-service/listBilledUsage`, _ctx);
 
         try {
@@ -133,3 +141,27 @@ export class PromisifiedUsageServiceClient {
         };
     }
 }
+
+export namespace UsageService {
+    export function mapBilledSession(s: BilledSession): BillableSession {
+        function mandatory<T>(v: T, m: (v: T) => string = (s) => "" + s): string {
+            if (!v) {
+                throw new Error(`Empty value in usage.BilledSession for instanceId '${s.getInstanceId()}'`);
+            }
+            return m(v);
+        }
+        return {
+            attributionId: mandatory(s.getAttributionId()),
+            userId: s.getUserId() || undefined,
+            teamId: s.getTeamId() || undefined,
+            projectId: s.getProjectId() || undefined,
+            workspaceId: mandatory(s.getWorkspaceId()),
+            instanceId: mandatory(s.getInstanceId()),
+            workspaceType: mandatory(s.getWorkspaceType()) as WorkspaceType,
+            workspaceClass: s.getWorkspaceClass(),
+            startTime: mandatory(s.getStartTime(), (t) => t!.toDate().toISOString()),
+            endTime: s.getEndTime()?.toDate().toISOString(),
+            credits: s.getCredits(), // optional
+        };
+    }
+}
