feat: add user level blob quota (#4114)

This commit is contained in:
DarkSky
2023-09-01 20:56:27 +08:00
committed by GitHub
parent e76cdf4d71
commit 7d3b1ad2b9
8 changed files with 99 additions and 20 deletions
Generated
+5 -5
View File
@@ -1695,7 +1695,7 @@ dependencies = [
[[package]]
name = "jwst"
version = "0.1.1"
source = "git+https://github.com/toeverything/OctoBase.git#b026b44f67e1043e83f31a21360a0baee122e819"
source = "git+https://github.com/toeverything/OctoBase.git#58f3bbdf97f391a535e772d32828a484376c4159"
dependencies = [
"async-trait",
"base64 0.21.3",
@@ -1723,7 +1723,7 @@ dependencies = [
[[package]]
name = "jwst-codec"
version = "0.1.0"
source = "git+https://github.com/toeverything/OctoBase.git#b026b44f67e1043e83f31a21360a0baee122e819"
source = "git+https://github.com/toeverything/OctoBase.git#58f3bbdf97f391a535e772d32828a484376c4159"
dependencies = [
"arbitrary",
"bitvec",
@@ -1743,7 +1743,7 @@ dependencies = [
[[package]]
name = "jwst-logger"
version = "0.1.0"
source = "git+https://github.com/toeverything/OctoBase.git#b026b44f67e1043e83f31a21360a0baee122e819"
source = "git+https://github.com/toeverything/OctoBase.git#58f3bbdf97f391a535e772d32828a484376c4159"
dependencies = [
"chrono",
"nu-ansi-term",
@@ -1756,7 +1756,7 @@ dependencies = [
[[package]]
name = "jwst-storage"
version = "0.1.0"
source = "git+https://github.com/toeverything/OctoBase.git#b026b44f67e1043e83f31a21360a0baee122e819"
source = "git+https://github.com/toeverything/OctoBase.git#58f3bbdf97f391a535e772d32828a484376c4159"
dependencies = [
"anyhow",
"async-trait",
@@ -1786,7 +1786,7 @@ dependencies = [
[[package]]
name = "jwst-storage-migration"
version = "0.1.0"
source = "git+https://github.com/toeverything/OctoBase.git#b026b44f67e1043e83f31a21360a0baee122e819"
source = "git+https://github.com/toeverything/OctoBase.git#58f3bbdf97f391a535e772d32828a484376c4159"
dependencies = [
"sea-orm-migration",
"tokio",
+5
View File
@@ -186,6 +186,11 @@ export interface AFFiNEConfig {
fs: {
path: string;
};
/**
* Free user storage quota
* @default 10 * 1024 * 1024 (10GB)
*/
quota: number;
};
/**
+2
View File
@@ -55,6 +55,7 @@ export const getDefaultAFFiNEConfig: () => AFFiNEConfig = () => {
AFFINE_SERVER_HOST: 'host',
AFFINE_SERVER_SUB_PATH: 'path',
AFFINE_ENV: 'affineEnv',
AFFINE_FREE_USER_QUOTA: 'objectStorage.quota',
DATABASE_URL: 'db.url',
ENABLE_R2_OBJECT_STORAGE: ['objectStorage.r2.enabled', 'boolean'],
R2_OBJECT_STORAGE_ACCOUNT_ID: 'objectStorage.r2.accountId',
@@ -170,6 +171,7 @@ export const getDefaultAFFiNEConfig: () => AFFiNEConfig = () => {
fs: {
path: join(homedir(), '.affine-storage'),
},
quota: 10 * 1024 * 1024,
},
rateLimiter: {
ttl: 60,
+30 -7
View File
@@ -27,6 +27,7 @@ import type { User, Workspace } from '@prisma/client';
import GraphQLUpload from 'graphql-upload/GraphQLUpload.mjs';
import { applyUpdate, Doc } from 'yjs';
import { Config } from '../../config';
import { PrismaService } from '../../prisma';
import { StorageProvide } from '../../storage';
import { CloudThrottlerGuard, Throttle } from '../../throttler';
@@ -130,6 +131,7 @@ export class UpdateWorkspaceInput extends PickType(
export class WorkspaceResolver {
constructor(
private readonly auth: AuthService,
private readonly config: Config,
private readonly mailer: MailService,
private readonly prisma: PrismaService,
private readonly permissionProvider: PermissionService,
@@ -610,17 +612,28 @@ export class WorkspaceResolver {
) {
await this.permissionProvider.check(workspaceId, user.id);
return this.storage.blobsSize(workspaceId).then(size => ({ size }));
return this.storage.blobsSize([workspaceId]).then(size => ({ size }));
}
@Query(() => WorkspaceBlobSizes)
async collectAllBlobSizes(@CurrentUser() user: User) {
const workspaces = await this.workspaces(user);
const size = (
await Promise.all(workspaces.map(({ id }) => this.storage.blobsSize(id)))
).reduce((prev, curr) => prev + curr, 0);
async collectAllBlobSizes(@CurrentUser() user: UserType) {
const workspaces = await this.prisma.userWorkspacePermission
.findMany({
where: {
userId: user.id,
accepted: true,
},
select: {
workspace: {
select: {
id: true,
},
},
},
})
.then(data => data.map(({ workspace }) => workspace.id));
const size = await this.storage.blobsSize(workspaces);
return { size };
}
@@ -632,6 +645,12 @@ export class WorkspaceResolver {
blob: FileUpload
) {
await this.permissionProvider.check(workspaceId, user.id, Permission.Write);
const quota = this.config.objectStorage.quota;
const { size } = await this.collectAllBlobSizes(user);
if (size > quota) {
throw new ForbiddenException('storage size limit exceeded');
}
const buffer = await new Promise<Buffer>((resolve, reject) => {
const stream = blob.createReadStream();
@@ -645,6 +664,10 @@ export class WorkspaceResolver {
});
});
if (size + buffer.length > quota) {
throw new ForbiddenException('storage size limit exceeded');
}
return this.storage.uploadBlob(workspaceId, buffer);
}
+21 -1
View File
@@ -339,7 +339,6 @@ async function collectBlobSizes(
const res = await request(app.getHttpServer())
.post(gql)
.auth(token, { type: 'bearer' })
.set({ 'x-request-id': 'test', 'x-operation-name': 'test' })
.send({
query: `
query {
@@ -353,6 +352,26 @@ async function collectBlobSizes(
return res.body.data.collectBlobSizes.size;
}
async function collectAllBlobSizes(
app: INestApplication,
token: string
): Promise<number> {
const res = await request(app.getHttpServer())
.post(gql)
.auth(token, { type: 'bearer' })
.send({
query: `
query {
collectAllBlobSizes {
size
}
}
`,
})
.expect(200);
return res.body.data.collectAllBlobSizes.size;
}
async function setBlob(
app: INestApplication,
token: string,
@@ -447,6 +466,7 @@ async function getInviteInfo(
export {
acceptInvite,
acceptInviteById,
collectAllBlobSizes,
collectBlobSizes,
createTestApp,
createWorkspace,
@@ -10,6 +10,7 @@ import request from 'supertest';
import { AppModule } from '../app';
import {
collectAllBlobSizes,
collectBlobSizes,
createWorkspace,
listBlobs,
@@ -108,4 +109,25 @@ describe('Workspace Module - Blobs', () => {
const size = await collectBlobSizes(app, u1.token.token, workspace.id);
ok(size === 4, 'failed to collect blob sizes');
});
it('should calc all blobs size', async () => {
const u1 = await signUp(app, 'u1', 'u1@affine.pro', '1');
const workspace1 = await createWorkspace(app, u1.token.token);
const buffer1 = Buffer.from([0, 0]);
await setBlob(app, u1.token.token, workspace1.id, buffer1);
const buffer2 = Buffer.from([0, 1]);
await setBlob(app, u1.token.token, workspace1.id, buffer2);
const workspace2 = await createWorkspace(app, u1.token.token);
const buffer3 = Buffer.from([0, 0]);
await setBlob(app, u1.token.token, workspace2.id, buffer3);
const buffer4 = Buffer.from([0, 1]);
await setBlob(app, u1.token.token, workspace2.id, buffer4);
const size = await collectAllBlobSizes(app, u1.token.token);
ok(size === 8, 'failed to collect all blob sizes');
});
});
+5 -2
View File
@@ -16,7 +16,7 @@ export class Storage {
/** Delete a blob from workspace storage. */
deleteBlob(workspaceId: string, hash: string): Promise<boolean>;
/** Workspace size taken by blobs. */
blobsSize(workspaceId: string): Promise<number>;
blobsSize(workspaces: Array<string>): Promise<number>;
}
export interface Blob {
@@ -26,5 +26,8 @@ export interface Blob {
data: Buffer;
}
/** Merge updates in form like `Y.applyUpdate(doc, update)` way and return the result binary. */
/**
* Merge updates in form like `Y.applyUpdate(doc, update)` way and return the
* result binary.
*/
export function mergeUpdatesInApplyWay(updates: Array<Buffer>): Buffer;
+9 -5
View File
@@ -9,7 +9,6 @@ use std::{
use jwst::BlobStorage;
use jwst_codec::Doc;
use jwst_storage::{BlobStorageType, JwstStorage, JwstStorageError};
use napi::{bindgen_prelude::*, Error, Result, Status};
#[macro_use]
@@ -112,7 +111,11 @@ impl Storage {
(id, ext.map(|ext| HashMap::from([("format".into(), ext)])))
};
let Ok(meta) = self.blobs().get_metadata(Some(workspace_id.clone()), id.clone(), params.clone()).await else {
let Ok(meta) = self
.blobs()
.get_metadata(Some(workspace_id.clone()), id.clone(), params.clone())
.await
else {
return Ok(None);
};
@@ -144,12 +147,13 @@ impl Storage {
/// Workspace size taken by blobs.
#[napi]
pub async fn blobs_size(&self, workspace_id: String) -> Result<i64> {
map_err!(self.blobs().get_blobs_size(workspace_id).await)
pub async fn blobs_size(&self, workspaces: Vec<String>) -> Result<i64> {
map_err!(self.blobs().get_blobs_size(workspaces).await)
}
}
/// Merge updates in form like `Y.applyUpdate(doc, update)` way and return the result binary.
/// Merge updates in form like `Y.applyUpdate(doc, update)` way and return the
/// result binary.
#[napi(catch_unwind)]
pub fn merge_updates_in_apply_way(updates: Vec<Buffer>) -> Result<Buffer> {
let mut doc = Doc::default();