DOP-2565: Create cron job to fail stuck jobs (#766)

Author: rayangler

api/controllers/v1/jobs.ts

@@ -62,12 +62,22 @@ export const HandleJobs = async (event: any = {}): Promise<any> => {
             // start the task , don't start the process before processing the notification
             const ecsServices = new ECSContainer(c, consoleLogger);
             const res = await ecsServices.execute(jobId);
+            if (res) {
+              await saveTaskId(jobId, res, consoleLogger);
+            }
             consoleLogger.info(jobId, JSON.stringify(res));
             break;
           case JobStatus[JobStatus.inProgress]:
             queueUrl = c.get('jobUpdatesQueueUrl');
             await NotifyBuildProgress(jobId);
             break;
+          case JobStatus[JobStatus.timedOut]:
+            await NotifyBuildSummary(jobId);
+            const taskId = body['taskId'];
+            if (taskId) {
+              await stopECSTask(taskId, consoleLogger);
+            }
+            break;
           case JobStatus[JobStatus.failed]:
           case JobStatus[JobStatus.completed]:
             queueUrl = c.get('jobUpdatesQueueUrl');
@@ -86,6 +96,44 @@ export const HandleJobs = async (event: any = {}): Promise<any> => {
   );
 };
 
+export const FailStuckJobs = async () => {
+  const client = new mongodb.MongoClient(c.get('dbUrl'));
+  await client.connect();
+  const db = client.db(c.get('dbName'));
+  const consoleLogger = new ConsoleLogger();
+  const jobRepository = new JobRepository(db, c, consoleLogger);
+
+  try {
+    const hours = 8;
+    await jobRepository.failStuckJobs(hours);
+  } catch (err) {
+    consoleLogger.error('FailStuckJobs', err);
+  }
+};
+
+async function saveTaskId(jobId: string, taskExecutionRes: any, consoleLogger: ConsoleLogger): Promise<void> {
+  const taskArn = taskExecutionRes?.tasks[0]?.taskArn;
+  if (!taskArn) return;
+
+  const client = new mongodb.MongoClient(c.get('dbUrl'));
+  await client.connect();
+  const db = client.db(c.get('dbName'));
+  const jobRepository = new JobRepository(db, c, consoleLogger);
+
+  try {
+    // Only interested in the actual task ID since the whole ARN might have sensitive information
+    const taskId = taskArn.split('/').pop();
+    await jobRepository.addTaskIdToJob(jobId, taskId);
+  } catch (err) {
+    consoleLogger.error('saveTaskId', err);
+  }
+}
+
+async function stopECSTask(taskId: string, consoleLogger: ConsoleLogger) {
+  const ecs = new ECSContainer(c, consoleLogger);
+  await ecs.stopZombieECSTask(taskId);
+}
+
 async function retry(message: JobQueueMessage, consoleLogger: ConsoleLogger, url: string): Promise<any> {
   try {
     const tries = message.tries;
@@ -113,10 +161,6 @@ async function NotifyBuildSummary(jobId: string): Promise<any> {
   const jobRepository = new JobRepository(db, c, consoleLogger);
   // TODO: Make fullDocument be of type Job, validate existence
   const fullDocument = await jobRepository.getJobById(jobId);
-  // TODO: Remove unused vars, and validate existing vars
-  const branchesRepo = new BranchRepository(db, c, consoleLogger);
-  const slackMsgs = fullDocument.comMessage;
-  const jobTitle = fullDocument.title;
   const repoName = fullDocument.payload.repoName;
   const username = fullDocument.user;
   const slackConnector = new SlackConnector(consoleLogger, c);
@@ -126,7 +170,7 @@ async function NotifyBuildSummary(jobId: string): Promise<any> {
     consoleLogger.error(username, 'Entitlement failed');
     return;
   }
-  const resp = await slackConnector.sendMessage(
+  await slackConnector.sendMessage(
     await prepSummaryMessage(
       env,
       fullDocument,

serverless.yml

@@ -214,8 +214,11 @@ functions:
       - Effect: Allow
         Action:
           - "ecs:RunTask"
+          - "ecs:StopTask"
+          - "ecs:DescribeTasks"
         Resource:
           - arn:aws:ecs:${aws:region}:${aws:accountId}:task-definition/docs-worker-pool-${self:provider.stage}:*
+          - arn:aws:ecs:${aws:region}:${aws:accountId}:task/docs-worker-pool-${self:provider.stage}/*
       - Effect: Allow
         Action:
           - ecs:DescribeTaskDefinition
@@ -238,6 +241,13 @@ functions:
       environment:
         <<: *webhook-env-core
 
+    v1FailStuckJobs:
+      handler: api/controllers/v1/jobs.FailStuckJobs
+      events:
+        - schedule: rate(8 hours)
+      environment:
+        <<: *webhook-env-core
+
 Outputs:
   JobsQueueURL:
     Description: Jobs Queue Url

src/entities/job.ts

@@ -6,6 +6,7 @@ export enum JobStatus {
   inProgress = 'inProgress',
   completed = 'completed',
   failed = 'failed',
+  timedOut = 'timedOut',
 }
 
 // TODO: Formalize JobTypes

src/entities/queueMessage.ts

@@ -1,13 +1,16 @@
-import { timingSafeEqual } from 'crypto';
 import { JobStatus } from './job';
 
 export class JobQueueMessage {
   jobId: string;
   jobStatus: JobStatus;
   tries: number;
-  constructor(jobId: string, status: JobStatus, tries = 0) {
+  taskId?: string;
+
+  constructor(jobId: string, status: JobStatus, tries = 0, taskId?: string) {
     this.jobId = jobId;
     this.jobStatus = status;
     this.tries = tries;
+
+    if (taskId) this.taskId = taskId;
   }
 }

src/repositories/baseRepository.ts

@@ -79,6 +79,40 @@ export abstract class BaseRepository {
       throw error;
     }
   }
+
+  protected async find(query: any, errorMsg: string, options?: mongodb.FindOptions): Promise<mongodb.FindCursor> {
+    try {
+      return this.promiseTimeoutS(this._config.get('MONGO_TIMEOUT_S'), this._collection.find(query, options), errorMsg);
+    } catch (error) {
+      this._logger.error(`${this._repoName}:find`, `Failed to find (${JSON.stringify(query)}) error: ${error}`);
+      throw error;
+    }
+  }
+
+  protected async updateMany(query: any, update: any, errorMsg: string): Promise<boolean> {
+    try {
+      const updateResult = await this.promiseTimeoutS(
+        this._config.get('MONGO_TIMEOUT_S'),
+        this._collection.updateMany(query, update),
+        errorMsg
+      );
+      // If no documents were found, this is not necessarily an error.
+      if ((updateResult?.matchedCount ?? 0) < 1) {
+        return false;
+      }
+      if ((updateResult?.modifiedCount ?? 0) < 1) {
+        throw new DBError(`Failed to modify jobs with query ${JSON.stringify(query)}`);
+      }
+    } catch (error) {
+      this._logger.error(
+        `${this._repoName}:updateMany`,
+        `Failed to update many with (${JSON.stringify(query)}) error: ${error}`
+      );
+      throw error;
+    }
+    return true;
+  }
+
   protected async updateOne(query: any, update: any, errorMsg: string): Promise<boolean> {
     try {
       const updateResult = await this.update(query, update, errorMsg);

src/repositories/jobRepository.ts

@@ -89,8 +89,8 @@ export class JobRepository extends BaseRepository {
     return await this.findOneAndUpdateJob(query);
   }
 
-  async notify(jobId: string, url: string, status: JobStatus, delay: number) {
-    await this._queueConnector.sendMessage(new JobQueueMessage(jobId, status), url, delay);
+  async notify(jobId: string, url: string, status: JobStatus, delay: number, taskId?: string) {
+    await this._queueConnector.sendMessage(new JobQueueMessage(jobId, status, 0, taskId), url, delay);
   }
 
   async findOneAndUpdateJob(query): Promise<Job | null> {
@@ -205,4 +205,63 @@ export class JobRepository extends BaseRepository {
     }
     return bRet;
   }
+
+  async failStuckJobs(hours: number) {
+    const hourInMS = 1000 * 60 * 60;
+    const currentTime = new Date();
+    const failReason = `Job timeout error: Job has been running for at least ${hours} hours`;
+
+    const query = {
+      status: JobStatus.inProgress,
+      startTime: {
+        $lte: new Date(currentTime.getTime() - hourInMS * hours),
+      },
+    };
+    const findOptions = {
+      projection: {
+        _id: 1,
+        taskId: 1,
+      },
+    };
+    const update = {
+      $set: {
+        status: JobStatus.timedOut,
+        endTime: currentTime,
+        error: { time: currentTime.toString(), reason: failReason },
+      },
+    };
+
+    // Mongo's updateMany does not return the IDs of documents changed, so we find them first
+    const stuckJobsCursor = await this.find(query, `Mongo Timeout Error: Timed out finding stuck jobs.`, findOptions);
+    const stuckJobs = await stuckJobsCursor.toArray();
+    this._logger.info('failStuckJobs', `Found ${stuckJobs.length} jobs.`);
+    // No stuck jobs found
+    if (!stuckJobs.length) return;
+
+    // Since the same query is split into 2 operations, there's a very small chance
+    // that there might be a mismatch between jobs updated in the db vs. in the queue.
+    const bRet = await this.updateMany(query, update, `Mongo Timeout Error: Timed out updating stuck jobs.`);
+    if (!bRet) {
+      throw new DBError('failStuckJobs: Unable to update stuck jobs.');
+    }
+
+    const jobUpdatesQueueUrl: string = this._config.get('jobUpdatesQueueUrl');
+    await Promise.all(
+      stuckJobs.map((stuckJob: any) => {
+        const id: string = stuckJob._id.toString();
+        return this.notify(id, jobUpdatesQueueUrl, JobStatus.timedOut, 0, stuckJob.taskId);
+      })
+    );
+  }
+
+  async addTaskIdToJob(id: string, taskId: string): Promise<void> {
+    const query = { _id: new objectId(id) };
+    const update = {
+      $set: {
+        taskId: taskId,
+      },
+    };
+
+    await this.updateOne(query, update, `Mongo Timeout Error: Timed out while updating taskId for jobId: ${id}`);
+  }
 }

src/services/containerServices.ts

@@ -87,4 +87,24 @@ export class ECSContainer implements IContainerServices {
     }
     return '';
   }
+
+  async stopZombieECSTask(taskId: string) {
+    const clusterName = this._config.get<string>('taskDefinitionFamily');
+    try {
+      const res = await this._client.describeTasks({
+        cluster: clusterName,
+        tasks: [taskId],
+      });
+
+      // Only stop ECS task if it's still running
+      if (res.tasks?.[0].lastStatus === 'RUNNING') {
+        await this._client.stopTask({
+          cluster: clusterName,
+          task: taskId,
+        });
+      }
+    } catch (error) {
+      throw error;
+    }
+  }
 }

tests/unit/repositories/jobRepository.test.ts

@@ -2,6 +2,7 @@ import { JobRepository } from '../../../src/repositories/jobRepository';
 import { getBuildJobPlain } from '../../data/jobDef';
 import { DBRepositoryHelper } from '../../utils/repositoryHelper';
 import { TestDataProvider } from '../../data/data';
+import { ObjectId } from 'mongodb';
 
 describe('Job Repository Tests', () => {
   let jobRepo: JobRepository;
@@ -179,6 +180,14 @@ describe('Job Repository Tests', () => {
     });
   });
 
+  describe('failStuckJobs Tests', () => {
+    test('failStuckJobs succeeds', async () => {
+      setupForFindSuccess();
+      setupForUpdateManySuccess();
+      await expect(jobRepo.failStuckJobs(8)).resolves.toEqual(undefined);
+    });
+  });
+
   function validateSuccessfulUpdate(testData: any) {
     expect(dbRepoHelper.collection.updateOne).toBeCalledTimes(1);
     expect(dbRepoHelper.collection.updateOne).toBeCalledWith(testData.query, testData.update);
@@ -190,4 +199,21 @@ describe('Job Repository Tests', () => {
     jest.spyOn(jobRepo, 'notify').mockResolvedValueOnce(true);
     dbRepoHelper.config.get.calledWith('MONGO_TIMEOUT_S').mockReturnValueOnce(1);
   }
+
+  function setupForFindSuccess() {
+    dbRepoHelper.collection.find.mockReturnValueOnce({
+      toArray: () => [
+        { _id: new ObjectId(), status: 'inProgress' },
+        { _id: new ObjectId(), status: 'inQueue' },
+        { _id: new ObjectId(), status: 'inProgress' },
+      ],
+    });
+    dbRepoHelper.config.get.calledWith('MONGO_TIMEOUT_S').mockReturnValueOnce(1);
+  }
+
+  function setupForUpdateManySuccess() {
+    dbRepoHelper.collection.updateMany.mockReturnValueOnce({ matchedCount: 2, modifiedCount: 2 });
+    jest.spyOn(jobRepo, 'notify').mockResolvedValue(true);
+    dbRepoHelper.config.get.calledWith('MONGO_TIMEOUT_S').mockReturnValueOnce(1);
+  }
 });

tests/utils/repositoryHelper.ts

@@ -1,6 +1,6 @@
 import { IConfig } from 'config';
 import { mockDeep } from 'jest-mock-extended';
-import { Db } from 'mongodb';
+import { Db, FindCursor, FindOptions } from 'mongodb';
 import { JobRepository } from '../../src/repositories/jobRepository';
 import { RepoEntitlementsRepository } from '../../src/repositories/repoEntitlementsRepository';
 import { RepoBranchesRepository } from '../../src/repositories/repoBranchesRepository';
@@ -11,6 +11,8 @@ export class DBRepositoryHelper {
   logger: ILogger;
   db: Db;
   updateOne: (query: any, update: any, errorMsg: string) => Promise<boolean>;
+  updateMany: (query: any, update: any, errorMsg: string) => Promise<boolean>;
+  find: (query: any, errorMsg: string, options?: FindOptions) => Promise<FindCursor>;
   findOne: (query: any, errorMsg: string) => Promise<any>;
   findOneAndUpdate: (query: any, update: any, options: any, errorMsg: string) => Promise<any>;
   collection: any;
@@ -29,10 +31,14 @@ export class DBRepositoryHelper {
     this.logger = mockDeep<ILogger>();
     this.db = mockDeep<Db>();
     this.updateOne = jest.fn();
+    this.updateMany = jest.fn();
+    this.find = jest.fn();
     this.findOne = jest.fn();
     this.findOneAndUpdate = jest.fn();
     this.collection = {
       updateOne: this.updateOne,
+      updateMany: this.updateMany,
+      find: this.find,
       findOne: this.findOne,
       findOneAndUpdate: this.findOneAndUpdate,
     };