Skip to content

Commit

Permalink
server: fix drive quota for remote users
Browse files Browse the repository at this point in the history
  • Loading branch information
atsu1125 committed Feb 4, 2024
1 parent f1d24bb commit b081dd3
Showing 1 changed file with 40 additions and 2 deletions.
42 changes: 40 additions & 2 deletions src/services/drive/add-file.ts
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
import * as fs from 'fs';

import { v4 as uuid } from 'uuid';
import { In, IsNull, getConnection } from 'typeorm';

import { publishMainStream, publishDriveStream } from '../stream';
import { deleteFile } from './delete-file';
Expand Down Expand Up @@ -289,6 +290,40 @@ async function deleteOldFile(user: IRemoteUser) {
}
}

async function expireOldFiles(user: IRemoteUser, driveCapacity: number): Promise<void> {
// Delete as many files as necessary so the total usage is below driveCapacity,
// oldest files first, and exclude avatar and banner.
//
// Using a window function, i.e. `OVER (ORDER BY "createdAt" DESC)` means that
// the `SUM` will be a running total.
const conn = await getConnection();
const exceededFileIds = await conn.query('SELECT "id" FROM ('
+ 'SELECT "id", SUM("size") OVER (ORDER BY "createdAt" DESC) AS "total" FROM "drive_file" WHERE "userId" = $1 AND NOT "isLink"'
+ (user.avatarId ? ' AND "id" != $2' : '')
+ (user.bannerId ? ' AND "id" != $3' : '')
+ ') AS "totals" WHERE "total" > $4',
[
user.id,
user.avatarId ?? '',
user.bannerId ?? '',
driveCapacity,
]
);

if (exceededFileIds.length === 0) {
// no files to expire, avatar and banner if present are already the only files
throw new Error('remote user drive quota met by avatar and banner');
}

const files = await DriveFiles.find({
id: In(exceededFileIds.map(x => x.id)),
});

for (const file of files) {
deleteFile(file, true);
}
}

/**
* Add file to drive
*
Expand Down Expand Up @@ -350,6 +385,7 @@ export default async function(
//#region Check drive usage
if (user && !isLink) {
const usage = await DriveFiles.calcDriveUsageOf(user);
const isLocalUser = Users.isLocalUser(user);

const instance = await fetchMeta();
const driveCapacity = 1024 * 1024 * getDriveCapacity(user, instance);
Expand All @@ -358,11 +394,13 @@ export default async function(

// If usage limit exceeded
if (usage + info.size > driveCapacity) {
if (Users.isLocalUser(user)) {
if (isLocalUser) {
// throw new IdentifiableError('c6244ed2-a39a-4e1c-bf93-f0fbd7764fa6', 'No free space.');
} else {
// (アバターまたはバナーを含まず)最も古いファイルを削除する
deleteOldFile(user as IRemoteUser);
//deleteOldFile(user as IRemoteUser);
// delete older files to make space for new file
expireOldFiles(await Users.findOneOrFail({ id: user.id }) as IRemoteUser, driveCapacity - info.size);
}
}
}
Expand Down

0 comments on commit b081dd3

Please sign in to comment.