@protontech/drive-sdk 0.8.0 → 0.9.1

package/src/internal/upload/streamUploader.test.ts

@@ -1,5 +1,6 @@
-import { Thumbnail, ThumbnailType, UploadMetadata } from '../../interface';
+import { Logger, Thumbnail, ThumbnailType, UploadMetadata } from '../../interface';
 import { IntegrityError } from '../../errors';
+import { getMockLogger } from '../../tests/logger';
 import { APIHTTPError, HTTPErrorCode } from '../apiService';
 import { FILE_CHUNK_SIZE, StreamUploader } from './streamUploader';
 import { UploadTelemetry } from './telemetry';
@@ -40,6 +41,7 @@ function mockUploadBlock(
 }
 
 describe('StreamUploader', () => {
+    let logger: Logger;
     let telemetry: UploadTelemetry;
     let apiService: jest.Mocked<UploadAPIService>;
     let cryptoService: UploadCryptoService;
@@ -54,14 +56,11 @@ describe('StreamUploader', () => {
     let uploader: StreamUploader;
 
     beforeEach(() => {
+        logger = getMockLogger();
+
         // @ts-expect-error No need to implement all methods for mocking
         telemetry = {
-            getLoggerForRevision: jest.fn().mockReturnValue({
-                debug: jest.fn(),
-                info: jest.fn(),
-                warn: jest.fn(),
-                error: jest.fn(),
-            }),
+            getLoggerForRevision: jest.fn().mockReturnValue(logger),
             logBlockVerificationError: jest.fn(),
             uploadFailed: jest.fn(),
             uploadFinished: jest.fn(),
@@ -403,6 +402,80 @@ describe('StreamUploader', () => {
             await verifyOnProgress([1024, 4 * 1024 * 1024, 2 * 1024 * 1024, 4 * 1024 * 1024]);
         });
 
+        it('should handle timeout when uploading block', async () => {
+            const error = new Error('TimeoutError');
+            error.name = 'TimeoutError';
+
+            let count = 0;
+            apiService.uploadBlock = jest.fn().mockImplementation(async function (bareUrl, token, block, onProgress) {
+                if (token === 'token/block:1' && count === 0) {
+                    count++;
+                    throw error;
+                }
+                return mockUploadBlock(bareUrl, token, block, onProgress);
+            });
+
+            expect((uploader as any).maxUploadingBlocks).toEqual(5);
+            await verifySuccess();
+            expect(apiService.requestBlockUpload).toHaveBeenCalledTimes(1);
+            // 3 blocks + 1 timeout retry + 1 thumbnail
+            expect(apiService.uploadBlock).toHaveBeenCalledTimes(5);
+            expect(logger.warn).toHaveBeenCalledTimes(2);
+            expect(logger.warn).toHaveBeenCalledWith(
+                'block 1:token/block:1: Upload timeout, limiting upload capacity to 1 block',
+            );
+            expect(logger.warn).toHaveBeenCalledWith('block 1:token/block:1: Upload timeout, retrying');
+            expect((uploader as any).maxUploadingBlocks).toEqual(1);
+        });
+
+        it('limitUploadCapacity should wait for the previous blocks to finish', async () => {
+            const error = new Error('TimeoutError');
+            error.name = 'TimeoutError';
+
+            const events: string[] = [];
+            let block1Resolver: (() => void) | undefined;
+            let block2FirstAttempt = true;
+
+            apiService.uploadBlock = jest.fn().mockImplementation(async function (bareUrl, token, block, onProgress) {
+                if (token === 'token/block:1') {
+                    events.push('block1:upload:start');
+                    await new Promise<void>((resolve) => {
+                        block1Resolver = resolve;
+                    });
+                    events.push('block1:upload:end');
+                    return mockUploadBlock(bareUrl, token, block, onProgress);
+                }
+                if (token === 'token/block:2') {
+                    if (block2FirstAttempt) {
+                        block2FirstAttempt = false;
+                        events.push('block2:timeout');
+                        // Resolve block 1 after a small delay to simulate real-world conditions
+                        setTimeout(() => block1Resolver?.(), 100);
+                        throw error;
+                    }
+                    events.push('block2:retry');
+                    return mockUploadBlock(bareUrl, token, block, onProgress);
+                }
+                // Block 3 and thumbnails proceed normally
+                return mockUploadBlock(bareUrl, token, block, onProgress);
+            });
+
+            await verifySuccess();
+
+            expect(events).toMatchObject([
+                'block1:upload:start',
+                'block2:timeout',
+                'block1:upload:end',
+                'block2:retry',
+            ]);
+
+            // Also verify the warning messages were logged
+            expect(logger.warn).toHaveBeenCalledWith(
+                'block 2:token/block:2: Upload timeout, limiting upload capacity to 1 block',
+            );
+            expect(logger.warn).toHaveBeenCalledWith('block 2:token/block:2: Upload timeout, retrying');
+        });
+
         it('should handle expired token when uploading block', async () => {
             let count = 0;
             apiService.uploadBlock = jest.fn().mockImplementation(async function (bareUrl, token, block, onProgress) {

package/src/internal/upload/streamUploader.ts

@@ -55,6 +55,8 @@ const MAX_BLOCK_UPLOAD_RETRIES = 3;
  * that the upload process is efficient and does not overload the server.
  */
 export class StreamUploader {
+    protected maxUploadingBlocks = MAX_UPLOADING_BLOCKS;
+
     protected logger: Logger;
 
     protected digests: UploadDigests;
@@ -67,6 +69,7 @@ export class StreamUploader {
     protected ongoingUploads = new Map<
         string,
         {
+            index?: number;
             uploadPromise: Promise<void>;
             encryptedBlock: EncryptedBlock | EncryptedThumbnail;
         }
@@ -181,7 +184,7 @@ export class StreamUploader {
             await this.controller.waitWhilePaused();
             await this.waitForUploadCapacityAndBufferedBlocks();
 
-            if (this.isEncryptionFullyFinished) {
+            if (this.isEncryptionFullyFinished || this.isUploadAborted) {
                 break;
             }
 
@@ -376,6 +379,7 @@ export class StreamUploader {
 
             const uploadKey = `block:${blockToken.index}`;
             this.ongoingUploads.set(uploadKey, {
+                index: blockToken.index,
                 uploadPromise: this.uploadBlock(blockToken, encryptedBlock, onProgress).finally(() => {
                     this.ongoingUploads.delete(uploadKey);
 
@@ -500,6 +504,13 @@ export class StreamUploader {
                     blockProgress = 0;
                 }
 
+                if (error instanceof Error && error.name === 'TimeoutError') {
+                    logger.warn(`Upload timeout, limiting upload capacity to 1 block`);
+                    await this.limitUploadCapacity(uploadToken.index);
+                    logger.warn(`Upload timeout, retrying`);
+                    continue;
+                }
+
                 if (
                     (error instanceof APIHTTPError && error.statusCode === HTTPErrorCode.NOT_FOUND) ||
                     error instanceof NotFoundAPIError
@@ -542,6 +553,27 @@ export class StreamUploader {
         logger.info(`Uploaded`);
     }
 
+    private async limitUploadCapacity(index: number) {
+        this.maxUploadingBlocks = 1;
+
+        // This ensures that when the upload is downscaled, all ongoing block
+        // uploads are waiting for their turn one by one.
+        try {
+            await waitForCondition(() => {
+                const ongoingIndexes = Array.from(this.ongoingUploads.values())
+                    .map(({ index: ongoingIndex }) => ongoingIndex)
+                    .filter((ongoingIndex) => ongoingIndex !== undefined);
+                ongoingIndexes.sort((a, b) => a - b);
+                return ongoingIndexes[0] === index;
+            }, this.abortController.signal);
+        } catch (error: unknown) {
+            if (error instanceof AbortError) {
+                return;
+            }
+            throw error;
+        }
+    }
+
     private async waitForBufferCapacity() {
         if (this.encryptedBlocks.size >= MAX_BUFFERED_BLOCKS) {
             try {
@@ -559,7 +591,7 @@ export class StreamUploader {
     }
 
     private async waitForUploadCapacityAndBufferedBlocks() {
-        while (this.ongoingUploads.size >= MAX_UPLOADING_BLOCKS) {
+        while (this.ongoingUploads.size >= this.maxUploadingBlocks) {
             await Promise.race(this.ongoingUploads.values().map(({ uploadPromise }) => uploadPromise));
         }
         try {

package/src/internal/upload/telemetry.ts

@@ -13,13 +13,12 @@ export class UploadTelemetry {
         private sharesService: SharesService,
     ) {
         this.telemetry = telemetry;
-        this.logger = this.telemetry.getLogger('download');
+        this.logger = this.telemetry.getLogger('upload');
         this.sharesService = sharesService;
     }
 
     getLoggerForRevision(revisionUid: string) {
-        const logger = this.telemetry.getLogger('upload');
-        return new LoggerWithPrefix(logger, `revision ${revisionUid}`);
+        return new LoggerWithPrefix(this.logger, `revision ${revisionUid}`);
     }
 
     logBlockVerificationError(retryHelped: boolean) {

package/src/protonDriveClient.ts

@@ -514,7 +514,7 @@ export class ProtonDriveClient {
 
     async emptyTrash(): Promise<void> {
         this.logger.info('Emptying trash');
-        throw new Error('Method not implemented');
+        return this.nodes.management.emptyTrash();
     }
 
     /**