fix(core): revalidate quota and subscription when subscribing

toeverything · Dec 20, 2024 · f95b062 · f95b062
1 parent a53e231
commit f95b062
Show file tree

Hide file tree

Showing 5 changed files with 47 additions and 28 deletions.
diff --git a/packages/frontend/core/src/components/affine/quota-reached-modal/cloud-quota-modal.tsx b/packages/frontend/core/src/components/affine/quota-reached-modal/cloud-quota-modal.tsx
@@ -7,10 +7,10 @@ import { WorkspaceQuotaService } from '@affine/core/modules/quota';
 import { type I18nString, useI18n } from '@affine/i18n';
 import { track } from '@affine/track';
 import { useLiveData, useService, WorkspaceService } from '@toeverything/infra';
-import bytes from 'bytes';
 import { useAtom } from 'jotai';
 import { useCallback, useEffect, useMemo } from 'react';
 
+import { useAsyncCallback } from '../../hooks/affine-async-hooks';
 import * as styles from './cloud-quota-modal.css';
 
 export const CloudQuotaModal = () => {
@@ -66,22 +66,32 @@ export const CloudQuotaModal = () => {
     }
   }, [userQuota, isOwner, workspaceQuota, t]);
 
+  const onAbortLargeBlob = useAsyncCallback(
+    async (blob: Blob) => {
+      // wait for quota revalidation
+      await workspaceQuotaService.quota.waitForRevalidation();
+      if (
+        blob.size > (workspaceQuotaService.quota.quota$.value?.blobLimit ?? 0)
+      ) {
+        setOpen(true);
+      }
+    },
+    [setOpen, workspaceQuotaService]
+  );
+
   useEffect(() => {
     if (!workspaceQuota) {
       return;
     }
-    // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
-    currentWorkspace.engine.blob.singleBlobSizeLimit = bytes.parse(
-      workspaceQuota.blobLimit.toString()
-    )!;
+
+    currentWorkspace.engine.blob.singleBlobSizeLimit = workspaceQuota.blobLimit;
 
-    const disposable = currentWorkspace.engine.blob.onAbortLargeBlob.on(() => {
-      setOpen(true);
-    });
+    const disposable =
+      currentWorkspace.engine.blob.onAbortLargeBlob.on(onAbortLargeBlob);
     return () => {
       disposable?.dispose();
     };
-  }, [currentWorkspace.engine.blob, setOpen, workspaceQuota]);
+  }, [currentWorkspace.engine.blob, onAbortLargeBlob, workspaceQuota]);
 
   return (
     <ConfirmModal

diff --git a/packages/frontend/core/src/desktop/dialogs/setting/general-setting/plans/checkout-slot.tsx b/packages/frontend/core/src/desktop/dialogs/setting/general-setting/plans/checkout-slot.tsx
@@ -1,12 +1,16 @@
 import { useAsyncCallback } from '@affine/core/components/hooks/affine-async-hooks';
-import { SubscriptionService } from '@affine/core/modules/cloud';
+import {
+  SubscriptionService,
+  UserQuotaService,
+} from '@affine/core/modules/cloud';
 import { UrlService } from '@affine/core/modules/url';
 import type { CreateCheckoutSessionInput } from '@affine/graphql';
 import { useService } from '@toeverything/infra';
 import { nanoid } from 'nanoid';
 import {
   type PropsWithChildren,
   type ReactNode,
+  useCallback,
   useEffect,
   useState,
 } from 'react';
@@ -35,26 +39,23 @@ export const CheckoutSlot = ({
   const urlService = useService(UrlService);
 
   const subscriptionService = useService(SubscriptionService);
+  const userQuotaService = useService(UserQuotaService);
+
+  const revalidate = useCallback(() => {
+    subscriptionService.subscription.revalidate();
+    userQuotaService.quota.revalidate();
+  }, [subscriptionService, userQuotaService]);
 
-  useEffect(() => {
-    subscriptionService.prices.revalidate();
-  }, [subscriptionService]);
   useEffect(() => {
     if (isOpenedExternalWindow) {
       // when the external window is opened, revalidate the subscription when window get focus
-      window.addEventListener(
-        'focus',
-        subscriptionService.subscription.revalidate
-      );
+      window.addEventListener('focus', revalidate);
       return () => {
-        window.removeEventListener(
-          'focus',
-          subscriptionService.subscription.revalidate
-        );
+        window.removeEventListener('focus', revalidate);
       };
     }
     return;
-  }, [isOpenedExternalWindow, subscriptionService]);
+  }, [isOpenedExternalWindow, revalidate, subscriptionService]);
 
   const subscribe = useAsyncCallback(async () => {
     setMutating(true);

diff --git a/...gs/setting/workspace-setting/new-workspace-setting-detail/members/cloud-members-panel.tsx b/...gs/setting/workspace-setting/new-workspace-setting-detail/members/cloud-members-panel.tsx
@@ -75,7 +75,7 @@ export const CloudWorkspaceMembersPanel = ({
   useEffect(() => {
     workspaceQuotaService.quota.revalidate();
   }, [workspaceQuotaService]);
-  const isLoading = useLiveData(workspaceQuotaService.quota.isLoading$);
+  const isLoading = useLiveData(workspaceQuotaService.quota.isRevalidating$);
   const error = useLiveData(workspaceQuotaService.quota.error$);
   const workspaceQuota = useLiveData(workspaceQuotaService.quota.quota$);
   const subscriptionService = useService(SubscriptionService);

diff --git a/...esktop/dialogs/setting/workspace-setting/new-workspace-setting-detail/workspace-quota.tsx b/...esktop/dialogs/setting/workspace-setting/new-workspace-setting-detail/workspace-quota.tsx
@@ -29,7 +29,7 @@ export const StorageProgress = () => {
   ).permission;
   const workspaceQuotaService = useService(WorkspaceQuotaService).quota;
   const isTeam = useLiveData(workspacePermissionService.isTeam$);
-  const isLoading = useLiveData(workspaceQuotaService.isLoading$);
+  const isLoading = useLiveData(workspaceQuotaService.isRevalidating$);
   const usedFormatted = useLiveData(workspaceQuotaService.usedFormatted$);
   const maxFormatted = useLiveData(workspaceQuotaService.maxFormatted$);
   const percent = useLiveData(workspaceQuotaService.percent$);

diff --git a/packages/frontend/core/src/modules/quota/entities/quota.ts b/packages/frontend/core/src/modules/quota/entities/quota.ts
@@ -25,7 +25,7 @@ const logger = new DebugLogger('affine:workspace-permission');
 
 export class WorkspaceQuota extends Entity {
   quota$ = new LiveData<QuotaType | null>(null);
-  isLoading$ = new LiveData(false);
+  isRevalidating$ = new LiveData(false);
   error$ = new LiveData<any>(null);
 
   /** Used storage in bytes */
@@ -106,18 +106,26 @@ export class WorkspaceQuota extends Entity {
           catchErrorInto(this.error$, error => {
             logger.error('Failed to fetch workspace quota', error);
           }),
-          onStart(() => this.isLoading$.setValue(true)),
-          onComplete(() => this.isLoading$.setValue(false))
+          onStart(() => this.isRevalidating$.setValue(true)),
+          onComplete(() => this.isRevalidating$.setValue(false))
         );
       }
     )
   );
 
+  waitForRevalidation(signal?: AbortSignal) {
+    this.revalidate();
+    return this.isRevalidating$.waitFor(
+      isRevalidating => !isRevalidating,
+      signal
+    );
+  }
+
   reset() {
     this.quota$.next(null);
     this.used$.next(null);
     this.error$.next(null);
-    this.isLoading$.next(false);
+    this.isRevalidating$.next(false);
   }
 
   override dispose(): void {