[ubp] reset usage on chargebee cancellation

Author: svenefftinge

components/ee/payment-endpoint/BUILD.yaml

@@ -10,6 +10,7 @@ packages:
     deps:
       - components/gitpod-db:lib
       - components/gitpod-protocol:lib
+      - components/usage-api/typescript:lib
     config:
       packaging: offline-mirror
       yarnLock: ${coreYarnLockBase}/../yarn.lock
@@ -24,6 +25,7 @@ packages:
     deps:
       - components/gitpod-db:lib
       - components/gitpod-protocol:lib
+      - components/usage-api/typescript:lib
       - :dbtest
     config:
       packaging: library
@@ -54,6 +56,7 @@ packages:
       - components/gitpod-db:dbtest-init
       - components/gitpod-db:lib
       - components/gitpod-protocol:lib
+      - components/usage-api/typescript:lib
     config:
       packaging: library
       yarnLock: ${coreYarnLockBase}/../yarn.lock

components/ee/payment-endpoint/package.json

@@ -7,6 +7,7 @@
   "dependencies": {
     "@gitpod/gitpod-db": "0.1.5",
     "@gitpod/gitpod-protocol": "0.1.5",
+    "@gitpod/usage-api": "0.1.5",
     "@octokit/rest": "18.5.6",
     "@octokit/webhooks": "9.17.0",
     "body-parser": "^1.19.2",

components/ee/payment-endpoint/src/accounting/team-subscription-service.ts

@@ -27,6 +27,7 @@ import { SubscriptionModel } from "./subscription-model";
 import { SubscriptionService } from "./subscription-service";
 import { AccountService } from "./account-service";
 import { log } from "@gitpod/gitpod-protocol/lib/util/logging";
+import { UbpResetOnCancel } from "../chargebee/ubp-reset-on-cancel";
 
 type TS = string | TeamSubscription;
 
@@ -36,6 +37,7 @@ export class TeamSubscriptionService {
     @inject(AccountingDB) protected readonly accountingDb: AccountingDB;
     @inject(AccountService) protected readonly accountService: AccountService;
     @inject(SubscriptionService) protected readonly subscriptionService: SubscriptionService;
+    @inject(UbpResetOnCancel) protected readonly ubpResetOnCancel: UbpResetOnCancel;
 
     /**
      * Adds new, free slots to the given TeamSubscription
@@ -178,6 +180,7 @@ export class TeamSubscriptionService {
             // Cancel assignee's subscription (if present)
             if (state === "assigned") {
                 const assignedSlot = slot as TeamSubscriptionSlotAssigned;
+                await this.ubpResetOnCancel.resetUsage(assignedSlot.assigneeId);
                 await this.cancelSubscription(db, assignedSlot.assigneeId, ts.planId, slot.id, now);
             }
 
@@ -453,6 +456,7 @@ export class TeamSubscriptionService {
             throw new Error(`Cannot find subscription for Team Subscription Slot ${slotId}!`);
         }
         model.cancel(subscription, cancellationDate, cancellationDate);
+        await this.ubpResetOnCancel.resetUsage(assigneeId);
         await this.subscriptionService.store(db, model);
     }
 

components/ee/payment-endpoint/src/accounting/team-subscription2-service.ts

@@ -9,6 +9,7 @@ import { AssignedTeamSubscription2, Subscription } from "@gitpod/gitpod-protocol
 import { Plans } from "@gitpod/gitpod-protocol/lib/plans";
 import { TeamSubscription2 } from "@gitpod/gitpod-protocol/lib/team-subscription-protocol";
 import { inject, injectable } from "inversify";
+import { UbpResetOnCancel } from "../chargebee/ubp-reset-on-cancel";
 import { SubscriptionModel } from "./subscription-model";
 import { SubscriptionService } from "./subscription-service";
 
@@ -17,6 +18,7 @@ export class TeamSubscription2Service {
     @inject(TeamDB) protected readonly teamDB: TeamDB;
     @inject(AccountingDB) protected readonly accountingDb: AccountingDB;
     @inject(SubscriptionService) protected readonly subscriptionService: SubscriptionService;
+    @inject(UbpResetOnCancel) protected readonly ubpResetOnCancel: UbpResetOnCancel;
 
     async addAllTeamMemberSubscriptions(ts2: TeamSubscription2): Promise<void> {
         const members = await this.teamDB.findMembersByTeam(ts2.teamId);
@@ -102,6 +104,7 @@ export class TeamSubscription2Service {
         teamMembershipId: string,
         cancellationDate: string,
     ) {
+        await this.ubpResetOnCancel.resetUsage(userId);
         const model = await this.loadSubscriptionModel(db, userId);
         const subscription = model.findSubscriptionByTeamMembershipId(teamMembershipId);
         if (!subscription) {

components/ee/payment-endpoint/src/chargebee/subscription-handler.ts

@@ -18,13 +18,15 @@ import { formatDate } from "@gitpod/gitpod-protocol/lib/util/date-time";
 import { getUpdatedAt } from "./chargebee-subscription-helper";
 import { UserPaidSubscription } from "@gitpod/gitpod-protocol/lib/accounting-protocol";
 import { DBSubscriptionAdditionalData } from "@gitpod/gitpod-db/lib/typeorm/entity/db-subscription";
+import { UbpResetOnCancel } from "./ubp-reset-on-cancel";
 
 @injectable()
 export class SubscriptionHandler implements EventHandler<chargebee.SubscriptionEventV2> {
     @inject(SubscriptionService) protected readonly subscriptionService: SubscriptionService;
     @inject(AccountingDB) protected readonly db: AccountingDB;
     @inject(SubscriptionMapperFactory) protected readonly mapperFactory: SubscriptionMapperFactory;
     @inject(UpgradeHelper) protected readonly upgradeHelper: UpgradeHelper;
+    @inject(UbpResetOnCancel) protected readonly ubpResetOnCancel: UbpResetOnCancel;
 
     canHandle(event: chargebee.Event<any>): boolean {
         if (event.event_type.startsWith("subscription")) {
@@ -55,6 +57,8 @@ export class SubscriptionHandler implements EventHandler<chargebee.SubscriptionE
 
             if (event.event_type === "subscription_changed") {
                 await this.checkAndChargeForUpgrade(userId, chargebeeSubscription);
+            } else if (event.event_type === "subscription_cancelled") {
+                await this.ubpResetOnCancel.resetUsage(userId);
             }
 
             await this.mapToGitpodSubscription(userId, event);

components/ee/payment-endpoint/src/chargebee/ubp-reset-on-cancel.ts

@@ -0,0 +1,38 @@
+/**
+ * Copyright (c) 2022 Gitpod GmbH. All rights reserved.
+ * Licensed under the GNU Affero General Public License (AGPL).
+ * See License.AGPL.txt in the project root for license information.
+ */
+
+import { AttributionId } from "@gitpod/gitpod-protocol/lib/attribution";
+import { log } from "@gitpod/gitpod-protocol/lib/util/logging";
+import { UsageServiceClient, UsageServiceDefinition } from "@gitpod/usage-api/lib/usage/v1/usage.pb";
+import { inject, injectable } from "inversify";
+
+injectable();
+export class UbpResetOnCancel {
+    @inject(UsageServiceDefinition.name) protected readonly usageService: UsageServiceClient;
+
+    async resetUsage(userId: string) {
+        try {
+            const attributionId = AttributionId.render({ kind: "user", userId: userId });
+            const balanceResponse = await this.usageService.getBalance({ attributionId });
+            const balance = balanceResponse.credits;
+            log.info({ userId }, "Chargbee subscription cancelled, adding credit note for remaining balance", {
+                balance,
+                attributionId,
+            });
+            if (balance > 0) {
+                await this.usageService.addUsageCreditNote({
+                    attributionId,
+                    credits: balance,
+                    description: "Resetting balance after chargebee subscription cancellation",
+                });
+            } else {
+                log.info({ userId }, "No balance to reset", { balance, attributionId });
+            }
+        } catch (err) {
+            log.error({ userId }, "Failed to reset usage balance", err);
+        }
+    }
+}

components/ee/payment-endpoint/src/config.ts

@@ -41,6 +41,8 @@ export GITPOD_GITHUB_APP_MKT_NAME=gitpod-draft-development-app
     readonly maxTeamSlotsOnCreation: number = !!process.env.TS_MAX_SLOTS_ON_CREATION
         ? parseInt(process.env.TS_MAX_SLOTS_ON_CREATION)
         : 1000;
+
+    readonly usageServiceAddr: string = process.env.GITPOD_USAGE_SERVICE_ADDR || "usage.default.svc.cluster.local:9001";
 }
 
 export interface ChargebeeWebhook {

components/ee/payment-endpoint/src/container-module.ts

@@ -6,23 +6,27 @@
 
 import { ContainerModule } from "inversify";
 
+import { log } from "@gitpod/gitpod-protocol/lib/util/logging";
+import { UsageServiceClient, UsageServiceDefinition } from "@gitpod/usage-api/lib/usage/v1/usage.pb";
+import { createChannel, createClient } from "nice-grpc";
+import { AccountService } from "./accounting/account-service";
+import { AccountServiceImpl } from "./accounting/account-service-impl";
+import { SubscriptionService } from "./accounting/subscription-service";
+import { TeamSubscriptionService } from "./accounting/team-subscription-service";
+import { TeamSubscription2Service } from "./accounting/team-subscription2-service";
+import { CompositeEventHandler, EventHandler } from "./chargebee/chargebee-event-handler";
 import { ChargebeeProvider, ChargebeeProviderOptions } from "./chargebee/chargebee-provider";
 import { EndpointController } from "./chargebee/endpoint-controller";
-import { SubscriptionService } from "./accounting/subscription-service";
-import { Config } from "./config";
-import { Server } from "./server";
 import { SubscriptionHandler } from "./chargebee/subscription-handler";
-import { SubscriptionMapperFactory, SubscriptionMapper } from "./chargebee/subscription-mapper";
+import { SubscriptionMapper, SubscriptionMapperFactory } from "./chargebee/subscription-mapper";
 import { TeamSubscriptionHandler } from "./chargebee/team-subscription-handler";
-import { CompositeEventHandler, EventHandler } from "./chargebee/chargebee-event-handler";
+import { UbpResetOnCancel } from "./chargebee/ubp-reset-on-cancel";
 import { UpgradeHelper } from "./chargebee/upgrade-helper";
-import { TeamSubscriptionService } from "./accounting/team-subscription-service";
-import { TeamSubscription2Service } from "./accounting/team-subscription2-service";
-import { AccountService } from "./accounting/account-service";
-import { AccountServiceImpl } from "./accounting/account-service-impl";
+import { Config } from "./config";
 import { GithubEndpointController } from "./github/endpoint-controller";
 import { GithubSubscriptionMapper } from "./github/subscription-mapper";
 import { GithubSubscriptionReconciler } from "./github/subscription-reconciler";
+import { Server } from "./server";
 
 export const productionContainerModule = new ContainerModule((bind, unbind, isBound, rebind) => {
     bind(Config).toSelf().inSingletonScope();
@@ -38,7 +42,6 @@ export const productionContainerModule = new ContainerModule((bind, unbind, isBo
         };
     });
     bind(TeamSubscriptionHandler).toSelf().inSingletonScope();
-
     bind(CompositeEventHandler).toSelf().inSingletonScope();
     bind(EventHandler).to(SubscriptionHandler).inSingletonScope();
     bind(EventHandler).to(TeamSubscriptionHandler).inSingletonScope();
@@ -60,4 +63,13 @@ export const productionContainerModule = new ContainerModule((bind, unbind, isBo
     bind(GithubEndpointController).toSelf().inSingletonScope();
     bind(GithubSubscriptionMapper).toSelf().inSingletonScope();
     bind(GithubSubscriptionReconciler).toSelf().inSingletonScope();
+
+    bind<UsageServiceClient>(UsageServiceDefinition.name)
+        .toDynamicValue((ctx) => {
+            const config = ctx.container.get<Config>(Config);
+            log.info("Connecting to usage service at", { addr: config.usageServiceAddr });
+            return createClient(UsageServiceDefinition, createChannel(config.usageServiceAddr));
+        })
+        .inSingletonScope();
+    bind(UbpResetOnCancel).toSelf().inSingletonScope();
 });

components/licensor/typescript/ee/src/api.ts

@@ -3,7 +3,6 @@
  * Licensed under the GNU Affero General Public License (AGPL).
  * See License.AGPL.txt in the project root for license information.
  */
-
 // generated using github.com/32leaves/bel
 // DO NOT MODIFY
 export enum Feature {
@@ -14,23 +13,23 @@ export enum Feature {
     FeatureWorkspaceSharing = "workspace-sharing",
 }
 export interface LicenseData {
-    type: LicenseType;
-    payload: LicensePayload;
-    plan: LicenseSubscriptionLevel;
-    fallbackAllowed: boolean;
+    type: LicenseType
+    payload: LicensePayload
+    plan: LicenseSubscriptionLevel
+    fallbackAllowed: boolean
 }
 
 export enum LicenseLevel {
     LevelTeam = 0,
     LevelEnterprise = 1,
 }
 export interface LicensePayload {
-    id: string;
-    domain: string;
-    level: LicenseLevel;
-    validUntil: string;
-    seats: number;
-    customerID?: string;
+    id: string
+    domain: string
+    level: LicenseLevel
+    validUntil: string
+    seats: number
+    customerID?: string
 }
 
 export enum LicenseSubscriptionLevel {

install/installer/pkg/common/constants.go

@@ -41,6 +41,7 @@ const (
 	SlowServerComponent         = "slow-server"
 	ServerInstallationAdminPort = 9000
 	SystemNodeCritical          = "system-node-critical"
+	PaymentEndpointComponent    = "payment-endpoint"
 	PublicApiComponent          = "public-api-server"
 	WSManagerComponent          = "ws-manager"
 	WSManagerBridgeComponent    = "ws-manager-bridge"

install/installer/pkg/components/usage/networkpolicy.go

@@ -50,6 +50,13 @@ func networkpolicy(ctx *common.RenderContext) ([]runtime.Object, error) {
 									},
 								},
 							},
+							{
+								PodSelector: &metav1.LabelSelector{
+									MatchLabels: map[string]string{
+										"component": common.PaymentEndpointComponent,
+									},
+								},
+							},
 							{
 								PodSelector: &metav1.LabelSelector{
 									MatchLabels: map[string]string{