trace.ts•2.21 kB
/**
* Copyright 2024 Google LLC
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
import type { NestedSpanData, TraceData } from '../types';
/** Transforms flat spans from the trace data into a tree of spans. */
export function stackTraceSpans(trace: TraceData): NestedSpanData | undefined {
if (!trace.spans || Object.keys(trace.spans).length === 0) return undefined;
let rootSpan: NestedSpanData | undefined = undefined;
const treeSpans: Map<string, NestedSpanData> = new Map();
Object.values(trace.spans).forEach((span) => {
treeSpans.set(span.spanId, { spans: [], ...span } as NestedSpanData);
if (!span.parentSpanId) {
rootSpan = treeSpans.get(span.spanId);
}
});
// Build the tree of spans.
treeSpans.forEach((span) => {
if (span.parentSpanId && span.spanId !== rootSpan?.spanId) {
const parent = treeSpans.get(span.parentSpanId);
if (parent) {
parent.spans?.push(span);
} else if (!rootSpan) {
// This shouldn't happen, but there is a parentSpanId that we cannot
// find. In this case, fallback to this span being the root so the UI
// can still render properly.
rootSpan = span;
}
}
});
// Sort children by start times for each span.
treeSpans.forEach((span) => {
span.spans?.sort((a, b) => a.startTime - b.startTime);
});
// Re-position the root node to the first non-internal span
let bestRoot = rootSpan!;
while (
bestRoot.attributes['genkit:metadata:genkit-dev-internal'] ||
bestRoot.attributes['genkit:metadata:flow:wrapperAction']
) {
if (!bestRoot.spans?.length) {
break;
}
bestRoot = bestRoot.spans[0];
}
return bestRoot;
}