pass more identifiers to provider and apply labels

nicktrn · nicktrn · commit 9354443ea4dc · 2024-03-28T14:08:31.000Z
diff --git a/apps/kubernetes-provider/src/index.ts b/apps/kubernetes-provider/src/index.ts
@@ -7,7 +7,7 @@ import {
   TaskOperationsIndexOptions,
   TaskOperationsRestoreOptions,
 } from "@trigger.dev/core-apps";
-import { Machine, PostStartCauses, PreStopCauses } from "@trigger.dev/core/v3";
+import { Machine, PostStartCauses, PreStopCauses, EnvironmentType } from "@trigger.dev/core/v3";
 import { randomUUID } from "crypto";
 
 const RUNTIME_ENV = process.env.KUBERNETES_PORT ? "kubernetes" : "local";
@@ -56,6 +56,12 @@ class KubernetesTaskOperations implements TaskOperations {
             metadata: {
               labels: {
                 app: "task-index",
+                "app.kubernetes.io/part-of": "trigger-worker",
+                "app.kubernetes.io/component": "index",
+                env: opts.envId,
+                envtype: this.#envTypeToLabelValue(opts.envType),
+                org: opts.orgId,
+                project: opts.projectId,
               },
             },
             spec: {
@@ -162,6 +168,13 @@ class KubernetesTaskOperations implements TaskOperations {
           namespace: this.#namespace.metadata.name,
           labels: {
             app: "task-run",
+            "app.kubernetes.io/part-of": "trigger-worker",
+            "app.kubernetes.io/component": "create",
+            env: opts.envId,
+            envtype: this.#envTypeToLabelValue(opts.envType),
+            org: opts.orgId,
+            project: opts.projectId,
+            run: opts.runId,
           },
         },
         spec: {
@@ -276,6 +289,14 @@ class KubernetesTaskOperations implements TaskOperations {
           namespace: this.#namespace.metadata.name,
           labels: {
             app: "task-run",
+            "app.kubernetes.io/part-of": "trigger-worker",
+            "app.kubernetes.io/component": "restore",
+            env: opts.envId,
+            envtype: this.#envTypeToLabelValue(opts.envType),
+            org: opts.orgId,
+            project: opts.projectId,
+            run: opts.runId,
+            checkpoint: opts.checkpointId,
           },
         },
         spec: {
@@ -372,6 +393,19 @@ class KubernetesTaskOperations implements TaskOperations {
     await this.#getPod(opts.runId, this.#namespace);
   }
 
+  #envTypeToLabelValue(type: EnvironmentType) {
+    switch (type) {
+      case "PRODUCTION":
+        return "prod";
+      case "STAGING":
+        return "stg";
+      case "DEVELOPMENT":
+        return "dev";
+      case "PREVIEW":
+        return "preview";
+    }
+  }
+
   #getResourcesFromMachineConfig(config: Machine) {
     return {
       cpu: `${config.cpu}`,
diff --git a/apps/webapp/app/v3/marqs/sharedQueueConsumer.server.ts b/apps/webapp/app/v3/marqs/sharedQueueConsumer.server.ts
@@ -449,7 +449,6 @@ export class SharedQueueConsumer {
         const { machineConfig } = taskRunAttempt.backgroundWorkerTask;
         const machine = Machine.safeParse(machineConfig ?? {});
 
-
         if (!machine.success) {
           logger.error("Failed to parse machine config", {
             queueMessage: message.data,
@@ -489,12 +488,16 @@ export class SharedQueueConsumer {
               backgroundWorkerId: deployment.worker.friendlyId,
               data: {
                 type: "SCHEDULE_ATTEMPT",
-                id: taskRunAttempt.id,
                 image: deployment.imageReference,
-                envId: environment.id,
-                runId: taskRunAttempt.taskRunId,
                 version: deployment.version,
                 machine: machine.data,
+                // identifiers
+                id: taskRunAttempt.id,
+                envId: environment.id,
+                envType: environment.type,
+                orgId: environment.organizationId,
+                projectId: environment.projectId,
+                runId: taskRunAttempt.taskRunId,
               },
             });
           }
diff --git a/apps/webapp/app/v3/services/indexDeployment.server.ts b/apps/webapp/app/v3/services/indexDeployment.server.ts
@@ -49,9 +49,13 @@ export class IndexDeploymentService extends BaseService {
         version: "v1",
         shortCode: deployment.shortCode,
         imageTag: deployment.imageReference,
-        envId: deployment.environmentId,
         apiKey: deployment.environment.apiKey,
         apiUrl: env.APP_ORIGIN,
+        // identifiers
+        envId: deployment.environmentId,
+        envType: deployment.environment.type,
+        projectId: deployment.projectId,
+        orgId: deployment.environment.organizationId,
       });
 
       logger.debug("Index ACK received", { responses });
diff --git a/apps/webapp/app/v3/services/restoreCheckpoint.server.ts b/apps/webapp/app/v3/services/restoreCheckpoint.server.ts
@@ -38,6 +38,7 @@ export class RestoreCheckpointService extends BaseService {
                 },
               },
             },
+            runtimeEnvironment: true,
           },
         },
       },
@@ -85,13 +86,18 @@ export class RestoreCheckpointService extends BaseService {
 
     socketIo.providerNamespace.emit("RESTORE", {
       version: "v1",
-      checkpointId: checkpoint.id,
-      runId: checkpoint.runId,
       type: checkpoint.type,
       location: checkpoint.location,
       reason: checkpoint.reason ?? undefined,
       imageRef: checkpoint.imageRef,
       machine: machine.data,
+      // identifiers
+      checkpointId: checkpoint.id,
+      envId: checkpoint.runtimeEnvironment.id,
+      envType: checkpoint.runtimeEnvironment.type,
+      orgId: checkpoint.runtimeEnvironment.organizationId,
+      projectId: checkpoint.runtimeEnvironment.projectId,
+      runId: checkpoint.runId,
     });
 
     return checkpoint;
diff --git a/packages/core-apps/src/provider.ts b/packages/core-apps/src/provider.ts
@@ -2,6 +2,7 @@ import { createServer } from "node:http";
 import {
   ClientToSharedQueueMessages,
   clientWebsocketMessages,
+  EnvironmentType,
   Machine,
   PlatformToProviderMessages,
   ProviderToPlatformMessages,
@@ -25,24 +26,39 @@ const logger = new SimpleLogger(`[${MACHINE_NAME}]`);
 export interface TaskOperationsIndexOptions {
   shortCode: string;
   imageRef: string;
-  envId: string;
   apiKey: string;
   apiUrl: string;
+  // identifiers
+  envId: string;
+  envType: EnvironmentType;
+  orgId: string;
+  projectId: string;
 }
 
 export interface TaskOperationsCreateOptions {
-  runId: string;
   image: string;
   machine: Machine;
-  envId: string;
   version: string;
+  // identifiers
+  envId: string;
+  envType: EnvironmentType;
+  orgId: string;
+  projectId: string;
+  runId: string;
+  attemptId: string;
 }
 
 export interface TaskOperationsRestoreOptions {
-  runId: string;
   imageRef: string;
   checkpointRef: string;
   machine: Machine;
+  // identifiers
+  envId: string;
+  envType: EnvironmentType;
+  orgId: string;
+  projectId: string;
+  runId: string;
+  checkpointId: string;
 }
 
 export interface TaskOperations {
@@ -103,11 +119,16 @@ export class ProviderShell implements Provider {
           if (message.data.type === "SCHEDULE_ATTEMPT") {
             try {
               this.tasks.create({
-                envId: message.data.envId,
-                runId: message.data.runId,
                 image: message.data.image,
                 machine: message.data.machine,
-                version: message.version,
+                version: message.data.version,
+                // identifiers
+                envId: message.data.envId,
+                envType: message.data.envType,
+                orgId: message.data.orgId,
+                projectId: message.data.projectId,
+                runId: message.data.runId,
+                attemptId: message.data.id,
               });
             } catch (error) {
               logger.error("create failed", error);
@@ -168,9 +189,13 @@ export class ProviderShell implements Provider {
             await this.tasks.index({
               shortCode: message.shortCode,
               imageRef: message.imageTag,
-              envId: message.envId,
               apiKey: message.apiKey,
               apiUrl: message.apiUrl,
+              // identifiers
+              envId: message.envId,
+              envType: message.envType,
+              orgId: message.orgId,
+              projectId: message.projectId,
             });
           } catch (error) {
             logger.error("index failed", error);
@@ -209,10 +234,16 @@ export class ProviderShell implements Provider {
 
           try {
             await this.tasks.restore({
-              runId: message.runId,
               checkpointRef: message.location,
               machine: message.machine,
               imageRef: message.imageRef,
+              // identifiers
+              envId: message.envId,
+              envType: message.envType,
+              orgId: message.orgId,
+              projectId: message.projectId,
+              runId: message.runId,
+              checkpointId: message.checkpointId,
             });
           } catch (error) {
             logger.error("restore failed", error);
diff --git a/packages/core/src/v3/schemas/messages.ts b/packages/core/src/v3/schemas/messages.ts
@@ -1,6 +1,9 @@
 import { z } from "zod";
 import { TaskRunExecution, TaskRunExecutionResult } from "./common";
 
+export const EnvironmentType = z.enum(["PRODUCTION", "STAGING", "DEVELOPMENT", "PREVIEW"])
+export type EnvironmentType = z.infer<typeof EnvironmentType>;
+
 export const MachineCpu = z
   .union([z.literal(0.25), z.literal(0.5), z.literal(1), z.literal(2), z.literal(4)])
   .default(0.5);
@@ -59,12 +62,16 @@ export const BackgroundWorkerServerMessages = z.discriminatedUnion("type", [
   }),
   z.object({
     type: z.literal("SCHEDULE_ATTEMPT"),
-    id: z.string(),
     image: z.string(),
-    envId: z.string(),
-    runId: z.string(),
     version: z.string(),
     machine: Machine,
+    // identifiers
+    id: z.string(), // attempt
+    envId: z.string(),
+    envType: EnvironmentType,
+    orgId: z.string(),
+    projectId: z.string(),
+    runId: z.string(),
   }),
 ]);
 
diff --git a/packages/core/src/v3/schemas/schemas.ts b/packages/core/src/v3/schemas/schemas.ts
@@ -7,15 +7,16 @@ import {
   ProdTaskRunExecution,
   ProdTaskRunExecutionPayload,
   RetryOptions,
-  Machine
+  Machine,
+  EnvironmentType,
 } from "./messages";
 import { TaskResource } from "./resources";
 
 export const PostStartCauses = z.enum(["index", "create", "restore"]);
-export type PostStartCauses = z.infer<typeof PostStartCauses>
+export type PostStartCauses = z.infer<typeof PostStartCauses>;
 
 export const PreStopCauses = z.enum(["terminate"]);
-export type PreStopCauses = z.infer<typeof PreStopCauses>
+export type PreStopCauses = z.infer<typeof PreStopCauses>;
 
 const RegexSchema = z.custom<RegExp>((val) => {
   try {
@@ -85,9 +86,13 @@ export const PlatformToProviderMessages = {
       version: z.literal("v1").default("v1"),
       imageTag: z.string(),
       shortCode: z.string(),
-      envId: z.string(),
       apiKey: z.string(),
       apiUrl: z.string(),
+      // identifiers
+      envId: z.string(),
+      envType: EnvironmentType,
+      orgId: z.string(),
+      projectId: z.string(),
     }),
     callback: z.discriminatedUnion("success", [
       z.object({
@@ -107,13 +112,18 @@ export const PlatformToProviderMessages = {
   RESTORE: {
     message: z.object({
       version: z.literal("v1").default("v1"),
-      checkpointId: z.string(),
-      runId: z.string(),
       type: z.enum(["DOCKER", "KUBERNETES"]),
       location: z.string(),
       reason: z.string().optional(),
       imageRef: z.string(),
       machine: Machine,
+      // identifiers
+      checkpointId: z.string(),
+      envId: z.string(),
+      envType: EnvironmentType,
+      orgId: z.string(),
+      projectId: z.string(),
+      runId: z.string(),
     }),
   },
   DELETE: {
