Signal/Signal-Server
1
0
Fork 0
mirror of https://github.com/signalapp/Signal-Server synced 2026-04-20 10:38:01 +01:00
Code Issues Packages Projects Releases Wiki Activity

Add a crawler to recalculate quota usage

Browse Source
This commit is contained in:
Ravi Khadiwala
2025-05-28 12:51:56 -05:00
committed by ravi-signal
parent 4dc3b19d2a
commit a7ea42adc3
6 changed files with 191 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

20
service/src/main/java/org/whispersystems/textsecuregcm/backup/BackupManager.java
View File

@@ -24,6 +24,7 @@ import java.util.Optional;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.CompletionStage;
import java.util.function.Function;
import org.apache.commons.lang3.StringUtils;
import org.signal.libsignal.protocol.ecc.Curve;
import org.signal.libsignal.protocol.ecc.ECPublicKey;
import org.signal.libsignal.zkgroup.GenericServerSecretParams;
@@ -337,6 +338,19 @@ public class BackupManager {
});
}
public record RecalculationResult(UsageInfo oldUsage, UsageInfo newUsage) {}
public CompletionStage<Optional<RecalculationResult>> recalculateQuota(final StoredBackupAttributes storedBackupAttributes) {
if (StringUtils.isBlank(storedBackupAttributes.backupDir()) || StringUtils.isBlank(storedBackupAttributes.mediaDir())) {
return CompletableFuture.completedFuture(Optional.empty());
}
final String cdnPath = cdnMediaDirectory(storedBackupAttributes.backupDir(), storedBackupAttributes.mediaDir());
return this.remoteStorageManager.calculateBytesUsed(cdnPath).thenCompose(usage ->
backupsDb.setMediaUsage(storedBackupAttributes, usage).thenApply(ignored ->
Optional.of(new RecalculationResult(
new UsageInfo(storedBackupAttributes.bytesUsed(), storedBackupAttributes.numObjects()),
usage))));
}
/**
* @return the largest index i such that sum(ts[0],...ts[i - 1]) <= max
*/
@@ -735,8 +749,12 @@ public class BackupManager {
return "%s/%s".formatted(backupUser.backupDir(), MESSAGE_BACKUP_NAME);
}
private static String cdnMediaDirectory(final String backupDir, final String mediaDir) {
return "%s/%s/".formatted(backupDir, mediaDir);
}
private static String cdnMediaDirectory(final AuthenticatedBackupUser backupUser) {
return "%s/%s/".formatted(backupUser.backupDir(), backupUser.mediaDir());
return cdnMediaDirectory(backupUser.backupDir(), backupUser.mediaDir());
}
private static String cdnMediaPath(final AuthenticatedBackupUser backupUser, final byte[] mediaId) {

19
service/src/main/java/org/whispersystems/textsecuregcm/backup/BackupsDb.java
View File

@@ -402,8 +402,16 @@ public class BackupsDb {
}
CompletableFuture<Void> setMediaUsage(final AuthenticatedBackupUser backupUser, UsageInfo usageInfo) {
return setMediaUsage(UpdateBuilder.forUser(backupTableName, backupUser), usageInfo);
}
CompletableFuture<Void> setMediaUsage(final StoredBackupAttributes backupAttributes, UsageInfo usageInfo) {
return setMediaUsage(new UpdateBuilder(backupTableName, BackupLevel.PAID, backupAttributes.hashedBackupId()), usageInfo);
}
private CompletableFuture<Void> setMediaUsage(final UpdateBuilder updateBuilder, UsageInfo usageInfo) {
return dynamoClient.updateItem(
UpdateBuilder.forUser(backupTableName, backupUser)
updateBuilder
.addSetExpression("#mediaBytesUsed = :mediaBytesUsed",
Map.entry("#mediaBytesUsed", ATTR_MEDIA_BYTES_USED),
Map.entry(":mediaBytesUsed", AttributeValues.n(usageInfo.bytesUsed())))
@@ -496,13 +504,18 @@ public class BackupsDb {
"#refresh", ATTR_LAST_REFRESH,
"#mediaRefresh", ATTR_LAST_MEDIA_REFRESH,
"#bytesUsed", ATTR_MEDIA_BYTES_USED,
"#numObjects", ATTR_MEDIA_COUNT))
.projectionExpression("#backupIdHash, #refresh, #mediaRefresh, #bytesUsed, #numObjects")
"#numObjects", ATTR_MEDIA_COUNT,
"#backupDir", ATTR_BACKUP_DIR,
"#mediaDir", ATTR_MEDIA_DIR))
.projectionExpression("#backupIdHash, #refresh, #mediaRefresh, #bytesUsed, #numObjects, #backupDir, #mediaDir")
.build())
.items())
.sequential()
.filter(item -> item.containsKey(KEY_BACKUP_ID_HASH))
.map(item -> new StoredBackupAttributes(
AttributeValues.getByteArray(item, KEY_BACKUP_ID_HASH, null),
AttributeValues.getString(item, ATTR_BACKUP_DIR, null),
AttributeValues.getString(item, ATTR_MEDIA_DIR, null),
Instant.ofEpochSecond(AttributeValues.getLong(item, ATTR_LAST_REFRESH, 0L)),
Instant.ofEpochSecond(AttributeValues.getLong(item, ATTR_LAST_MEDIA_REFRESH, 0L)),
AttributeValues.getLong(item, ATTR_MEDIA_BYTES_USED, 0L),

12
service/src/main/java/org/whispersystems/textsecuregcm/backup/StoredBackupAttributes.java
View File

@@ -9,11 +9,19 @@ import java.time.Instant;
/**
* Attributes stored in the backups table for a single backup id
*
* @param hashedBackupId The hashed backup-id of this entry
* @param backupDir The cdn backupDir of this entry
* @param mediaDir The cdn mediaDir (within the backupDir) of this entry
* @param lastRefresh The last time the record was updated with a messages or media tier credential
* @param lastMediaRefresh The last time the record was updated with a media tier credential
* @param bytesUsed The number of media bytes used by the backup
* @param numObjects The number of media objects used byt the backup
*/
public record StoredBackupAttributes(
Instant lastRefresh, Instant lastMediaRefresh,
long bytesUsed, long numObjects) {}
byte[] hashedBackupId,
String backupDir,
String mediaDir,
Instant lastRefresh,
Instant lastMediaRefresh,
long bytesUsed,
long numObjects) {}