v3: Cancel dev runs when the dev CLI exits (#916)

* Cancel dev runs when the dev CLI exits

* Clear out in progress attempts after cancelling them
This commit is contained in:
Eric Allam
2024-02-28 16:00:39 +00:00
committed by GitHub
parent 479445f741
commit b6517af522
19 changed files with 905 additions and 354 deletions
@@ -1,9 +1,13 @@
import {
isExceptionSpanEvent,
type ExceptionEventProperties,
type SpanEvent as OtelSpanEvent,
} from "@trigger.dev/core/v3";
import { CodeBlock } from "~/components/code/CodeBlock";
import { Callout } from "~/components/primitives/Callout";
import { DateTime, DateTimeAccurate } from "~/components/primitives/DateTime";
import { Header2, Header3 } from "~/components/primitives/Headers";
import { DateTimeAccurate } from "~/components/primitives/DateTime";
import { Header2 } from "~/components/primitives/Headers";
import { Paragraph } from "~/components/primitives/Paragraph";
import type { OtelExceptionProperty, OtelSpanEvent } from "~/presenters/v3/SpanPresenter.server";
type SpanEventsProps = {
spanEvents: OtelSpanEvent[];
@@ -39,7 +43,7 @@ function SpanEventHeader({
}
function SpanEvent({ spanEvent }: { spanEvent: OtelSpanEvent }) {
if (spanEvent.properties?.exception) {
if (isExceptionSpanEvent(spanEvent)) {
return <SpanEventError spanEvent={spanEvent} exception={spanEvent.properties.exception} />;
}
@@ -58,7 +62,7 @@ function SpanEventError({
exception,
}: {
spanEvent: OtelSpanEvent;
exception: OtelExceptionProperty;
exception: ExceptionEventProperties;
}) {
return (
<div className="flex flex-col gap-2 rounded-sm border border-rose-500/50 p-3">
@@ -1,9 +1,7 @@
import { ChevronRightIcon } from "@heroicons/react/20/solid";
import { Span } from "@opentelemetry/sdk-trace-base";
import { TaskEventStyle } from "@trigger.dev/core/v3";
import { TaskEventLevel } from "@trigger.dev/database";
import type { TaskEventLevel } from "@trigger.dev/database";
import { Fragment } from "react";
import { Paragraph } from "~/components/primitives/Paragraph";
import { cn } from "~/utils/cn";
type SpanTitleProps = {
@@ -1,9 +1,7 @@
import { Attributes } from "@opentelemetry/api";
import { TaskEventStyle, unflattenAttributes } from "@trigger.dev/core/v3";
import { TaskEvent } from "@trigger.dev/database";
import { createTreeFromFlatItems, flattenTree } from "~/components/primitives/TreeView/TreeView";
import { PrismaClient, prisma } from "~/db.server";
import { getUsername } from "~/utils/username";
import { eventRepository } from "~/v3/eventRepository.server";
type Result = Awaited<ReturnType<RunPresenter["call"]>>;
export type Run = Result["run"];
@@ -62,46 +60,14 @@ export class RunPresenter {
});
// get the events
const events = await this.#prismaClient.$queryRaw<(TaskEvent & { rank: BigInt })[]>`
WITH ranked_events AS (
SELECT *,
ROW_NUMBER() OVER (PARTITION BY "spanId" ORDER BY "isPartial" ASC) as rank
FROM "TaskEvent"
WHERE "traceId" = ${run.traceId}
)
SELECT *
FROM ranked_events
WHERE rank = 1
ORDER BY "startTime" ASC;
`;
const traceSummary = await eventRepository.getTraceSummary(run.traceId);
const tree = createTreeFromFlatItems(
events.map((event) => {
const styleUnflattened = unflattenAttributes(event.style as Attributes);
const style = TaskEventStyle.parse(styleUnflattened);
return {
id: event.spanId,
parentId: event.parentId ?? undefined,
data: {
message: event.message,
style,
duration: Number(event.duration),
isError: event.isError,
isPartial: event.isPartial,
startTime: event.startTime,
level: event.level,
},
};
}),
run.spanId
);
const rootSpanId = events.find((event) => !event.parentId);
if (!rootSpanId) {
throw new Error("Root span not found");
if (!traceSummary) {
throw new Error("Trace not found");
}
const tree = createTreeFromFlatItems(traceSummary.spans, run.spanId);
return {
run: {
number: run.number,
@@ -114,7 +80,8 @@ export class RunPresenter {
},
},
events: tree ? flattenTree(tree) : [],
parentRunFriendlyId: tree?.id === rootSpanId.spanId ? undefined : rootSpanId.runId,
parentRunFriendlyId:
tree?.id === traceSummary.rootSpan.id ? undefined : traceSummary.rootSpan.runId,
};
}
}
@@ -1,40 +1,19 @@
import { Attributes } from "@opentelemetry/api";
import {
ExceptionEventProperties,
SemanticInternalAttributes,
TaskEventStyle,
SpanEvent,
SpanEvents,
correctErrorStackTrace,
isExceptionSpanEvent,
} from "@trigger.dev/core/v3";
import { unflattenAttributes } from "@trigger.dev/core/v3";
import { z } from "zod";
import { PrismaClient, prisma, Prisma } from "~/db.server";
import { PrismaClient, prisma } from "~/db.server";
import { eventRepository } from "~/v3/eventRepository.server";
type Result = Awaited<ReturnType<SpanPresenter["call"]>>;
export type Span = Result["event"];
const OtelExceptionProperty = z.object({
type: z.string().optional(),
message: z.string().optional(),
stacktrace: z.string().optional(),
});
export type OtelExceptionProperty = z.infer<typeof OtelExceptionProperty>;
const OtelSpanEvent = z.object({
name: z.string(),
time: z.coerce.date(),
properties: z
.object({
exception: OtelExceptionProperty.optional(),
})
.passthrough()
.optional(),
});
const OtelSpanEvents = z.array(OtelSpanEvent).optional();
type OtelSpanEvents = z.infer<typeof OtelSpanEvents>;
export type OtelSpanEvent = z.infer<typeof OtelSpanEvent>;
export class SpanPresenter {
#prismaClient: PrismaClient;
@@ -63,157 +42,20 @@ export class SpanPresenter {
throw new Error("Project not found");
}
// Find the project scoped to the organization
const matchingEvents = await this.#prismaClient.taskEvent.findMany({
where: {
spanId,
projectId: project.id,
},
});
const span = await eventRepository.getSpan(spanId);
const event =
matchingEvents.length > 1
? matchingEvents.find((event) => !event.isPartial)
: matchingEvents.at(0);
if (!event) {
throw new Error("Span not found");
if (!span) {
throw new Error("Event not found");
}
const styleUnflattened = unflattenAttributes(event.style as Attributes);
const style = TaskEventStyle.parse(styleUnflattened);
const eventsUnflattened = event.events
? (event.events as any[]).map((e) => ({
...e,
properties: unflattenAttributes(e.properties as Attributes),
}))
: undefined;
const events = OtelSpanEvents.parse(eventsUnflattened);
const payload = unflattenAttributes(
filteredAttributes(event.properties as Attributes, SemanticInternalAttributes.PAYLOAD)
)[SemanticInternalAttributes.PAYLOAD];
return {
event: {
...event,
events: transformEvents(events, event.metadata as Attributes),
output: isEmptyJson(event.output)
? null
: JSON.stringify(unflattenAttributes(event.output as Attributes), null, 2),
payload: payload ? JSON.stringify(payload, null, 2) : undefined,
properties: sanitizedAttributesStringified(event.properties),
style,
duration: Number(event.duration),
...span,
events: span.events,
output: span.output ? JSON.stringify(span.output, null, 2) : undefined,
payload: span.payload ? JSON.stringify(span.payload, null, 2) : undefined,
properties: span.properties ? JSON.stringify(span.properties, null, 2) : undefined,
},
};
}
}
function transformEvents(events: OtelSpanEvents, properties: Attributes): OtelSpanEvents {
return (events ?? []).map((event) => transformEvent(event, properties));
}
function transformEvent(event: OtelSpanEvent, properties: Attributes): OtelSpanEvent {
if (!event.properties?.exception) {
return event;
}
return {
...event,
properties: {
exception: transformException(event.properties.exception, properties),
},
};
}
function transformException(
exception: OtelExceptionProperty,
properties: Attributes
): OtelExceptionProperty {
const projectDirAttributeValue = properties[SemanticInternalAttributes.PROJECT_DIR];
if (typeof projectDirAttributeValue !== "string") {
return exception;
}
return {
...exception,
stacktrace: exception.stacktrace
? correctErrorStackTrace(exception.stacktrace, projectDirAttributeValue, {
removeFirstLine: true,
})
: undefined,
};
}
function filteredAttributes(attributes: Attributes, prefix: string): Attributes {
const result: Attributes = {};
for (const [key, value] of Object.entries(attributes)) {
if (key.startsWith(prefix)) {
result[key] = value;
}
}
return result;
}
function sanitizedAttributesStringified(json: Prisma.JsonValue): string | undefined {
const sanitizedAttributesValue = sanitizedAttributes(json);
if (!sanitizedAttributesValue) {
return;
}
return JSON.stringify(sanitizedAttributesValue, null, 2);
}
function sanitizedAttributes(json: Prisma.JsonValue): Record<string, unknown> | undefined {
if (json === null || json === undefined) {
return;
}
const withoutPrivateProperties = removePrivateProperties(json as Attributes);
if (!withoutPrivateProperties) {
return;
}
return unflattenAttributes(withoutPrivateProperties);
}
function isEmptyJson(json: Prisma.JsonValue) {
if (json === null) {
return true;
}
if (Object.keys(json).length === 0) {
return true;
}
return false;
}
// removes keys that start with a $ sign. If there are no keys left, return undefined
function removePrivateProperties(
attributes: Attributes | undefined | null
): Attributes | undefined {
if (!attributes) {
return undefined;
}
const result: Attributes = {};
for (const [key, value] of Object.entries(attributes)) {
if (key.startsWith("$")) {
continue;
}
result[key] = value;
}
if (Object.keys(result).length === 0) {
return undefined;
}
return result;
}
@@ -103,7 +103,7 @@ export default function Page() {
</div>
)}
{event.output !== null && (
{event.output && (
<div>
<Header2 spacing>Output</Header2>
<CodeBlock code={event.output} maxLines={20} />
@@ -3,12 +3,13 @@ import {
ChevronRightIcon,
ExclamationCircleIcon,
} from "@heroicons/react/20/solid";
import { Link, Outlet, useNavigate, useParams, useSubmit } from "@remix-run/react";
import { Link, Outlet, useNavigate } from "@remix-run/react";
import { LoaderFunctionArgs } from "@remix-run/server-runtime";
import { formatDuration, formatDurationNanoseconds } from "@trigger.dev/core/v3";
import { useCallback, useEffect, useRef, useState } from "react";
import { formatDurationNanoseconds } from "@trigger.dev/core/v3";
import { useRef, useState } from "react";
import { typedjson, useTypedLoaderData } from "remix-typedjson";
import { ShowParentIcon, ShowParentIconSelected } from "~/assets/icons/ShowParentIcon";
import { EnvironmentLabel } from "~/components/environments/EnvironmentLabel";
import { PageBody } from "~/components/layout/AppLayout";
import { Input } from "~/components/primitives/Input";
import {
@@ -26,21 +27,19 @@ import {
import { Spinner } from "~/components/primitives/Spinner";
import { Switch } from "~/components/primitives/Switch";
import { TreeView, useTree } from "~/components/primitives/TreeView/TreeView";
import { SpanTitle } from "~/components/runs/v3/SpanTitle";
import { LiveTimer } from "~/components/runs/v3/LiveTimer";
import { RunIcon } from "~/components/runs/v3/RunIcon";
import { SpanTitle } from "~/components/runs/v3/SpanTitle";
import { useDebounce } from "~/hooks/useDebounce";
import { useOrganization } from "~/hooks/useOrganizations";
import { usePathName } from "~/hooks/usePathName";
import { useProject } from "~/hooks/useProject";
import { useThrottle } from "~/hooks/useThrottle";
import { useUser } from "~/hooks/useUser";
import { RunEvent, RunPresenter } from "~/presenters/v3/RunPresenter.server";
import { getResizableRunSettings, setResizableRunSettings } from "~/services/resizablePanel";
import { requireUserId } from "~/services/session.server";
import { cn } from "~/utils/cn";
import { v3RunParamsSchema, v3RunPath, v3RunSpanPath } from "~/utils/pathBuilder";
import { EnvironmentLabel } from "~/components/environments/EnvironmentLabel";
import { useUser } from "~/hooks/useUser";
export const loader = async ({ request, params }: LoaderFunctionArgs) => {
const userId = await requireUserId(request);
@@ -286,6 +285,11 @@ function TasksTreeView({
) : node.data.duration > 0 ? (
<Duration duration={node.data.duration} />
) : null}
{node.data.isCancelled ? (
<Paragraph variant="extra-small" className="text-amber-500">
Cancelled
</Paragraph>
) : null}
</div>
</div>
</div>
@@ -16,7 +16,7 @@ export class AuthenticatedSocketConnection {
public onClose: Evt<CloseEvent> = new Evt();
private _sender: ZodMessageSender<typeof serverWebsocketMessages>;
private _environmentConsumer: DevQueueConsumer;
private _consumer: DevQueueConsumer;
private _messageHandler: ZodMessageHandler<typeof clientWebsocketMessages>;
constructor(public ws: WebSocket, public authenticatedEnv: AuthenticatedEnvironment) {
@@ -38,7 +38,7 @@ export class AuthenticatedSocketConnection {
},
});
this._environmentConsumer = new DevQueueConsumer(authenticatedEnv, this._sender);
this._consumer = new DevQueueConsumer(authenticatedEnv, this._sender);
ws.addEventListener("message", this.#handleMessage.bind(this));
ws.addEventListener("close", this.#handleClose.bind(this));
@@ -48,13 +48,13 @@ export class AuthenticatedSocketConnection {
schema: clientWebsocketMessages,
messages: {
READY_FOR_TASKS: async (payload) => {
await this._environmentConsumer.registerBackgroundWorker(payload.backgroundWorkerId);
await this._consumer.registerBackgroundWorker(payload.backgroundWorkerId);
},
BACKGROUND_WORKER_MESSAGE: async (payload) => {
switch (payload.data.type) {
case "TASK_RUN_COMPLETED": {
await this._environmentConsumer.taskRunCompleted(
await this._consumer.taskAttemptCompleted(
payload.backgroundWorkerId,
payload.data.completion,
payload.data.execution
@@ -62,10 +62,7 @@ export class AuthenticatedSocketConnection {
break;
}
case "TASK_HEARTBEAT": {
await this._environmentConsumer.taskHeartbeat(
payload.backgroundWorkerId,
payload.data.id
);
await this._consumer.taskHeartbeat(payload.backgroundWorkerId, payload.data.id);
break;
}
}
@@ -85,7 +82,7 @@ export class AuthenticatedSocketConnection {
}
async #handleClose(ev: CloseEvent) {
await this._environmentConsumer.stop();
await this._consumer.stop();
this.onClose.post(ev);
}
+516 -31
View File
@@ -1,14 +1,26 @@
import { Prisma, TaskEventStatus, type TaskEventKind } from "@trigger.dev/database";
import { PrismaClient, prisma } from "~/db.server";
import { Attributes } from "@opentelemetry/api";
import { RandomIdGenerator } from "@opentelemetry/sdk-trace-base";
import { Attributes, ROOT_CONTEXT, propagation, trace } from "@opentelemetry/api";
import { SemanticResourceAttributes } from "@opentelemetry/semantic-conventions";
import { SemanticInternalAttributes, PRIMARY_VARIANT } from "@trigger.dev/core/v3";
import { flattenAttributes } from "@trigger.dev/core/v3";
import {
ExceptionEventProperties,
PRIMARY_VARIANT,
SemanticInternalAttributes,
SpanEvent,
SpanEvents,
TaskEventStyle,
correctErrorStackTrace,
flattenAndNormalizeAttributes,
flattenAttributes,
isExceptionSpanEvent,
logger,
omit,
unflattenAttributes,
} from "@trigger.dev/core/v3";
import { Prisma, TaskEvent, TaskEventStatus, type TaskEventKind } from "@trigger.dev/database";
import { createHash } from "node:crypto";
import { PrismaClient, prisma } from "~/db.server";
import { AuthenticatedEnvironment } from "~/services/apiAuth.server";
import { DynamicFlushScheduler } from "./dynamicFlushScheduler.server";
import { logger } from "~/services/logger.server";
import { createHash } from "node:crypto";
export type CreatableEvent = Omit<
Prisma.TaskEventCreateInput,
@@ -50,6 +62,7 @@ export type TraceEventOptions = {
taskSlug: string;
startTime?: Date;
endTime?: Date;
immediate?: boolean;
};
export type EventBuilder = {
@@ -63,6 +76,44 @@ export type EventRepoConfig = {
batchInterval: number;
};
export type QueryOptions = Prisma.TaskEventWhereInput;
export type TaskEventRecord = TaskEvent;
export type QueriedEvent = TaskEvent;
export type PreparedEvent = Omit<TaskEventRecord, "events" | "style" | "duration"> & {
duration: number;
events: SpanEvents;
style: TaskEventStyle;
};
export type SpanSummary = {
recordId: string;
id: string;
parentId: string | undefined;
runId: string;
data: {
message: string;
style: TaskEventStyle;
events: SpanEvents;
startTime: Date;
duration: number;
isError: boolean;
isPartial: boolean;
isCancelled: boolean;
level: NonNullable<CreatableEvent["level"]>;
};
};
export type TraceSummary = { rootSpan: SpanSummary; spans: Array<SpanSummary> };
export type UpdateEventOptions = {
attributes: TraceAttributes;
endTime?: Date;
immediate?: boolean;
};
export class EventRepository {
private readonly _flushScheduler: DynamicFlushScheduler<CreatableEvent>;
@@ -80,10 +131,231 @@ export class EventRepository {
this._flushScheduler.addToBatch([event]);
}
async insertImmediate(event: CreatableEvent) {
await this.db.taskEvent.create({
data: event as Prisma.TaskEventCreateInput,
});
}
async insertMany(events: CreatableEvent[]) {
this._flushScheduler.addToBatch(events);
}
async completeEvent(spanId: string, options?: UpdateEventOptions) {
const events = await this.queryIncompleteEvents({ spanId });
if (events.length === 0) {
return;
}
const event = events[0];
logger.debug("Completing event", { spanId, eventId: event.id });
await this.insert({
...omit(event, "id"),
isPartial: false,
isError: options?.attributes.isError ?? false,
isCancelled: false,
status: options?.attributes.isError ? "ERROR" : "OK",
links: event.links ?? [],
events: event.events ?? [],
duration:
((options?.endTime ?? new Date()).getTime() - event.startTime.getTime()) * 1_000_000, // convert to nanoseconds
properties: event.properties as Attributes,
metadata: event.metadata as Attributes,
style: event.style as Attributes,
output: options?.attributes.output
? flattenAndNormalizeAttributes(
options.attributes.output,
SemanticInternalAttributes.OUTPUT
)
: undefined,
});
}
async cancelEvent(event: TaskEventRecord, cancelledAt: Date, reason: string) {
if (!event.isPartial) {
return;
}
await this.insertImmediate({
...omit(event, "id"),
isPartial: false,
isError: false,
isCancelled: true,
status: "ERROR",
links: event.links ?? [],
events: [
{
name: "cancellation",
time: cancelledAt,
properties: {
reason,
},
},
...((event.events as any[]) ?? []),
],
duration: (cancelledAt.getTime() - event.startTime.getTime()) * 1_000_000, // convert to nanoseconds
properties: event.properties as Attributes,
metadata: event.metadata as Attributes,
style: event.style as Attributes,
output: event.output as Attributes,
});
}
async queryEvents(queryOptions: QueryOptions): Promise<TaskEventRecord[]> {
return await this.db.taskEvent.findMany({
where: queryOptions,
});
}
async queryIncompleteEvents(queryOptions: QueryOptions) {
// First we will find all the events that match the query options (selecting minimal data).
const taskEvents = await this.db.taskEvent.findMany({
where: queryOptions,
select: {
spanId: true,
isPartial: true,
isCancelled: true,
},
});
const filteredTaskEvents = taskEvents.filter((event) => {
// Event must be partial
if (!event.isPartial) return false;
// If the event is cancelled, it is not incomplete
if (event.isCancelled) return false;
// There must not be another complete event with the same spanId
const hasCompleteDuplicate = taskEvents.some(
(otherEvent) =>
otherEvent.spanId === event.spanId && !otherEvent.isPartial && !otherEvent.isCancelled
);
return !hasCompleteDuplicate;
});
return this.queryEvents({
spanId: {
in: filteredTaskEvents.map((event) => event.spanId),
},
});
}
public async getTraceSummary(traceId: string): Promise<TraceSummary | undefined> {
const events = await this.db.taskEvent.findMany({
where: {
traceId,
},
orderBy: {
startTime: "asc",
},
});
const preparedEvents = removeDuplicateEvents(events.map(prepareEvent));
const spans = preparedEvents.map((event) => {
const ancestorCancelled = isAncestorCancelled(preparedEvents, event.spanId);
const duration = calculateDurationIfAncestorIsCancelled(
preparedEvents,
event.spanId,
event.duration
);
return {
recordId: event.id,
id: event.spanId,
parentId: event.parentId ?? undefined,
runId: event.runId,
data: {
message: event.message,
style: event.style,
duration,
isError: event.isError,
isPartial: ancestorCancelled ? false : event.isPartial,
isCancelled: event.isCancelled === true ? true : event.isPartial && ancestorCancelled,
startTime: event.startTime,
level: event.level,
events: event.events,
},
};
});
const rootSpanId = events.find((event) => !event.parentId);
if (!rootSpanId) {
return;
}
const rootSpan = spans.find((span) => span.id === rootSpanId.spanId);
if (!rootSpan) {
return;
}
return {
rootSpan,
spans,
};
}
// A Span can be cancelled if it is partial and has a parent that is cancelled
// And a span's duration, if it is partial and has a cancelled parent, is the time between the start of the span and the time of the cancellation event of the parent
public async getSpan(spanId: string) {
const traceSearch = await this.db.taskEvent.findFirst({
where: {
spanId,
},
select: {
traceId: true,
},
});
if (!traceSearch) {
return;
}
const traceSummary = await this.getTraceSummary(traceSearch.traceId);
const span = traceSummary?.spans.find((span) => span.id === spanId);
if (!span) {
return;
}
const fullEvent = await this.db.taskEvent.findUnique({
where: {
id: span.recordId,
},
});
if (!fullEvent) {
return;
}
const payload = unflattenAttributes(
filteredAttributes(fullEvent.properties as Attributes, SemanticInternalAttributes.PAYLOAD)
)[SemanticInternalAttributes.PAYLOAD];
const output = isEmptyJson(fullEvent.output)
? null
: unflattenAttributes(fullEvent.output as Attributes);
const properties = sanitizedAttributes(fullEvent.properties);
const events = transformEvents(span.data.events, fullEvent.metadata as Attributes);
return {
...fullEvent,
...span.data,
payload,
output,
properties,
events,
};
}
public async recordEvent(message: string, options: TraceEventOptions) {
const propagatedContext = extractContextFromCarrier(options.context ?? {});
@@ -154,14 +426,18 @@ export class EventRepository {
output: undefined,
};
this._flushScheduler.addToBatch([event]);
if (options.immediate) {
await this.insertImmediate(event);
} else {
this._flushScheduler.addToBatch([event]);
}
return event;
}
public async traceEvent<TResult>(
message: string,
options: TraceEventOptions,
options: TraceEventOptions & { incomplete?: boolean },
callback: (
e: EventBuilder,
traceContext: Record<string, string | undefined>
@@ -179,15 +455,6 @@ export class EventRepository {
? this.#generateDeterministicSpanId(traceId, options.spanIdSeed)
: this.generateSpanId();
logger.info("traceEvent", {
traceId,
parentId,
tracestate,
spanId,
context: options.context,
propagatedContext,
});
const traceContext = {
traceparent: `00-${traceId}-${spanId}-01`,
};
@@ -242,7 +509,8 @@ export class EventRepository {
spanId,
parentId,
tracestate,
duration: duration,
duration: options.incomplete ? 0 : duration,
isPartial: options.incomplete,
message: message,
serviceName: "api server",
serviceNamespace: "trigger.dev",
@@ -265,13 +533,18 @@ export class EventRepository {
string,
string
>),
...flattenAttributes(options.attributes.properties),
},
metadata: metadata,
style: stripAttributePrefix(style, SemanticInternalAttributes.STYLE),
output: undefined,
};
this._flushScheduler.addToBatch([event]);
if (options.immediate) {
await this.insertImmediate(event);
} else {
this._flushScheduler.addToBatch([event]);
}
return result;
}
@@ -375,16 +648,228 @@ function parseTraceparent(traceparent?: string): { traceId: string; spanId: stri
return { traceId, spanId };
}
const SHARED_CHAR_CODES_ARRAY = Array(32);
function getIdGenerator(bytes: number): () => string {
return function generateId() {
for (let i = 0; i < bytes * 2; i++) {
SHARED_CHAR_CODES_ARRAY[i] = Math.floor(Math.random() * 16) + 48;
// valid hex characters in the range 48-57 and 97-102
if (SHARED_CHAR_CODES_ARRAY[i] >= 58) {
SHARED_CHAR_CODES_ARRAY[i] += 39;
}
}
return String.fromCharCode.apply(null, SHARED_CHAR_CODES_ARRAY.slice(0, bytes * 2));
function prepareEvent(event: QueriedEvent): PreparedEvent {
return {
...event,
duration: Number(event.duration),
events: parseEventsField(event.events),
style: parseStyleField(event.style),
};
}
function parseEventsField(events: Prisma.JsonValue): SpanEvents {
const eventsUnflattened = events
? (events as any[]).map((e) => ({
...e,
properties: unflattenAttributes(e.properties as Attributes),
}))
: undefined;
const spanEvents = SpanEvents.safeParse(eventsUnflattened);
if (spanEvents.success) {
return spanEvents.data;
}
return [];
}
function parseStyleField(style: Prisma.JsonValue): TaskEventStyle {
const parsedStyle = TaskEventStyle.safeParse(unflattenAttributes(style as Attributes));
if (parsedStyle.success) {
return parsedStyle.data;
}
return {};
}
function isAncestorCancelled(events: PreparedEvent[], spanId: string) {
const event = events.find((event) => event.spanId === spanId);
if (!event) {
return false;
}
if (event.isCancelled) {
return true;
}
if (event.parentId) {
return isAncestorCancelled(events, event.parentId);
}
return false;
}
function calculateDurationIfAncestorIsCancelled(
events: PreparedEvent[],
spanId: string,
defaultDuration: number
) {
const event = events.find((event) => event.spanId === spanId);
if (!event) {
return defaultDuration;
}
if (event.isCancelled) {
return defaultDuration;
}
if (!event.isPartial) {
return defaultDuration;
}
if (event.parentId) {
const cancelledAncestor = findFirstCancelledAncestor(events, event.parentId);
if (cancelledAncestor) {
// We need to get the cancellation time from the cancellation span event
const cancellationEvent = cancelledAncestor.events.find(
(event) => event.name === "cancellation"
);
if (cancellationEvent) {
return (cancellationEvent.time.getTime() - event.startTime.getTime()) * 1_000_000;
}
}
}
return defaultDuration;
}
function findFirstCancelledAncestor(events: PreparedEvent[], spanId: string) {
const event = events.find((event) => event.spanId === spanId);
if (!event) {
return;
}
if (event.isCancelled) {
return event;
}
if (event.parentId) {
return findFirstCancelledAncestor(events, event.parentId);
}
return;
}
// Prioritize spans with the same id, keeping the completed spans over partial spans
// Completed spans are either !isPartial or isCancelled
function removeDuplicateEvents(events: PreparedEvent[]) {
const dedupedEvents = new Map<string, PreparedEvent>();
for (const event of events) {
const existingEvent = dedupedEvents.get(event.spanId);
if (!existingEvent) {
dedupedEvents.set(event.spanId, event);
continue;
}
if (event.isCancelled || !event.isPartial) {
dedupedEvents.set(event.spanId, event);
}
}
return Array.from(dedupedEvents.values());
}
function isEmptyJson(json: Prisma.JsonValue) {
if (json === null) {
return true;
}
if (Object.keys(json).length === 0) {
return true;
}
return false;
}
function sanitizedAttributes(json: Prisma.JsonValue): Record<string, unknown> | undefined {
if (json === null || json === undefined) {
return;
}
const withoutPrivateProperties = removePrivateProperties(json as Attributes);
if (!withoutPrivateProperties) {
return;
}
return unflattenAttributes(withoutPrivateProperties);
}
// removes keys that start with a $ sign. If there are no keys left, return undefined
function removePrivateProperties(
attributes: Attributes | undefined | null
): Attributes | undefined {
if (!attributes) {
return undefined;
}
const result: Attributes = {};
for (const [key, value] of Object.entries(attributes)) {
if (key.startsWith("$")) {
continue;
}
result[key] = value;
}
if (Object.keys(result).length === 0) {
return undefined;
}
return result;
}
function transformEvents(events: SpanEvents, properties: Attributes): SpanEvents {
return (events ?? []).map((event) => transformEvent(event, properties));
}
function transformEvent(event: SpanEvent, properties: Attributes): SpanEvent {
if (isExceptionSpanEvent(event)) {
return {
...event,
properties: {
exception: transformException(event.properties.exception, properties),
},
};
}
return event;
}
function transformException(
exception: ExceptionEventProperties,
properties: Attributes
): ExceptionEventProperties {
const projectDirAttributeValue = properties[SemanticInternalAttributes.PROJECT_DIR];
if (typeof projectDirAttributeValue !== "string") {
return exception;
}
return {
...exception,
stacktrace: exception.stacktrace
? correctErrorStackTrace(exception.stacktrace, projectDirAttributeValue, {
removeFirstLine: true,
})
: undefined,
};
}
function filteredAttributes(attributes: Attributes, prefix: string): Attributes {
const result: Attributes = {};
for (const [key, value] of Object.entries(attributes)) {
if (key.startsWith(prefix)) {
result[key] = value;
}
}
return result;
}
@@ -20,8 +20,10 @@ import { marqs } from "../marqs.server";
import { attributesFromAuthenticatedEnv } from "../tracer.server";
import { eventRepository } from "../eventRepository.server";
import { EnvironmentVariablesRepository } from "../environmentVariables/environmentVariablesRepository.server";
import { CancelAttemptService } from "../services/cancelAttempt.server";
import { CompleteAttemptService } from "../services/completeAttempt.server";
const tracer = trace.getTracer("environmentQueueConsumer");
const tracer = trace.getTracer("devQueueConsumer");
const MessageBody = z.discriminatedUnion("type", [
z.object({
@@ -32,7 +34,7 @@ const MessageBody = z.discriminatedUnion("type", [
type BackgroundWorkerWithTasks = BackgroundWorker & { tasks: BackgroundWorkerTask[] };
export type EnvironmentQueueConsumerOptions = {
export type DevQueueConsumerOptions = {
maximumItemsPerTrace?: number;
traceTimeoutSeconds?: number;
};
@@ -40,7 +42,7 @@ export type EnvironmentQueueConsumerOptions = {
export class DevQueueConsumer {
private _backgroundWorkers: Map<string, BackgroundWorkerWithTasks> = new Map();
private _enabled = false;
private _options: Required<EnvironmentQueueConsumerOptions>;
private _options: Required<DevQueueConsumerOptions>;
private _perTraceCountdown: number | undefined;
private _lastNewTrace: Date | undefined;
private _currentSpanContext: Context | undefined;
@@ -48,11 +50,12 @@ export class DevQueueConsumer {
private _taskSuccesses: number = 0;
private _currentSpan: Span | undefined;
private _endSpanInNextIteration = false;
private _inProgressAttempts: Map<string, string> = new Map(); // Keys are task attempt friendly IDs, values are TaskRun ids/queue message ids
constructor(
public env: AuthenticatedEnvironment,
private _sender: ZodMessageSender<typeof serverWebsocketMessages>,
options: EnvironmentQueueConsumerOptions = {}
options: DevQueueConsumerOptions = {}
) {
this._options = {
maximumItemsPerTrace: options.maximumItemsPerTrace ?? 1_000, // 1k items per trace
@@ -80,90 +83,23 @@ export class DevQueueConsumer {
this.#enable();
}
public async taskRunCompleted(
public async taskAttemptCompleted(
workerId: string,
completion: TaskRunExecutionResult,
execution: TaskRunExecution
) {
logger.debug("Task run completed", { taskRunCompletion: completion, execution });
this._inProgressAttempts.delete(completion.id);
const taskRunAttempt = completion.ok
? await prisma.taskRunAttempt.update({
where: { friendlyId: completion.id },
data: {
status: "COMPLETED",
completedAt: new Date(),
output: completion.output,
outputType: completion.outputType,
},
include: {
taskRun: true,
backgroundWorkerTask: true,
},
})
: await prisma.taskRunAttempt.update({
where: { friendlyId: completion.id },
data: {
status: "FAILED",
completedAt: new Date(),
error: completion.error,
},
include: {
taskRun: true,
backgroundWorkerTask: true,
},
});
if (taskRunAttempt.status === "COMPLETED") {
if (completion.ok) {
this._taskSuccesses++;
} else {
this._taskFailures++;
}
if (!completion.ok && completion.retry !== undefined) {
const retryConfig = taskRunAttempt.backgroundWorkerTask.retryConfig
? {
...defaultRetryOptions,
...RetryOptions.parse(taskRunAttempt.backgroundWorkerTask.retryConfig),
}
: undefined;
logger.debug("Task run completed", { taskRunCompletion: completion, execution });
const retryAt = new Date(completion.retry.timestamp);
// Retry the task run
await eventRepository.recordEvent(
retryConfig?.maxAttempts
? `Retry ${execution.attempt.number}/${retryConfig?.maxAttempts - 1} delay`
: `Retry #${execution.attempt.number} delay`,
{
taskSlug: taskRunAttempt.taskRun.taskIdentifier,
environment: this.env,
attributes: {
metadata: this.#generateMetadataAttributesForNextAttempt(execution),
properties: {
retryAt: retryAt.toISOString(),
factor: retryConfig?.factor,
maxAttempts: retryConfig?.maxAttempts,
minTimeoutInMs: retryConfig?.minTimeoutInMs,
maxTimeoutInMs: retryConfig?.maxTimeoutInMs,
randomize: retryConfig?.randomize,
},
runId: taskRunAttempt.taskRunId,
style: {
icon: "schedule-attempt",
},
queueId: taskRunAttempt.queueId,
queueName: taskRunAttempt.taskRun.queue,
},
context: taskRunAttempt.taskRun.traceContext as Record<string, string | undefined>,
spanIdSeed: `retry-${taskRunAttempt.number + 1}`,
endTime: retryAt,
}
);
await marqs?.nackMessage(taskRunAttempt.taskRunId, completion.retry.timestamp);
} else {
await marqs?.acknowledgeMessage(taskRunAttempt.taskRunId);
}
const service = new CompleteAttemptService();
await service.call(completion, execution, this.env);
}
#generateMetadataAttributesForNextAttempt(execution: TaskRunExecution) {
@@ -189,8 +125,49 @@ export class DevQueueConsumer {
await marqs?.heartbeatMessage(taskRunAttempt.taskRunId, seconds);
}
public async stop() {
public async stop(reason: string = "CLI disconnected") {
if (!this._enabled) {
return;
}
logger.debug("Stopping dev queue consumer", { env: this.env });
this._enabled = false;
// We need to cancel all the in progress task run attempts and ack the messages so they will stop processing
await this.#cancelInProgressAttempts(reason);
}
async #cancelInProgressAttempts(reason: string) {
const service = new CancelAttemptService();
const cancelledAt = new Date();
const inProgressAttempts = new Map(this._inProgressAttempts);
this._inProgressAttempts.clear();
for (const [attemptId, messageId] of inProgressAttempts) {
await this.#cancelInProgressAttempt(attemptId, messageId, service, cancelledAt, reason);
}
}
async #cancelInProgressAttempt(
attemptId: string,
messageId: string,
cancelAttemptService: CancelAttemptService,
cancelledAt: Date,
reason: string
) {
try {
await cancelAttemptService.call(attemptId, messageId, cancelledAt, reason, this.env);
} catch (e) {
logger.error("Failed to cancel in progress attempt", {
attemptId,
messageId,
error: e,
});
}
}
#enable() {
@@ -228,7 +205,7 @@ export class DevQueueConsumer {
// Create a new trace
this._currentSpan = tracer.startSpan(
"EnvironmentQueueConsumer.doWork()",
"DevQueueConsumer.doWork()",
{
kind: SpanKind.CONSUMER,
attributes: {
@@ -456,6 +433,8 @@ export class DevQueueConsumer {
payloads: [payload],
},
});
this._inProgressAttempts.set(taskRunAttempt.friendlyId, message.messageId);
} catch (e) {
if (e instanceof Error) {
this._currentSpan?.recordException(e);
@@ -0,0 +1,58 @@
import { AuthenticatedEnvironment } from "~/services/apiAuth.server";
import { eventRepository } from "../eventRepository.server";
import { marqs } from "../marqs.server";
import { BaseService } from "./baseService.server";
import { logger } from "~/services/logger.server";
export class CancelAttemptService extends BaseService {
public async call(
attemptId: string,
taskRunId: string,
cancelledAt: Date,
reason: string,
environment: AuthenticatedEnvironment
) {
return await this.traceWithEnv("call()", environment, async (span) => {
span.setAttribute("taskRunId", taskRunId);
span.setAttribute("attemptId", attemptId);
const taskRunAttempt = await this._prisma.taskRunAttempt.findUnique({
where: {
friendlyId: attemptId,
},
include: {
taskRun: true,
},
});
if (!taskRunAttempt) {
return;
}
await marqs?.acknowledgeMessage(taskRunId);
await this._prisma.taskRunAttempt.update({
where: {
friendlyId: attemptId,
},
data: {
status: "CANCELED",
},
});
const inProgressEvents = await eventRepository.queryIncompleteEvents({
runId: taskRunAttempt.taskRun.friendlyId,
});
logger.debug("Cancelling in-progress events", {
inProgressEvents: inProgressEvents.map((event) => event.id),
});
await Promise.all(
inProgressEvents.map((event) => {
return eventRepository.cancelEvent(event, cancelledAt, reason);
})
);
});
}
}
@@ -0,0 +1,126 @@
import {
RetryOptions,
TaskRunContext,
TaskRunExecution,
TaskRunExecutionResult,
defaultRetryOptions,
flattenAttributes,
} from "@trigger.dev/core/v3";
import { AuthenticatedEnvironment } from "~/services/apiAuth.server";
import { eventRepository } from "../eventRepository.server";
import { marqs } from "../marqs.server";
import { BaseService } from "./baseService.server";
import { Attributes } from "@opentelemetry/api";
export class CompleteAttemptService extends BaseService {
public async call(
completion: TaskRunExecutionResult,
execution: TaskRunExecution,
env: AuthenticatedEnvironment
): Promise<"ACKNOWLEDGED" | "RETRIED"> {
const taskRunAttempt = completion.ok
? await this._prisma.taskRunAttempt.update({
where: { friendlyId: completion.id },
data: {
status: "COMPLETED",
completedAt: new Date(),
output: completion.output,
outputType: completion.outputType,
},
include: {
taskRun: true,
backgroundWorkerTask: true,
},
})
: await this._prisma.taskRunAttempt.update({
where: { friendlyId: completion.id },
data: {
status: "FAILED",
completedAt: new Date(),
error: completion.error,
},
include: {
taskRun: true,
backgroundWorkerTask: true,
},
});
if (!completion.ok && completion.retry !== undefined) {
const retryConfig = taskRunAttempt.backgroundWorkerTask.retryConfig
? {
...defaultRetryOptions,
...RetryOptions.parse(taskRunAttempt.backgroundWorkerTask.retryConfig),
}
: undefined;
const retryAt = new Date(completion.retry.timestamp);
// Retry the task run
await eventRepository.recordEvent(
retryConfig?.maxAttempts
? `Retry ${execution.attempt.number}/${retryConfig?.maxAttempts - 1} delay`
: `Retry #${execution.attempt.number} delay`,
{
taskSlug: taskRunAttempt.taskRun.taskIdentifier,
environment: env,
attributes: {
metadata: this.#generateMetadataAttributesForNextAttempt(execution),
properties: {
retryAt: retryAt.toISOString(),
factor: retryConfig?.factor,
maxAttempts: retryConfig?.maxAttempts,
minTimeoutInMs: retryConfig?.minTimeoutInMs,
maxTimeoutInMs: retryConfig?.maxTimeoutInMs,
randomize: retryConfig?.randomize,
},
runId: taskRunAttempt.taskRunId,
style: {
icon: "schedule-attempt",
},
queueId: taskRunAttempt.queueId,
queueName: taskRunAttempt.taskRun.queue,
},
context: taskRunAttempt.taskRun.traceContext as Record<string, string | undefined>,
spanIdSeed: `retry-${taskRunAttempt.number + 1}`,
endTime: retryAt,
}
);
await marqs?.nackMessage(taskRunAttempt.taskRunId, completion.retry.timestamp);
return "RETRIED";
} else {
await marqs?.acknowledgeMessage(taskRunAttempt.taskRunId);
// Now we need to "complete" the task run event/span
if (completion.ok) {
await eventRepository.completeEvent(taskRunAttempt.taskRun.spanId, {
endTime: new Date(),
attributes: {
isError: false,
output: JSON.parse(completion.output) as Attributes,
},
});
} else {
await eventRepository.completeEvent(taskRunAttempt.taskRun.spanId, {
endTime: new Date(),
attributes: {
isError: true,
},
});
}
return "ACKNOWLEDGED";
}
}
#generateMetadataAttributesForNextAttempt(execution: TaskRunExecution) {
const context = TaskRunContext.parse(execution);
// @ts-ignore
context.attempt = {
number: context.attempt.number + 1,
};
return flattenAttributes(context, "ctx");
}
}
@@ -47,21 +47,23 @@ export class TriggerTaskService extends BaseService {
}
return await eventRepository.traceEvent(
`${taskId}`,
taskId,
{
context: options.traceContext,
kind: "SERVER",
environment,
taskSlug: taskId,
attributes: {
metadata: {
...flattenAttributes(body.payload, SemanticInternalAttributes.PAYLOAD),
properties: {
[SemanticInternalAttributes.PAYLOAD]: body.payload,
},
style: {
icon: "play",
variant: PRIMARY_VARIANT,
},
},
incomplete: true,
immediate: true,
},
async (event, traceContext) => {
const lockId = taskIdentifierToLockId(taskId);
+6 -1
View File
@@ -33,7 +33,12 @@ export { TriggerTracer } from "./tracer";
export type { TaskLogger } from "./logger/taskLogger";
export { OtelTaskLogger } from "./logger/taskLogger";
export { ConsoleInterceptor } from "./consoleInterceptor";
export { flattenAttributes, unflattenAttributes } from "./utils/flattenAttributes";
export {
flattenAttributes,
unflattenAttributes,
flattenAndNormalizeAttributes,
} from "./utils/flattenAttributes";
export { defaultRetryOptions, calculateNextRetryDelay, calculateResetAt } from "./utils/retries";
export { accessoryAttributes } from "./utils/styleAttributes";
export { eventFilterMatches } from "../eventFilterMatches";
export { omit } from "./utils/omit";
+1
View File
@@ -6,3 +6,4 @@ export * from "./messages";
export * from "./style";
export * from "./fetch";
export * from "./eventFilter";
export * from "./openTelemetry";
@@ -0,0 +1,53 @@
import { z } from "zod";
export const ExceptionEventProperties = z.object({
type: z.string().optional(),
message: z.string().optional(),
stacktrace: z.string().optional(),
});
export type ExceptionEventProperties = z.infer<typeof ExceptionEventProperties>;
export const ExceptionSpanEvent = z.object({
name: z.literal("exception"),
time: z.coerce.date(),
properties: z.object({
exception: ExceptionEventProperties,
}),
});
export type ExceptionSpanEvent = z.infer<typeof ExceptionSpanEvent>;
export const CancellationSpanEvent = z.object({
name: z.literal("cancellation"),
time: z.coerce.date(),
properties: z.object({
reason: z.string(),
}),
});
export type CancellationSpanEvent = z.infer<typeof CancellationSpanEvent>;
export const OtherSpanEvent = z.object({
name: z.string(),
time: z.coerce.date(),
properties: z.record(z.unknown()),
});
export type OtherSpanEvent = z.infer<typeof OtherSpanEvent>;
export const SpanEvent = z.union([ExceptionSpanEvent, CancellationSpanEvent, OtherSpanEvent]);
export type SpanEvent = z.infer<typeof SpanEvent>;
export const SpanEvents = z.array(SpanEvent);
export type SpanEvents = z.infer<typeof SpanEvents>;
export function isExceptionSpanEvent(event: SpanEvent): event is ExceptionSpanEvent {
return event.name === "exception";
}
export function isCancellationSpanEvent(event: SpanEvent): event is CancellationSpanEvent {
return event.name === "cancellation";
}
@@ -94,3 +94,16 @@ export function unflattenAttributes(obj: Attributes): Record<string, unknown> {
return result;
}
export function flattenAndNormalizeAttributes(
obj: Record<string, unknown> | Array<unknown> | string | boolean | number | undefined,
prefix: string
): Attributes {
const attributes = flattenAttributes(obj, prefix);
if (typeof attributes[prefix] !== "undefined" && attributes[prefix] !== null) {
return attributes[prefix] as unknown as Attributes;
}
return attributes;
}
+14
View File
@@ -0,0 +1,14 @@
export function omit<T extends Record<string, any>, K extends keyof T>(
obj: T,
...keys: K[]
): Omit<T, K> {
const result: Record<string, any> = {};
for (const key in obj) {
if (!keys.includes(key as unknown as K)) {
result[key] = obj[key];
}
}
return result as Omit<T, K>;
}
@@ -0,0 +1,2 @@
-- AlterTable
ALTER TABLE "TaskEvent" ADD COLUMN "isCancelled" BOOLEAN NOT NULL DEFAULT false;
+3 -2
View File
@@ -1694,8 +1694,9 @@ model TaskEvent {
parentId String?
tracestate String?
isError Boolean @default(false)
isPartial Boolean @default(false)
isError Boolean @default(false)
isPartial Boolean @default(false)
isCancelled Boolean @default(false)
serviceName String
serviceNamespace String