Graceful shutdown: wait for session to be released (#35)

* Graceful shutdown: wait for session to be released * Disable puppeteer signal handling, remove hacks around it * Reset session info on release as it is implicitly restarted * Re-introduce signal handler removal for local dev env `npm run dev` will install its own signal handlers and break things
steel-dev · Jan 10, 2025 · 5a8dff7 · 5a8dff7
1 parent 62af6d7
commit 5a8dff7
Show file tree

Hide file tree

Showing 14 changed files with 240 additions and 120 deletions.
diff --git a/api/Dockerfile b/api/Dockerfile
@@ -1,4 +1,4 @@
-ARG NODE_VERSION=20.12.0
+ARG NODE_VERSION=22.13.0
 
 FROM node:${NODE_VERSION}-slim AS base
 

diff --git a/api/entrypoint.sh b/api/entrypoint.sh
@@ -42,5 +42,8 @@ export DISPLAY=:10
 export CDP_REDIRECT_PORT=9223
 export HOST=0.0.0.0
 
-# Run the npm start command
-exec npm run start
+# Run the `npm run start` command but without npm.
+# NPM will introduce its own signal handling
+# which will prevent the container from waiting
+# for a session to be released before stopping gracefully
+exec node ./build/index.js
diff --git a/api/package-lock.json b/api/package-lock.json
diff --git a/api/package.json b/api/package.json
@@ -32,6 +32,7 @@
     "@scalar/fastify-api-reference": "^1.25.59",
     "dotenv": "^16.4.5",
     "fastify": "^4.15.0",
+    "fastify-graceful-shutdown": "^4.0.1",
     "fastify-plugin": "^4.5.1",
     "fastify-zod": "^1.4.0",
     "fingerprint-generator": "^2.1.56",
@@ -51,4 +52,4 @@
     "zod": "^3.22.4",
     "zod-to-json-schema": "^3.21.4"
   }
-}
+}
diff --git a/api/src/build-server.ts b/api/src/build-server.ts
@@ -4,6 +4,7 @@ import fastifyCors from "@fastify/cors";
 import openAPIPlugin from "./plugins/schemas";
 import requestLogger from "./plugins/request-logger";
 import browserInstancePlugin from "./plugins/browser";
+import browserSessionPlugin from "./plugins/browser-session";
 import browserWebSocket from "./plugins/browser-socket";
 import seleniumPlugin from "./plugins/selenium";
 import customBodyParser from "./plugins/custom-body-parser";
@@ -21,6 +22,8 @@ export default function buildFastifyServer(options?: FastifyServerOptions) {
   server.register(seleniumPlugin);
   server.register(browserWebSocket);
   server.register(customBodyParser);
+  server.register(browserSessionPlugin);
+
   // Routes
   server.register(actionsRoutes, { prefix: "/v1" });
   server.register(sessionsRoutes, { prefix: "/v1" });

diff --git a/api/src/env.ts b/api/src/env.ts
@@ -14,6 +14,7 @@ const envSchema = z.object({
     .optional()
     .transform((val) => (val ? JSON.parse(val) : {}))
     .pipe(z.record(z.string()).optional().default({})),
+  KILL_TIMEOUT: z.string().optional().default("25"), // to fit in default 30 seconds of Heroku or ECS with some margin
 });
 
 export const env = envSchema.parse(process.env);
diff --git a/api/src/modules/sessions/sessions.controller.ts b/api/src/modules/sessions/sessions.controller.ts
@@ -2,28 +2,8 @@ import { CDPService } from "../../services/cdp.service";
 import { FastifyInstance, FastifyReply, FastifyRequest } from "fastify";
 import { getErrors } from "../../utils/errors";
 import { CreateSessionRequest, SessionDetails } from "./sessions.schema";
-import { BrowserLauncherOptions } from "../../types";
-import { SeleniumService } from "../../services/selenium.service";
-import { v4 as uuidv4 } from "uuid";
 import { env } from "../../env";
 
-let activeSession: SessionDetails = {
-  id: uuidv4(),
-  createdAt: new Date().toISOString(),
-  status: "live",
-  duration: 0,
-  eventCount: 0,
-  timeout: 0,
-  creditsUsed: 0,
-  websocketUrl: `ws://${env.DOMAIN ?? env.HOST}:${env.PORT}/`,
-  debugUrl: `http://${env.DOMAIN ?? env.HOST}:${env.PORT}/v1/devtools/inspector.html`,
-  sessionViewerUrl: `http://${env.DOMAIN ?? env.HOST}:${env.PORT}`,
-  userAgent: "",
-  isSelenium: false,
-  proxy: "",
-  solveCaptcha: false,
-};
-
 export const handleLaunchBrowserSession = async (
   server: FastifyInstance,
   request: CreateSessionRequest,
@@ -44,104 +24,26 @@ export const handleLaunchBrowserSession = async (
     } = request.body;
 
     // If there's an active session, close it first
-    if (activeSession) {
-      if (activeSession.isSelenium) {
-        server.seleniumService.close();
-      } else {
-        await server.cdpService.shutdown();
-      }
-    }
+    await server.sessionService.endSession();
 
-    const browserLauncherOptions: BrowserLauncherOptions = {
-      options: {
-        headless: true,
-        args: [userAgent ? `--user-agent=${userAgent}` : undefined].filter(Boolean) as string[],
-        proxyUrl,
-      },
-      cookies: sessionContext?.cookies || [],
-      userAgent: sessionContext?.userAgent,
-      blockAds,
-      extensions: extensions || [],
-      logSinkUrl,
-      timezone: timezone || "US/Pacific",
-      dimensions,
-    };
+    return await server.sessionService.startSession({
+      sessionId, proxyUrl, userAgent, sessionContext, extensions, logSinkUrl, timezone, dimensions, isSelenium, blockAds,
+    });
 
-    if (isSelenium) {
-      await server.cdpService.shutdown();
-      await server.seleniumService.launch(browserLauncherOptions);
-      const sessionDetails: SessionDetails = {
-        id: sessionId || uuidv4(),
-        createdAt: new Date().toISOString(),
-        status: "live",
-        duration: 0,
-        eventCount: 0,
-        timeout: 0,
-        creditsUsed: 0,
-        websocketUrl: "",
-        debugUrl: "",
-        sessionViewerUrl: "",
-        userAgent:
-          sessionContext?.userAgent ||
-          "Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/128.0.0.0 Safari/537.36",
-        proxy: proxyUrl,
-        solveCaptcha: false,
-        isSelenium,
-      };
-      activeSession = sessionDetails;
-      reply.send(sessionDetails);
-    } else {
-      await server.cdpService.startNewSession(browserLauncherOptions);
-      const sessionDetails: SessionDetails = {
-        id: sessionId || uuidv4(),
-        createdAt: new Date().toISOString(),
-        status: "live",
-        duration: 0,
-        eventCount: 0,
-        timeout: 0,
-        creditsUsed: 0,
-        websocketUrl: `ws://${env.DOMAIN ?? env.HOST}:${env.PORT}/`,
-        debugUrl: `http://${env.DOMAIN ?? env.HOST}:${env.PORT}/v1/devtools/inspector.html`,
-        sessionViewerUrl: `http://${env.DOMAIN ?? env.HOST}:${env.PORT}`,
-        userAgent: server.cdpService.getUserAgent(),
-        proxy: proxyUrl,
-        solveCaptcha: false,
-        isSelenium,
-      };
-      activeSession = sessionDetails;
-      return reply.send(sessionDetails);
-    }
   } catch (e: unknown) {
     const error = getErrors(e);
     return reply.code(500).send({ success: false, message: error });
   }
 };
 
 export const handleExitBrowserSession = async (
-  browserService: CDPService,
-  seleniumService: SeleniumService,
+  server: FastifyInstance,
   request: FastifyRequest,
   reply: FastifyReply,
 ) => {
   try {
-    seleniumService.close();
-    await browserService.endSession();
-    activeSession = {
-      id: uuidv4(),
-      createdAt: new Date().toISOString(),
-      status: "live",
-      duration: 0,
-      eventCount: 0,
-      timeout: 0,
-      creditsUsed: 0,
-      websocketUrl: `ws://${env.DOMAIN ?? env.HOST}/`,
-      debugUrl: `http://${env.DOMAIN ?? env.HOST}/debug`,
-      sessionViewerUrl: `http://${env.DOMAIN ?? env.HOST}`,
-      userAgent: "",
-      isSelenium: false,
-      proxy: "",
-      solveCaptcha: false,
-    };
+    await server.sessionService.endSession();
+
     reply.send({ success: true });
   } catch (e: unknown) {
     const error = getErrors(e);
@@ -159,11 +61,12 @@ export const handleGetBrowserContext = async (
 };
 
 export const handleGetSessionDetails = async (
+  server: FastifyInstance,
   request: FastifyRequest<{ Params: { sessionId: string } }>,
   reply: FastifyReply,
 ) => {
   const sessionId = request.params.sessionId;
-  if (sessionId !== activeSession.id) {
+  if (sessionId !== server.sessionService.activeSession.id) {
     return reply.send({
       id: sessionId,
       createdAt: new Date().toISOString(),
@@ -181,9 +84,9 @@ export const handleGetSessionDetails = async (
       solveCaptcha: false,
     });
   }
-  return reply.send(activeSession);
+  return reply.send(server.sessionService.activeSession);
 };
 
-export const handleGetSessions = async (request: FastifyRequest, reply: FastifyReply) => {
-  return reply.send([activeSession]);
+export const handleGetSessions = async (server: FastifyInstance, request: FastifyRequest, reply: FastifyReply) => {
+  return reply.send([server.sessionService.activeSession]);
 };
diff --git a/api/src/modules/sessions/sessions.routes.ts b/api/src/modules/sessions/sessions.routes.ts
@@ -58,7 +58,7 @@ async function routes(server: FastifyInstance) {
         },
       },
     },
-    async (request: FastifyRequest, reply: FastifyReply) => handleGetSessions(request, reply),
+    async (request: FastifyRequest, reply: FastifyReply) => handleGetSessions(server, request, reply),
   );
 
   server.get(
@@ -75,7 +75,7 @@ async function routes(server: FastifyInstance) {
       },
     },
     async (request: FastifyRequest<{ Params: { sessionId: string } }>, reply: FastifyReply) =>
-      handleGetSessionDetails(request, reply),
+      handleGetSessionDetails(server, request, reply),
   );
 
   server.get(
@@ -102,7 +102,7 @@ async function routes(server: FastifyInstance) {
       },
     },
     async (request: FastifyRequest, reply: FastifyReply) =>
-      handleExitBrowserSession(server.cdpService, server.seleniumService, request, reply),
+      handleExitBrowserSession(server, request, reply),
   );
 
   server.post(
@@ -116,7 +116,7 @@ async function routes(server: FastifyInstance) {
       },
     },
     async (request: FastifyRequest, reply: FastifyReply) =>
-      handleExitBrowserSession(server.cdpService, server.seleniumService, request, reply),
+      handleExitBrowserSession(server, request, reply),
   );
 
   server.post(

diff --git a/api/src/modules/sessions/sessions.schema.ts b/api/src/modules/sessions/sessions.schema.ts
@@ -24,7 +24,7 @@ const CreateSession = z.object({
 const SessionDetails = z.object({
   id: z.string().uuid().describe("Unique identifier for the session"),
   createdAt: z.string().datetime().describe("Timestamp when the session started"),
-  status: z.enum(["live", "released", "failed"]).describe("Status of the session"),
+  status: z.enum(["pending", "live", "released", "failed"]).describe("Status of the session"),
   duration: z.number().int().describe("Duration of the session in milliseconds"),
   eventCount: z.number().int().describe("Number of events processed in the session"),
   timeout: z.number().int().describe("Session timeout duration in milliseconds"),
@@ -49,7 +49,10 @@ const MultipleSessions = z.array(SessionDetails);
 export type CreateSessionBody = z.infer<typeof CreateSession>;
 export type CreateSessionRequest = FastifyRequest<{ Body: CreateSessionBody }>;
 
-export type SessionDetails = z.infer<typeof SessionDetails>;
+export type SessionDetails = z.infer<typeof SessionDetails> & {
+  completion: Promise<void>,
+  complete: (value: void) => void,
+};
 export type MultipleSessions = z.infer<typeof MultipleSessions>;
 export const browserSchemas = {
   CreateSession,

diff --git a/api/src/plugins/browser-session.ts b/api/src/plugins/browser-session.ts
@@ -0,0 +1,36 @@
+import fastifyGracefulShutdown from "fastify-graceful-shutdown";
+import { FastifyPluginAsync } from "fastify";
+import fp from "fastify-plugin";
+import { SessionService } from "../services/session.service";
+import { env } from "../env";
+
+const browserSessionPlugin: FastifyPluginAsync = async (fastify, options) => {
+  const sessionService = new SessionService({
+    cdpService: fastify.cdpService,
+    seleniumService: fastify.seleniumService,
+    logger: fastify.log,
+  });
+  fastify.decorate("sessionService", sessionService);
+
+  process.removeAllListeners("SIGINT");
+  process.removeAllListeners("SIGTERM");
+
+  const gracefulOptions = {
+    timeout: parseInt(env.KILL_TIMEOUT) * 1000,
+  };
+  fastify.register(fastifyGracefulShutdown, gracefulOptions).after((err) => {
+    if (err) {
+      fastify.log.error(err)
+    }
+
+    fastify.gracefulShutdown(async (_signal) => {
+      if (sessionService.activeSession.status === "live") {
+        fastify.log.info("Waiting for active session to be released...");
+        await sessionService.activeSession.completion;
+        fastify.log.info("Active session has been released...");
+      }      
+    });
+  });
+};
+
+export default fp(browserSessionPlugin, "4.x");
diff --git a/api/src/services/cdp.service.ts b/api/src/services/cdp.service.ts
@@ -393,6 +393,8 @@ export class CDPService extends EventEmitter {
       args: launchArgs,
       executablePath: this.chromeExecPath,
       timeout: 0,
+      handleSIGINT: false,
+      handleSIGTERM: false,
       // dumpio: true, //uncomment this line to see logs from chromium
     };