feat: add cronjob to check workspace limit which will pause workspace

pnpm-lock.yaml

@@ -543,6 +543,9 @@ importers:
       nodemailer:
         specifier: ^6.9.8
         version: 6.9.8
+      p-map:
+        specifier: 4.0.0
+        version: 4.0.0
       passport:
         specifier: ^0.7.0
         version: 0.7.0
@@ -661,9 +664,6 @@ importers:
       execa:
         specifier: ^5.1.1
         version: 5.1.1
-      p-map:
-        specifier: 4.0.0
-        version: 4.0.0
       prisma:
         specifier: 5.14.0
         version: 5.14.0
@@ -2341,10 +2341,6 @@ packages:
     resolution: {integrity: sha512-gM/FdNsK3BlrD6JRrhmiyqBXQsCpzSUdKSoZwJMQfXqfqcK321og+uMssc6HYcygUMrGvPnNJyJ1RqZPFDrgtg==}
     engines: {node: '>=20'}
 
-  '@ljharb/resumer@0.0.1':
-    resolution: {integrity: sha512-skQiAOrCfO7vRTq53cxznMpks7wS1va95UCidALlOVWqvBAzwPVErwizDwoMqNVMEn1mDq0utxZd02eIrvF1lw==}
-    engines: {node: '>= 0.4'}
-
   '@ljharb/through@2.3.11':
     resolution: {integrity: sha512-ccfcIDlogiXNq5KcbAwbaO7lMh3Tm1i3khMPYpxlK8hH/W53zN81KM9coerRLOnTGu3nfXIniAmQbRI9OxbC0w==}
     engines: {node: '>= 0.4'}
@@ -15298,10 +15294,6 @@ snapshots:
 
   '@lemonsqueezy/lemonsqueezy.js@3.3.1': {}
 
-  '@ljharb/resumer@0.0.1':
-    dependencies:
-      '@ljharb/through': 2.3.11
-
   '@ljharb/through@2.3.11':
     dependencies:
       call-bind: 1.0.7

src/server/cronjob/index.ts

@@ -9,6 +9,7 @@ import { token } from '../model/notification/token/index.js';
 import pMap from 'p-map';
 import { sendFeedEventsNotify } from '../model/feed/event.js';
 import { get } from 'lodash-es';
+import { checkWorkspaceUsage } from '../model/billing/cronjob.js';
 
 type WebsiteEventCountSqlReturn = {
   workspace_id: string;
@@ -29,6 +30,10 @@ export function initCronjob() {
         checkFeedEventsNotify(FeedChannelNotifyFrequency.day),
       ]);
 
+      if (env.billing.enable) {
+        await checkWorkspaceUsage();
+      }
+
       logger.info('Daily cronjob completed');
     } catch (err) {
       logger.error('Daily cronjob error:', err);
@@ -386,6 +391,9 @@ async function dailyHTTPCertCheckNotify() {
   );
 }
 
+/**
+ * Check feed events notify
+ */
 async function checkFeedEventsNotify(
   notifyFrequency: FeedChannelNotifyFrequency
 ) {

src/server/model/billing/cronjob.ts

@@ -0,0 +1,65 @@
+import pMap from 'p-map';
+import { prisma } from '../_client.js';
+import { WorkspaceSubscriptionTier } from '@prisma/client';
+import { logger } from '../../utils/logger.js';
+import { getTierLimit } from './limit.js';
+import { getWorkspaceUsage, pauseWorkspace } from './workspace.js';
+import dayjs from 'dayjs';
+import { getWorkspaceServiceCount } from '../workspace.js';
+
+/**
+ * Check workspace usage
+ * if over limit, pause workspace
+ */
+export async function checkWorkspaceUsage() {
+  logger.info('[checkWorkspaceUsage] Start run checkWorkspaceUsage');
+
+  const workspaces = await prisma.workspace.findMany({
+    where: {
+      paused: false,
+    },
+    include: {
+      subscription: true,
+    },
+  });
+
+  await pMap(
+    workspaces,
+    async (workspace) => {
+      const tier =
+        workspace.subscription?.tier ?? WorkspaceSubscriptionTier.FREE;
+
+      if (tier === WorkspaceSubscriptionTier.UNLIMITED) {
+        return;
+      }
+
+      const [usage, serviceCount] = await Promise.all([
+        getWorkspaceUsage(
+          workspace.id,
+          dayjs().startOf('month').valueOf(),
+          dayjs().valueOf()
+        ),
+        getWorkspaceServiceCount(workspace.id),
+      ]);
+
+      const limit = getTierLimit(tier);
+
+      const overUsage =
+        serviceCount.website > limit.maxWebsiteCount ||
+        usage.websiteEventCount > limit.maxWebsiteEventCount ||
+        usage.monitorExecutionCount > limit.maxMonitorExecutionCount ||
+        usage.websiteEventCount > limit.maxWebsiteEventCount ||
+        usage.surveyCount > limit.maxSurveyCount ||
+        serviceCount.feed > limit.maxFeedChannelCount ||
+        usage.feedEventCount > limit.maxFeedEventCount;
+
+      if (overUsage) {
+        // pause workspace
+        await pauseWorkspace(workspace.id);
+      }
+    },
+    {
+      concurrency: 5,
+    }
+  );
+}

src/server/model/billing/limit.ts

@@ -1,4 +1,4 @@
-import { TierType } from './types.js';
+import { WorkspaceSubscriptionTier } from '@prisma/client';
 
 interface TierLimit {
   maxWebsiteCount: number;
@@ -12,8 +12,8 @@ interface TierLimit {
 /**
  * Limit, Every month
  */
-export function getTierLimit(tier: TierType): TierLimit {
-  if (tier === 'free') {
+export function getTierLimit(tier: WorkspaceSubscriptionTier): TierLimit {
+  if (tier === WorkspaceSubscriptionTier.FREE) {
     return {
       maxWebsiteCount: 3,
       maxWebsiteEventCount: 100_000,
@@ -24,7 +24,7 @@ export function getTierLimit(tier: TierType): TierLimit {
     };
   }
 
-  if (tier === 'pro') {
+  if (tier === WorkspaceSubscriptionTier.PRO) {
     return {
       maxWebsiteCount: 10,
       maxWebsiteEventCount: 1_000_000,
@@ -35,7 +35,7 @@ export function getTierLimit(tier: TierType): TierLimit {
     };
   }
 
-  if (tier === 'team') {
+  if (tier === WorkspaceSubscriptionTier.TEAM) {
     return {
       maxWebsiteCount: -1,
       maxWebsiteEventCount: 20_000_000,
@@ -46,6 +46,7 @@ export function getTierLimit(tier: TierType): TierLimit {
     };
   }
 
+  // Unlimited
   return {
     maxWebsiteCount: -1,
     maxWebsiteEventCount: -1,

src/server/model/billing/types.ts

@@ -1 +0,0 @@
-export type TierType = 'free' | 'pro' | 'team' | 'unlimited';

src/server/model/billing/workspace.ts

@@ -0,0 +1,56 @@
+import { WorkspaceSubscriptionTier } from '@prisma/client';
+import { prisma } from '../_client.js';
+
+export async function getWorkspaceUsage(
+  workspaceId: string,
+  startAt: number,
+  endAt: number
+) {
+  const res = await prisma.workspaceDailyUsage.aggregate({
+    where: {
+      workspaceId,
+      date: {
+        gte: new Date(startAt),
+        lte: new Date(endAt),
+      },
+    },
+    _sum: {
+      websiteAcceptedCount: true,
+      websiteEventCount: true,
+      monitorExecutionCount: true,
+      surveyCount: true,
+      feedEventCount: true,
+    },
+  });
+
+  return {
+    websiteAcceptedCount: res._sum.websiteAcceptedCount ?? 0,
+    websiteEventCount: res._sum.websiteEventCount ?? 0,
+    monitorExecutionCount: res._sum.monitorExecutionCount ?? 0,
+    surveyCount: res._sum.surveyCount ?? 0,
+    feedEventCount: res._sum.feedEventCount ?? 0,
+  };
+}
+
+export async function getWorkspaceSubscription(
+  workspaceId: string
+): Promise<WorkspaceSubscriptionTier> {
+  const subscription = await prisma.workspaceSubscription.findFirst({
+    where: {
+      workspaceId,
+    },
+  });
+
+  return subscription?.tier ?? WorkspaceSubscriptionTier.FREE;
+}
+
+export async function pauseWorkspace(workspaceId: string) {
+  await prisma.workspace.update({
+    where: {
+      id: workspaceId,
+    },
+    data: {
+      paused: true,
+    },
+  });
+}

src/server/model/workspace.ts

@@ -72,3 +72,47 @@ export async function getWorkspaceWebsiteDateRange(websiteId: string) {
     min: res._min.createdAt,
   };
 }
+
+export async function getWorkspaceServiceCount(workspaceId: string) {
+  const [website, monitor, telemetry, page, survey, feed] = await Promise.all([
+    prisma.website.count({
+      where: {
+        workspaceId,
+      },
+    }),
+    prisma.monitor.count({
+      where: {
+        workspaceId,
+      },
+    }),
+    prisma.telemetry.count({
+      where: {
+        workspaceId,
+      },
+    }),
+    prisma.monitorStatusPage.count({
+      where: {
+        workspaceId,
+      },
+    }),
+    prisma.survey.count({
+      where: {
+        workspaceId,
+      },
+    }),
+    prisma.feedChannel.count({
+      where: {
+        workspaceId,
+      },
+    }),
+  ]);
+
+  return {
+    website,
+    monitor,
+    telemetry,
+    page,
+    survey,
+    feed,
+  };
+}

src/server/package.json

@@ -60,6 +60,7 @@
     "morgan": "^1.10.0",
     "nanoid": "^5.0.4",
     "nodemailer": "^6.9.8",
+    "p-map": "4.0.0",
     "passport": "^0.7.0",
     "passport-jwt": "^4.0.1",
     "ping": "^0.4.4",
@@ -101,7 +102,6 @@
     "@types/tcp-ping": "^0.1.5",
     "@types/uuid": "^9.0.7",
     "execa": "^5.1.1",
-    "p-map": "4.0.0",
     "prisma": "5.14.0",
     "prisma-json-types-generator": "3.0.3",
     "prisma-zod-generator": "0.8.13",

src/server/prisma/migrations/20241106143832_add_workspace_paused_field/migration.sql

@@ -0,0 +1,2 @@
+-- AlterTable
+ALTER TABLE "Workspace" ADD COLUMN     "paused" BOOLEAN NOT NULL DEFAULT false;

src/server/prisma/schema.prisma

@@ -96,6 +96,7 @@ model Workspace {
   /// [CommonPayload]
   /// @zod.custom(imports.CommonPayloadSchema)
   settings        Json     @default("{}")
+  paused          Boolean  @default(false) // if workspace over billing, its will marked as pause and not receive and input.
   createdAt       DateTime @default(now()) @db.Timestamptz(6)
   updatedAt       DateTime @updatedAt @db.Timestamptz(6)
 

src/server/prisma/zod/workspace.ts

@@ -20,6 +20,7 @@ export const WorkspaceModelSchema = z.object({
    * [CommonPayload]
    */
   settings: imports.CommonPayloadSchema,
+  paused: z.boolean(),
   createdAt: z.date(),
   updatedAt: z.date(),
 })

src/server/trpc/routers/billing.ts

@@ -16,6 +16,7 @@ import {
   SubscriptionTierType,
 } from '../../model/billing/index.js';
 import { LemonSqueezySubscriptionModelSchema } from '../../prisma/zod/lemonsqueezysubscription.js';
+import { getWorkspaceUsage } from '../../model/billing/workspace.js';
 
 export const billingRouter = router({
   usage: workspaceProcedure
@@ -44,30 +45,7 @@ export const billingRouter = router({
     .query(async ({ input }) => {
       const { workspaceId, startAt, endAt } = input;
 
-      const res = await prisma.workspaceDailyUsage.aggregate({
-        where: {
-          workspaceId,
-          date: {
-            gte: new Date(startAt),
-            lte: new Date(endAt),
-          },
-        },
-        _sum: {
-          websiteAcceptedCount: true,
-          websiteEventCount: true,
-          monitorExecutionCount: true,
-          surveyCount: true,
-          feedEventCount: true,
-        },
-      });
-
-      return {
-        websiteAcceptedCount: res._sum.websiteAcceptedCount ?? 0,
-        websiteEventCount: res._sum.websiteEventCount ?? 0,
-        monitorExecutionCount: res._sum.monitorExecutionCount ?? 0,
-        surveyCount: res._sum.surveyCount ?? 0,
-        feedEventCount: res._sum.feedEventCount ?? 0,
-      };
+      return getWorkspaceUsage(workspaceId, startAt, endAt);
     }),
   currentSubscription: workspaceProcedure
     .meta(

src/server/trpc/routers/workspace.ts

@@ -28,6 +28,7 @@ import { WorkspacesOnUsersModelSchema } from '../../prisma/zod/workspacesonusers
 import { monitorManager } from '../../model/monitor/index.js';
 import { get, merge } from 'lodash-es';
 import { promWorkspaceCounter } from '../../utils/prometheus/client.js';
+import { getWorkspaceServiceCount } from '../../model/workspace.js';
 
 export const workspaceRouter = router({
   create: protectProedure
@@ -382,39 +383,8 @@ export const workspaceRouter = router({
     .query(async ({ input }) => {
       const { workspaceId } = input;
 
-      const [website, monitor, telemetry, page, survey, feed] =
-        await Promise.all([
-          prisma.website.count({
-            where: {
-              workspaceId,
-            },
-          }),
-          prisma.monitor.count({
-            where: {
-              workspaceId,
-            },
-          }),
-          prisma.telemetry.count({
-            where: {
-              workspaceId,
-            },
-          }),
-          prisma.monitorStatusPage.count({
-            where: {
-              workspaceId,
-            },
-          }),
-          prisma.survey.count({
-            where: {
-              workspaceId,
-            },
-          }),
-          prisma.feedChannel.count({
-            where: {
-              workspaceId,
-            },
-          }),
-        ]);
+      const { website, monitor, telemetry, page, survey, feed } =
+        await getWorkspaceServiceCount(workspaceId);
 
       const server = getServerCount(workspaceId);
 

src/server/utils/env.ts

@@ -47,6 +47,7 @@ export const env = {
     },
   },
   billing: {
+    enable: process.env.ENABLE_BILLING,
     lemonSqueezy: {
       signatureSecret: process.env.LEMON_SQUEEZY_SIGNATURE_SECRET ?? '',
       apiKey: process.env.LEMON_SQUEEZY_API_KEY ?? '',