Merge remote-tracking branch 'origin/codex/use-aws-sdk-for-workspace-locking' into cloud-runner-develop

src/model/cloud-runner/services/core/shared-workspace-locking.ts

@@ -1,23 +1,53 @@
-import { CloudRunnerSystem } from './cloud-runner-system';
-import fs from 'node:fs';
 import CloudRunnerLogger from './cloud-runner-logger';
 import BuildParameters from '../../../build-parameters';
 import CloudRunner from '../../cloud-runner';
+import Input from '../../../input';
+import { DeleteObjectCommand, ListObjectsV2Command, PutObjectCommand, S3 } from '@aws-sdk/client-s3';
 export class SharedWorkspaceLocking {
+  private static _s3: S3;
+  private static get s3(): S3 {
+    if (!SharedWorkspaceLocking._s3) {
+      const region = Input.region || process.env.AWS_REGION || process.env.AWS_DEFAULT_REGION || 'us-east-1';
+      SharedWorkspaceLocking._s3 = new S3({ region });
+    }
+    return SharedWorkspaceLocking._s3;
+  }
+  private static get bucket() {
+    return CloudRunner.buildParameters.awsStackName;
+  }
   public static get workspaceBucketRoot() {
-    return `s3://${CloudRunner.buildParameters.awsStackName}/`;
+    return `s3://${SharedWorkspaceLocking.bucket}/`;
   }
   public static get workspaceRoot() {
     return `${SharedWorkspaceLocking.workspaceBucketRoot}locks/`;
   }
+  private static get workspacePrefix() {
+    return `locks/`;
+  }
+  private static async listObjects(prefix: string, bucket = SharedWorkspaceLocking.bucket): Promise<string[]> {
+    if (prefix !== '' && !prefix.endsWith('/')) {
+      prefix += '/';
+    }
+    const result = await SharedWorkspaceLocking.s3.send(
+      new ListObjectsV2Command({ Bucket: bucket, Prefix: prefix, Delimiter: '/' }),
+    );
+    const entries: string[] = [];
+    for (const p of result.CommonPrefixes || []) {
+      if (p.Prefix) entries.push(p.Prefix.slice(prefix.length));
+    }
+    for (const c of result.Contents || []) {
+      if (c.Key && c.Key !== prefix) entries.push(c.Key.slice(prefix.length));
+    }
+    return entries;
+  }
   public static async GetAllWorkspaces(buildParametersContext: BuildParameters): Promise<string[]> {
     if (!(await SharedWorkspaceLocking.DoesCacheKeyTopLevelExist(buildParametersContext))) {
       return [];
     }
 
     return (
-      await SharedWorkspaceLocking.ReadLines(
+      await SharedWorkspaceLocking.listObjects(
-        `aws s3 ls ${SharedWorkspaceLocking.workspaceRoot}${buildParametersContext.cacheKey}/`,
+        `${SharedWorkspaceLocking.workspacePrefix}${buildParametersContext.cacheKey}/`,
       )
     )
       .map((x) => x.replace(`/`, ``))
@@ -26,13 +56,11 @@ export class SharedWorkspaceLocking {
   }
   public static async DoesCacheKeyTopLevelExist(buildParametersContext: BuildParameters) {
     try {
-      const rootLines = await SharedWorkspaceLocking.ReadLines(
+      const rootLines = await SharedWorkspaceLocking.listObjects('');
-        `aws s3 ls ${SharedWorkspaceLocking.workspaceBucketRoot}`,
-      );
       const lockFolderExists = rootLines.map((x) => x.replace(`/`, ``)).includes(`locks`);
 
       if (lockFolderExists) {
-        const lines = await SharedWorkspaceLocking.ReadLines(`aws s3 ls ${SharedWorkspaceLocking.workspaceRoot}`);
+        const lines = await SharedWorkspaceLocking.listObjects(SharedWorkspaceLocking.workspacePrefix);
 
         return lines.map((x) => x.replace(`/`, ``)).includes(buildParametersContext.cacheKey);
       } else {
@@ -55,8 +83,8 @@ export class SharedWorkspaceLocking {
     }
 
     return (
-      await SharedWorkspaceLocking.ReadLines(
+      await SharedWorkspaceLocking.listObjects(
-        `aws s3 ls ${SharedWorkspaceLocking.workspaceRoot}${buildParametersContext.cacheKey}/`,
+        `${SharedWorkspaceLocking.workspacePrefix}${buildParametersContext.cacheKey}/`,
       )
     )
       .map((x) => x.replace(`/`, ``))
@@ -182,8 +210,8 @@ export class SharedWorkspaceLocking {
     }
 
     return (
-      await SharedWorkspaceLocking.ReadLines(
+      await SharedWorkspaceLocking.listObjects(
-        `aws s3 ls ${SharedWorkspaceLocking.workspaceRoot}${buildParametersContext.cacheKey}/`,
+        `${SharedWorkspaceLocking.workspacePrefix}${buildParametersContext.cacheKey}/`,
       )
     )
       .map((x) => x.replace(`/`, ``))
@@ -195,8 +223,8 @@ export class SharedWorkspaceLocking {
     if (!(await SharedWorkspaceLocking.DoesWorkspaceExist(workspace, buildParametersContext))) {
       throw new Error(`workspace doesn't exist ${workspace}`);
     }
-    const files = await SharedWorkspaceLocking.ReadLines(
+    const files = await SharedWorkspaceLocking.listObjects(
-      `aws s3 ls ${SharedWorkspaceLocking.workspaceRoot}${buildParametersContext.cacheKey}/`,
+      `${SharedWorkspaceLocking.workspacePrefix}${buildParametersContext.cacheKey}/`,
     );
 
     const lockFilesExist =
@@ -212,14 +240,10 @@ export class SharedWorkspaceLocking {
       throw new Error(`${workspace} already exists`);
     }
     const timestamp = Date.now();
-    const file = `${timestamp}_${workspace}_workspace`;
+    const key = `${SharedWorkspaceLocking.workspacePrefix}${buildParametersContext.cacheKey}/${timestamp}_${workspace}_workspace`;
-    fs.writeFileSync(file, '');
+    await SharedWorkspaceLocking.s3.send(
-    await CloudRunnerSystem.Run(
+      new PutObjectCommand({ Bucket: SharedWorkspaceLocking.bucket, Key: key, Body: '' }),
-      `aws s3 cp ./${file} ${SharedWorkspaceLocking.workspaceRoot}${buildParametersContext.cacheKey}/${file}`,
-      false,
-      true,
     );
-    fs.rmSync(file);
 
     const workspaces = await SharedWorkspaceLocking.GetAllWorkspaces(buildParametersContext);
 
@@ -241,24 +265,20 @@ export class SharedWorkspaceLocking {
   ): Promise<boolean> {
     const existingWorkspace = workspace.endsWith(`_workspace`);
     const ending = existingWorkspace ? workspace : `${workspace}_workspace`;
-    const file = `${Date.now()}_${runId}_${ending}_lock`;
+    const key = `${SharedWorkspaceLocking.workspacePrefix}${
-    fs.writeFileSync(file, '');
+      buildParametersContext.cacheKey
-    await CloudRunnerSystem.Run(
+    }/${Date.now()}_${runId}_${ending}_lock`;
-      `aws s3 cp ./${file} ${SharedWorkspaceLocking.workspaceRoot}${buildParametersContext.cacheKey}/${file}`,
+    await SharedWorkspaceLocking.s3.send(
-      false,
+      new PutObjectCommand({ Bucket: SharedWorkspaceLocking.bucket, Key: key, Body: '' }),
-      true,
     );
-    fs.rmSync(file);
 
     const hasLock = await SharedWorkspaceLocking.HasWorkspaceLock(workspace, runId, buildParametersContext);
 
     if (hasLock) {
       CloudRunner.lockedWorkspace = workspace;
     } else {
-      await CloudRunnerSystem.Run(
+      await SharedWorkspaceLocking.s3.send(
-        `aws s3 rm ${SharedWorkspaceLocking.workspaceRoot}${buildParametersContext.cacheKey}/${file}`,
+        new DeleteObjectCommand({ Bucket: SharedWorkspaceLocking.bucket, Key: key }),
-        false,
-        true,
       );
     }
 
@@ -275,25 +295,34 @@ export class SharedWorkspaceLocking {
     CloudRunnerLogger.log(`All Locks ${files} ${workspace} ${runId}`);
     CloudRunnerLogger.log(`Deleting lock ${workspace}/${file}`);
     CloudRunnerLogger.log(`rm ${SharedWorkspaceLocking.workspaceRoot}${buildParametersContext.cacheKey}/${file}`);
-    await CloudRunnerSystem.Run(
+    if (file) {
-      `aws s3 rm ${SharedWorkspaceLocking.workspaceRoot}${buildParametersContext.cacheKey}/${file}`,
+      await SharedWorkspaceLocking.s3.send(
-      false,
+        new DeleteObjectCommand({
-      true,
+          Bucket: SharedWorkspaceLocking.bucket,
-    );
+          Key: `${SharedWorkspaceLocking.workspacePrefix}${buildParametersContext.cacheKey}/${file}`,
+        }),
+      );
+    }
 
     return !(await SharedWorkspaceLocking.HasWorkspaceLock(workspace, runId, buildParametersContext));
   }
 
   public static async CleanupWorkspace(workspace: string, buildParametersContext: BuildParameters) {
-    await CloudRunnerSystem.Run(
+    const prefix = `${SharedWorkspaceLocking.workspacePrefix}${buildParametersContext.cacheKey}/`;
-      `aws s3 rm ${SharedWorkspaceLocking.workspaceRoot}${buildParametersContext.cacheKey} --exclude "*" --include "*_${workspace}_*"`,
+    const files = await SharedWorkspaceLocking.listObjects(prefix);
-      false,
+    for (const file of files.filter((x) => x.includes(`_${workspace}_`))) {
-      true,
+      await SharedWorkspaceLocking.s3.send(
-    );
+        new DeleteObjectCommand({ Bucket: SharedWorkspaceLocking.bucket, Key: `${prefix}${file}` }),
+      );
+    }
   }
 
   public static async ReadLines(command: string): Promise<string[]> {
-    return CloudRunnerSystem.RunAndReadLines(command);
+    const path = command.replace('aws s3 ls', '').trim();
+    const withoutScheme = path.replace('s3://', '');
+    const [bucket, ...rest] = withoutScheme.split('/');
+    const prefix = rest.join('/');
+    return SharedWorkspaceLocking.listObjects(prefix, bucket);
   }
 }