fix: worker

bjsi · Feb 28, 2024 · 39be70c · 39be70c
1 parent d603df4
commit 39be70c
Show file tree

Hide file tree

Showing 5 changed files with 75 additions and 34 deletions.
diff --git a/packages/client/src/components/AdminPage.tsx b/packages/client/src/components/AdminPage.tsx
@@ -42,37 +42,42 @@ export const AdminPage = () => {
     return <div>Loading...</div>;
   }
   return (
-    <TableContainer component={Paper}>
-      <Table aria-label="pipelines table">
-        <TableHead>
-          <TableRow>
-            <TableCell>
-              <span className="font-semibold">Username</span>
-            </TableCell>
-            <TableCell>
-              <span className="font-semibold">Created At</span>
-            </TableCell>
-            <TableCell>
-              <span className="font-semibold">Tasks</span>
-            </TableCell>
-            <TableCell>
-              <span className="font-semibold">Actions</span>
-            </TableCell>
-          </TableRow>
-        </TableHead>
-        <TableBody>
-          {sortBy(pipelines, (x) => dayjs(x.createdAt).valueOf()).map(
-            (pipeline) => (
-              <PipelineRow
-                key={pipeline.id}
-                pipeline={pipeline}
-                setForce={setForce}
-              />
-            )
-          )}
-        </TableBody>
-      </Table>
-    </TableContainer>
+    <>
+      <Button onClick={() => trpcAdmin.unlockWorkers.mutate()}>
+        Force unlock workers
+      </Button>
+      <TableContainer component={Paper}>
+        <Table aria-label="pipelines table">
+          <TableHead>
+            <TableRow>
+              <TableCell>
+                <span className="font-semibold">Username</span>
+              </TableCell>
+              <TableCell>
+                <span className="font-semibold">Created At</span>
+              </TableCell>
+              <TableCell>
+                <span className="font-semibold">Tasks</span>
+              </TableCell>
+              <TableCell>
+                <span className="font-semibold">Actions</span>
+              </TableCell>
+            </TableRow>
+          </TableHead>
+          <TableBody>
+            {sortBy(pipelines, (x) => dayjs(x.createdAt).valueOf()).map(
+              (pipeline) => (
+                <PipelineRow
+                  key={pipeline.id}
+                  pipeline={pipeline}
+                  setForce={setForce}
+                />
+              )
+            )}
+          </TableBody>
+        </Table>
+      </TableContainer>
+    </>
   );
 };
 

diff --git a/packages/server/src/lib/getNumRunningPipelines.ts b/packages/server/src/lib/getNumRunningPipelines.ts
@@ -1,6 +1,7 @@
 import { UserModel } from "shared/src/schemas";
 import { z } from "zod";
 import { prisma } from "../db";
+import { sortBy } from "remeda";
 
 export async function getNumRunningPipelines(
   authenticatedUser: z.infer<typeof UserModel>
@@ -20,7 +21,11 @@ export async function getNumRunningPipelines(
   )
     // TODO: check
     // slice(1) because the first task is the pipeline itself, doesn't get updated properly
-    .filter((p) => p.tasks.slice(1).some((t) => t.status === "running"));
+    .filter((p) =>
+      sortBy(p.tasks, (x) => x.createdAt.valueOf())
+        .slice(1)
+        .some((t) => t.status === "running")
+    );
 
   return pipelines.length;
 }
diff --git a/packages/server/src/lib/jobsPrisma.ts b/packages/server/src/lib/jobsPrisma.ts
@@ -33,10 +33,18 @@ interface Job {
   created_at: Date;
   updated_at: Date;
   key: string | null;
+  locked_by: string | null;
 }
 
 export async function getGraphileJobs(): Promise<any[]> {
   return await prisma.$queryRaw<any[]>(
     Prisma.sql`SELECT * FROM graphile_worker.jobs`
   );
 }
+
+export async function getRunningWorkerIds(): Promise<string[]> {
+  // unique list of locked_by ids
+  return await prisma.$queryRaw<string[]>(
+    Prisma.sql`SELECT DISTINCT locked_by FROM graphile_worker.jobs`
+  );
+}
diff --git a/packages/server/src/routers/adminRouter.ts b/packages/server/src/routers/adminRouter.ts
@@ -12,6 +12,7 @@ import {
   getGraphileJobs,
   getPipelineJobByKey,
   getPipelineTaskJobById,
+  getRunningWorkerIds,
 } from "../lib/jobsPrisma";
 import { indexBy, omit } from "remeda";
 
@@ -206,4 +207,11 @@ export const adminRouter = router({
         },
       });
     }),
+
+  unlockWorkers: publicProcedure.mutation(async () => {
+    const utils = await workerUtils();
+    const workerIds = await getRunningWorkerIds();
+    await utils.forceUnlockWorkers(workerIds);
+    console.log("attempted to unlock workers");
+  }),
 });
diff --git a/packages/server/src/tasks/worker.ts b/packages/server/src/tasks/worker.ts
@@ -1,4 +1,10 @@
-import { Logger, TaskList, parseCronItem, run } from "graphile-worker";
+import {
+  JobHelpers,
+  Logger,
+  TaskList,
+  parseCronItem,
+  run,
+} from "graphile-worker";
 import dotenv from "dotenv";
 import { workerUtils } from "./workerUtils";
 import { prisma } from "../db";
@@ -18,7 +24,7 @@ dotenv.config();
 // if they have < max pipelines, add a pipeline to get recommendations
 export const addPipelineCronJob = {
   id: "recurring-add-pipeline-cron",
-  handler: async () => {
+  handler: async (_: any, helpers: JobHelpers) => {
     const users = shuffle(await prisma.user.findMany());
     for (const user of users) {
       const numActivePipelines = await getNumRunningPipelines(user);
@@ -31,6 +37,14 @@ export const addPipelineCronJob = {
           runId: uuidv4(),
           emailResults: true,
         });
+      } else {
+        helpers.logger.info(
+          "Skipping user " +
+            user.username +
+            " as they have " +
+            numActivePipelines +
+            " active pipelines"
+        );
       }
     }
   },
@@ -100,6 +114,7 @@ export async function startWorker() {
     taskList: taskList as TaskList,
     connectionString: process.env.DATABASE_URL,
   });
+
   const getTaskOrPipeline = async (jobOrKeyId: string) => {
     const pipeline = await prisma.pipelineRun.findFirst({
       where: { jobKeyId: jobOrKeyId },