perf: use in-memory cache to speed up job status polling

Author: angristan

README.md

@@ -25,18 +25,18 @@ npm start
 
 ### Submission jobs
 
-| Component | Action                                                                                                              |
-| --------- | ------------------------------------------------------------------------------------------------------------------- |
-| `Worker`  | Launch and maintain pool of X prebooted microVMs in the background                                                  |
-| `Worker`  | Bind to RabbitMQ queue `jobs_q` on `jobs_ex` exchange and `jobs_rk` routing key                                     |
-| `API`     | Bind to RabbitMQ queue `jobs_status_q` on `jobs_status_ex` exchange and `jobs_status_rk` routing key                |
-| `Front`   | `POST /submissions`                                                                                                 |
-| `Front`   | poll `GET /submissions/:id` every 500 ms                                                                            |
-| `API`     | Send job to RabbitMQ on `jobs_ex` direct exchange with `jobs_rk` routing key.                                       |
-| `Worker`  | Receive job, get ready microVM from pool, send job to agent in microVM                                              |
-| `Agent`   | Compile/Run code, return result                                                                                     |
-| `Worker`  | During two previous steips, send status of jobs on `jobs_status_ex` exchange with `jobs_status_rk` routing key      |
-| `API`     | Receive each new status on `jobs_status_q`, and insert them into DB, so that the polling can get the latest status. |
+| Component | Action                                                                                                                                  |
+| --------- | --------------------------------------------------------------------------------------------------------------------------------------- |
+| `Worker`  | Launch and maintain pool of X prebooted microVMs in the background                                                                      |
+| `Worker`  | Bind to RabbitMQ queue `jobs_q` on `jobs_ex` exchange and `jobs_rk` routing key                                                         |
+| `API`     | Bind to RabbitMQ queue `jobs_status_q` on `jobs_status_ex` exchange and `jobs_status_rk` routing key                                    |
+| `Front`   | `POST /submissions`                                                                                                                     |
+| `Front`   | Poll `GET /submissions/:id` every 500 ms                                                                                                |
+| `API`     | Send job to RabbitMQ on `jobs_ex` direct exchange with `jobs_rk` routing key.                                                           |
+| `Worker`  | Receive job, get ready microVM from pool, send job to agent in microVM                                                                  |
+| `Agent`   | Compile/Run code, return result                                                                                                         |
+| `Worker`  | During two previous steips, send status of jobs on `jobs_status_ex` exchange with `jobs_status_rk` routing key                          |
+| `API`     | Receive each new status on `jobs_status_q`, and insert them into DB and in-memory cache, so that the polling can get the latest status. |
 
 ## Contribute
 

package-lock.json

@@ -46,6 +46,7 @@
     "amqp-connection-manager": "^3.2.2",
     "amqplib": "^0.8.0",
     "argon2": "^0.27.2",
+    "cache-manager": "^3.4.3",
     "class-transformer": "^0.4.0",
     "class-validator": "^0.13.1",
     "dotenv": "^10.0.0",
@@ -69,6 +70,7 @@
     "@nestjs/schematics": "7.3.1",
     "@nestjs/testing": "7.6.17",
     "@tsconfig/node16": "^1.0.1",
+    "@types/cache-manager": "^3.4.0",
     "@types/jest": "26.0.23",
     "@types/node": "15.6.1",
     "@types/passport-jwt": "^3.0.5",

package.json

@@ -1,6 +1,7 @@
 /* eslint-disable @typescript-eslint/no-unused-vars */
 import { ApiProperty } from '@nestjs/swagger';
 import { User } from 'src/users/user.entity';
+import { jsonMember, jsonObject } from 'typedjson';
 import {
   BaseEntity,
   Column,
@@ -11,6 +12,7 @@ import {
   UpdateDateColumn,
 } from 'typeorm';
 
+@jsonObject
 @Entity('submissions')
 export class Submission extends BaseEntity {
   constructor(partial: Partial<Submission>) {
@@ -19,27 +21,34 @@ export class Submission extends BaseEntity {
   }
 
   @ApiProperty()
+  @jsonMember
   @PrimaryGeneratedColumn('uuid')
   id: string;
 
   @ApiProperty()
+  @jsonMember
   @Column()
   language: string;
 
   @ApiProperty()
+  @jsonMember
   @Column()
   code: string;
 
+  @jsonMember
   @Column()
   status: string;
 
+  @jsonMember
   @Column({ nullable: true })
   output: string;
 
+  @jsonMember
   @ApiProperty()
   @CreateDateColumn()
   createdAt: Date;
 
+  @jsonMember
   @ApiProperty()
   @UpdateDateColumn()
   updatedAt: Date;
@@ -49,5 +58,6 @@ export class Submission extends BaseEntity {
     eager: true,
   })
   @ApiProperty({ type: () => User })
+  @jsonMember(() => User)
   user: User;
 }

src/submissions/submission.entity.ts

@@ -49,8 +49,8 @@ export class SubmissionsController {
   async findOne(
     @Param() findSubmissionDTO: FindSubmissionDTO,
   ): Promise<Submission> {
-    const submission: Submission | undefined =
-      await this.submissionsService.findOne(findSubmissionDTO);
+    // Since this endpoint is used for polling, the service will fetch from cache first and fallback to DB
+    const submission = await this.submissionsService.findOne(findSubmissionDTO);
 
     if (!submission) {
       throw NotFoundException;

src/submissions/submissions.controller.ts

@@ -1,5 +1,5 @@
 import { RabbitMQModule } from '@golevelup/nestjs-rabbitmq';
-import { forwardRef, Module } from '@nestjs/common';
+import { CacheModule, forwardRef, Module } from '@nestjs/common';
 import { TypeOrmModule } from '@nestjs/typeorm';
 import { UsersModule } from 'src/users/users.module';
 import { Submission } from './submission.entity';
@@ -10,6 +10,7 @@ import { SubmissionsService } from './submissions.service';
   imports: [
     TypeOrmModule.forFeature([Submission]),
     forwardRef(() => UsersModule),
+    CacheModule.register(),
     RabbitMQModule.forRoot(RabbitMQModule, {
       exchanges: [
         {

src/submissions/submissions.module.ts

@@ -1,6 +1,7 @@
 import { RabbitSubscribe } from '@golevelup/nestjs-rabbitmq';
-import { Injectable } from '@nestjs/common';
+import { CACHE_MANAGER, Inject, Injectable } from '@nestjs/common';
 import { InjectRepository } from '@nestjs/typeorm';
+import { Cache } from 'cache-manager';
 import { TypedJSON } from 'typedjson';
 import { Repository } from 'typeorm';
 import { FindSubmissionDTO } from './dto/find-submission.dto';
@@ -13,6 +14,7 @@ export class SubmissionsService {
   constructor(
     @InjectRepository(Submission)
     private submissionsRepository: Repository<Submission>,
+    @Inject(CACHE_MANAGER) private cacheManager: Cache,
   ) {}
 
   async create(insertSubmissionDTO: InsertSubmissionDTO): Promise<Submission> {
@@ -22,7 +24,11 @@ export class SubmissionsService {
     return submission.save();
   }
 
-  async setStatus(id: string, status: string, output?: string): Promise<void> {
+  async setStatus(
+    id: string,
+    status: string,
+    output?: string,
+  ): Promise<Submission | undefined> {
     const submission = await this.submissionsRepository.findOne({
       id,
     });
@@ -32,13 +38,28 @@ export class SubmissionsService {
         submission.output = output;
       }
       await submission.save();
+      return submission;
     }
+    return undefined;
   }
 
   async findOne(
-    submission: FindSubmissionDTO,
+    queriedSubmission: FindSubmissionDTO,
   ): Promise<Submission | undefined> {
-    return this.submissionsRepository.findOne({ id: submission.id });
+    // First, try to get from cache
+    const cachedSubmission = await this.cacheManager.get(queriedSubmission.id);
+
+    if (cachedSubmission !== '') {
+      const serializer = new TypedJSON(Submission);
+
+      const submission = serializer.parse(cachedSubmission);
+      if (submission) {
+        return submission;
+      }
+    }
+
+    // Fallback to DB
+    return this.submissionsRepository.findOne({ id: queriedSubmission.id });
   }
 
   @RabbitSubscribe({
@@ -50,11 +71,27 @@ export class SubmissionsService {
     // TODO: use logger instead
     console.log(`Received job status: ${JSON.stringify(msg)}`);
 
-    const serializer = new TypedJSON(JobStatusDTO);
+    const jobSerializer = new TypedJSON(JobStatusDTO);
+    const jobStatus = jobSerializer.parse(msg);
 
-    const jobStatus = serializer.parse(msg);
     if (jobStatus) {
-      await this.setStatus(jobStatus.id, jobStatus.status, jobStatus.stdout);
+      // Set in DB
+      const submission = await this.setStatus(
+        jobStatus.id,
+        jobStatus.status,
+        jobStatus.stdout,
+      );
+
+      if (submission) {
+        const submissionSerializer = new TypedJSON(Submission);
+
+        // Set in cache to speed up polling
+        await this.cacheManager.set(
+          jobStatus.id,
+          submissionSerializer.stringify(submission),
+          { ttl: 600 },
+        );
+      }
     }
   }
 }