Improving code as per reviews

Signed-off-by: Vidhan Shah <vidhans958@gmail.com>
2025-11-08 07:58:34 +05:30
parent c13704bf52
commit f2f84ad597
2 changed files with 166 additions and 68 deletions
@@ -22,6 +22,7 @@ import {
  ListObjectsV2Command,
  PutObjectCommand,
  S3Client,
+  S3ServiceException,
  UploadPartCommand,
 } from '@aws-sdk/client-s3';
 import { Entity, DEFAULT_NAMESPACE } from '@backstage/catalog-model';
@@ -196,16 +197,16 @@ describe('AwsS3Publish', () => {
      [directory]: files,
    });

-    s3Mock = mockClient(S3Client as any);
+    s3Mock = mockClient(S3Client);

-    s3Mock.on(HeadObjectCommand).callsFake((input: any) => {
+    s3Mock.on(HeadObjectCommand).callsFake((input: { Key: string }) => {
      if (!fs.pathExistsSync(mockDir.resolve(input.Key))) {
        throw new Error('File does not exist');
      }
      return {};
    });

-    s3Mock.on(GetObjectCommand).callsFake((input: any) => {
+    s3Mock.on(GetObjectCommand).callsFake((input: { Key: string }) => {
      if (fs.pathExistsSync(mockDir.resolve(input.Key))) {
        return {
          Body: Readable.from(fs.readFileSync(mockDir.resolve(input.Key))),
@@ -215,14 +216,14 @@ describe('AwsS3Publish', () => {
      throw new Error(`The file ${input.Key} does not exist!`);
    });

-    s3Mock.on(HeadBucketCommand).callsFake((input: any) => {
+    s3Mock.on(HeadBucketCommand).callsFake((input: { Bucket: string }) => {
      if (input.Bucket === 'errorBucket') {
        throw new Error('Bucket does not exist');
      }
      return {};
    });

-    s3Mock.on(ListObjectsV2Command).callsFake((input: any) => {
+    s3Mock.on(ListObjectsV2Command).callsFake((input: { Bucket: string }) => {
      if (
        input.Bucket === 'delete_stale_files_success' ||
        input.Bucket === 'delete_stale_files_error'
@@ -234,7 +235,7 @@ describe('AwsS3Publish', () => {
      return {};
    });

-    s3Mock.on(DeleteObjectCommand).callsFake((input: any) => {
+    s3Mock.on(DeleteObjectCommand).callsFake((input: { Bucket: string }) => {
      if (input.Bucket === 'delete_stale_files_error') {
        throw new Error('Message');
      }
@@ -242,9 +243,11 @@ describe('AwsS3Publish', () => {
    });

    s3Mock.on(UploadPartCommand).rejects();
-    s3Mock.on(PutObjectCommand).callsFake((input: any) => {
-      mockDir.addContent({ [input.Key]: input.Body });
-    });
+    s3Mock
+      .on(PutObjectCommand)
+      .callsFake((input: { Key: string; Body: any }) => {
+        mockDir.addContent({ [input.Key]: input.Body });
+      });
  });

  afterEach(() => {
@@ -299,6 +302,122 @@ describe('AwsS3Publish', () => {
    });
  });

+  describe('retry mechanism', () => {
+    it('should retry with custom retry strategy', async () => {
+      const publisher = (await createPublisherFromConfig()) as AwsS3Publish;
+      const customRetryStrategy = jest.fn((error: any) => {
+        return error.name === 'NetworkingError';
+      });
+
+      s3Mock
+        .on(ListObjectsV2Command)
+        .rejectsOnce(
+          new S3ServiceException({
+            name: 'NetworkingError',
+            $fault: 'client',
+            $metadata: {},
+          }),
+        )
+        .resolvesOnce({ Contents: [] });
+
+      await (publisher as any).retryOperation(
+        async () => {
+          return s3Mock.send(
+            new ListObjectsV2Command({ Bucket: 'bucketName' }),
+          );
+        },
+        'TestOperation',
+        3,
+        customRetryStrategy,
+      );
+
+      expect(customRetryStrategy).toHaveBeenCalled();
+    });
+
+    it('should use default retry strategy when no custom strategy provided', async () => {
+      const publisher = (await createPublisherFromConfig()) as AwsS3Publish;
+      s3Mock
+        .on(ListObjectsV2Command)
+        .rejectsOnce('RequestTimeout')
+        .resolvesOnce({ Contents: [] });
+
+      await (publisher as any).retryOperation(
+        async () => {
+          return s3Mock.send(
+            new ListObjectsV2Command({ Bucket: 'bucketName' }),
+          );
+        },
+        'TestOperation',
+        3,
+      );
+    });
+
+    it('should retry on server errors (5xx)', async () => {
+      const publisher = (await createPublisherFromConfig()) as AwsS3Publish;
+      s3Mock
+        .on(ListObjectsV2Command)
+        .rejectsOnce(
+          new S3ServiceException({
+            name: 'InternalError',
+            $fault: 'server',
+            $metadata: { httpStatusCode: 500 },
+          }),
+        )
+        .resolvesOnce({ Contents: [] });
+
+      await (publisher as any).retryOperation(
+        async () => {
+          return s3Mock.send(
+            new ListObjectsV2Command({ Bucket: 'bucketName' }),
+          );
+        },
+        'TestOperation',
+        3,
+      );
+    });
+
+    it('should retry on specific 4xx errors', async () => {
+      const publisher = (await createPublisherFromConfig()) as AwsS3Publish;
+      s3Mock
+        .on(ListObjectsV2Command)
+        .rejectsOnce('RequestTimeout')
+        .resolvesOnce({ Contents: [] });
+
+      await (publisher as any).retryOperation(
+        async () => {
+          return s3Mock.send(
+            new ListObjectsV2Command({ Bucket: 'bucketName' }),
+          );
+        },
+        'TestOperation',
+        3,
+      );
+    });
+
+    it('should not retry on non-retriable 4xx errors', async () => {
+      const publisher = (await createPublisherFromConfig()) as AwsS3Publish;
+      s3Mock.on(ListObjectsV2Command).rejectsOnce(
+        new S3ServiceException({
+          name: 'BadRequest',
+          $fault: 'client',
+          $metadata: { httpStatusCode: 400 },
+        }),
+      );
+
+      await expect(
+        (publisher as any).retryOperation(
+          async () => {
+            return s3Mock.send(
+              new ListObjectsV2Command({ Bucket: 'bucketName' }),
+            );
+          },
+          'TestOperation',
+          3,
+        ),
+      ).rejects.toHaveProperty('name', 'BadRequest');
+    });
+  });
+
  describe('getReadiness', () => {
    it('should validate correct config', async () => {
      const publisher = await createPublisherFromConfig();
@@ -16,6 +16,10 @@
 import { Entity, CompoundEntityRef } from '@backstage/catalog-model';
 import { Config } from '@backstage/config';
 import { assertError, ForwardedError } from '@backstage/errors';
+
+// Maximum size in bytes for a single upload part (5MB)
+const MAX_SINGLE_UPLOAD_BYTES = 5 * 1024 * 1024;
+
 import {
  AwsCredentialsManager,
  DefaultAwsCredentialsManager,
@@ -275,67 +279,48 @@ export class AwsS3Publish implements PublisherBase {
    operation: () => Promise<TOutput>,
    operationName: string,
    maxAttempts: number = 3,
+    shouldRetry: (error: S3ServiceException) => boolean = this
+      .defaultShouldRetry,
  ): Promise<TOutput> {
-    let attempts = maxAttempts;
-    let LastError: S3ServiceException;
-
-    while (attempts > 0) {
+    for (let attempt = 1; attempt < maxAttempts; attempt++) {
      try {
        return await operation();
-      } catch (error: unknown) {
-        LastError = error as S3ServiceException;
-        attempts--;
+      } catch (error) {
+        const e = error as S3ServiceException;
+        if (!shouldRetry(e)) {
+          this.logger.error(`${operationName} failed: ${e.message}`);
+          throw e;
+        }

-        const httpStatusCode = LastError.$metadata?.httpStatusCode;
-        const errorCode = LastError.name;
-
-        this.logger.warn(`${operationName} failed.`, {
-          errorCode,
-          httpStatusCode,
-          attemptsRemaining: attempts,
-          currentAttempt: maxAttempts - attempts,
-          totalAttempts: maxAttempts,
-          error: LastError.message,
+        this.logger.warn(`${operationName} failed, retrying...`, {
+          attempt,
+          maxAttempts,
+          error: e.message,
+          errorCode: e.name,
+          httpStatusCode: e.$metadata?.httpStatusCode,
        });
-        // Determine if we should retry based on error type
-        const shouldRetry = this.shouldRetryOperation(LastError, attempts);
-        if (!shouldRetry || attempts === 0) {
-          this.logger.error(
-            `${operationName} failed after all retries: ${LastError.message}`,
-          );
-          throw LastError;
-        }
-        // Enhanced exponential backoff with jitter for upload operation
-        let baseDelay = 1000;
-        if (operationName.startsWith('Upload-')) {
-          // for uploads use longer base delay due to potential multipart commplexity
-          baseDelay = 2000;
-        }
-        const backoffDelay = Math.min(
-          baseDelay * Math.pow(2, maxAttempts - attempts),
-          30000,
-        );
+
+        // Enhanced exponential backoff with jitter
+        const baseDelay = operationName.startsWith('Upload-') ? 2000 : 1000;
+        const backoffDelay = Math.min(baseDelay * Math.pow(2, attempt), 30000);
        const jitter = Math.random() * 1000;
-        const totalDelay = backoffDelay + jitter;
-        await new Promise(resolve => setTimeout(resolve, totalDelay));
+        await new Promise(resolve =>
+          setTimeout(resolve, backoffDelay + jitter),
+        );
      }
    }
-    // Final attempt without retry wrapper
-    return operation();
+    return await operation();
  }

  /**
   * Determines if an S3 operation should be retried based on the error details.
   */
-  private shouldRetryOperation(
-    error: S3ServiceException,
-    attemptsRemaining: number,
-  ): boolean {
+  private defaultShouldRetry(error: S3ServiceException): boolean {
    const httpStatusCode = error.$metadata?.httpStatusCode;
    const errorCode = error.name;
    // Handle invalid part errors first - these are retriable for multipart uploads
    if (errorCode === 'InvalidPart') {
-      return attemptsRemaining > 0;
+      return true;
    }
    // Dont retry for client errors (4xx) except specific cases
    if (httpStatusCode && httpStatusCode >= 400 && httpStatusCode < 500) {
@@ -345,7 +330,7 @@ export class AwsS3Publish implements PublisherBase {
        'RequestTimeoutException',
        'PriorRequestNotComplete',
        'ConnectionError',
-        'RequestTimeToooSkewed',
+        'RequestTimeTooSkewed',
        'InvalidPart',
        'NoSuchUpload',
      ];
@@ -355,7 +340,7 @@ export class AwsS3Publish implements PublisherBase {
    }
    // Always retry for server errors (5xx)
    if (httpStatusCode && httpStatusCode >= 500) {
-      return attemptsRemaining > 0;
+      return true;
    }
    // Retry specific network/connection errors and multipart upload errors
    const retriableErrors = [
@@ -380,12 +365,10 @@ export class AwsS3Publish implements PublisherBase {
      'IncompleteBody',
      'RequestTimeout',
    ];
-    return (
-      retriableErrors.some(
-        retriableError =>
-          errorCode.includes(retriableError) ||
-          error.message.includes(retriableError),
-      ) && attemptsRemaining > 0
+    return retriableErrors.some(
+      retriableError =>
+        errorCode.includes(retriableError) ||
+        error.message.includes(retriableError),
    );
  }

@@ -511,16 +494,12 @@ export class AwsS3Publish implements PublisherBase {
                    new PutObjectCommand(putParams),
                  );
                }
-                // For files 5MB and larger, use multipart upload with enhanced configuration
-                const calaculatedPartSize = Math.max(
-                  fiveMB,
-                  Math.ceil(fileSizeInBytes / 10000),
-                );
+                // For files larger than MAX_SINGLE_UPLOAD_BYTES, use multipart upload
                const upload = new Upload({
                  client: this.storageClient,
                  params,
                  // Configure miltipart upload option for better reliability
-                  partSize: calaculatedPartSize,
+                  partSize: MAX_SINGLE_UPLOAD_BYTES,
                  queueSize: 3,
                  leavePartsOnError: false,
                });
@@ -536,7 +515,7 @@ export class AwsS3Publish implements PublisherBase {

            // Check if this is a multipart upload failure that we can handle
            if (
-              fileSizeInBytes >= 5 * 1024 * 1024 &&
+              fileSizeInBytes >= MAX_SINGLE_UPLOAD_BYTES &&
              (errorName === 'InvalidPart' || errorName === 'NoSuchUpload')
            ) {
              this.logger.warn(