Skip to content
Merged
Show file tree
Hide file tree
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -5,8 +5,8 @@ import * as fs from 'fs';
import * as path from 'path';
import { Construct } from 'constructs';
import * as cxschema from 'aws-cdk-lib/cloud-assembly-schema';
import { App, CfnParameter, CfnResource, Lazy, Stack, TreeInspector } from 'aws-cdk-lib';
import { TreeFile } from 'aws-cdk-lib/core/lib/private/tree-metadata';
import { App, AssumptionError, CfnParameter, CfnResource, Lazy, Stack, TreeInspector } from 'aws-cdk-lib';
import { ForestFile, TreeFile } from 'aws-cdk-lib/core/lib/private/tree-metadata';

abstract class AbstractCfnResource extends CfnResource {
constructor(scope: Construct, id: string) {
Expand Down Expand Up @@ -219,16 +219,31 @@ describe('tree metadata', () => {
expect(treeArtifact).toBeDefined();

// THEN - does not explode, and file sizes are correctly limited
const sizes: Record<string, number> = {};
recurseVisit(assembly.directory, treeArtifact!.file, sizes);
const treeSizes: Record<string, number> = {};
recurseVisit(assembly.directory, treeArtifact!.file, undefined, treeSizes);

// `treeSizes` has keys of the form `<file>#<tree id>`. Combine them.
const fileSizes: Record<string, number> = {};
for (const [treeName, count] of Object.entries(treeSizes)) {
const fileName = treeName.split('#')[0];
if (fileName in fileSizes) {
fileSizes[fileName] += count;
} else {
fileSizes[fileName ] = count;
}
}

for (const size of Object.values(sizes)) {
for (const size of Object.values(treeSizes)) {
expect(size).toBeLessThanOrEqual(MAX_NODES);
}
for (const size of Object.values(fileSizes)) {
expect(size).toBeLessThanOrEqual(MAX_NODES);
}

expect(Object.keys(sizes).length).toBeGreaterThan(1);
expect(Object.keys(treeSizes).length).toBeGreaterThan(1);
expect(Object.keys(fileSizes).length).toBeLessThan(Object.keys(treeSizes).length);

const foundNodes = sum(Object.values(sizes));
const foundNodes = sum(Object.values(treeSizes));
expect(foundNodes).toEqual(addedNodes + 2); // App, Tree
} finally {
fs.rmSync(assembly.directory, { force: true, recursive: true });
Expand All @@ -253,16 +268,29 @@ describe('tree metadata', () => {
return ret;
}

function recurseVisit(directory: string, fileName: string, files: Record<string, number>) {
let nodes = 0;
const treeJson: TreeFile = readJson(directory, fileName);
rec(treeJson.tree);
files[fileName] = nodes;
function recurseVisit(directory: string, fileName: string, treeId: string | undefined, files: Record<string, number>) {
let nodes: number;

if (treeId) {
// Assume a forest file
const treeJson: ForestFile = readJson(directory, fileName);
for (const [tid, tree] of Object.entries(treeJson.forest)) {
nodes = 0;
rec(tree);
files[`${fileName}#${tid}`] = nodes;
}
} else {
// Assume a tree file
const treeJson: TreeFile = readJson(directory, fileName);
nodes = 0;
rec(treeJson.tree);
files[fileName] = nodes;
}

function rec(x: TreeFile['tree']) {
if (isSubtreeReference(x)) {
// We'll count this node as part of our visit to the "real" node
recurseVisit(directory, x.fileName, files);
recurseVisit(directory, x.fileName, x.treeId, files);
} else {
nodes += 1;
for (const child of Object.values(x.children ?? {})) {
Expand Down Expand Up @@ -437,7 +465,7 @@ describe('tree metadata', () => {
test('failing nodes', () => {
class MyCfnResource extends CfnResource {
public inspect(_: TreeInspector) {
throw new Error('Forcing an inspect error');
throw new AssumptionError('Forcing an inspect error');
}
}

Expand Down
112 changes: 95 additions & 17 deletions packages/aws-cdk-lib/core/lib/private/tree-metadata.ts
Original file line number Diff line number Diff line change
Expand Up @@ -113,6 +113,11 @@ export interface TreeFile {
tree: TreeNode;
}

export interface ForestFile {
version: 'forest-0.1';
forest: Record<string, TreeNode>;
}

type TreeNode = Node | SubTreeReference;

/**
Expand All @@ -121,7 +126,15 @@ type TreeNode = Node | SubTreeReference;
interface SubTreeReference {
readonly id: string;
readonly path: string;
readonly fileName: string;
fileName: string;

/**
* If set, indicates the subtree in the forest file
*
* If this is set then `fileName` must point to a ForestFile, and this indicates
* the tree inside the forest.
*/
treeId?: string;
}

/**
Expand Down Expand Up @@ -149,6 +162,17 @@ interface SubTreeReference {
* file (490 bytes/node). We'll estimate the size of a node to be 1000 bytes.
*/
class FragmentedTreeWriter {
/**
* We only care about the identify of this object.
*
* Whatever tree in the forest is "pointed to" by this pointer is the main tree.
*/
private readonly mainTreePointer: SubTreeReference = {
fileName: 'yyy',
id: 'id',
path: 'path',
};

private readonly forest = new Array<Tree>();

/**
Expand All @@ -168,8 +192,6 @@ class FragmentedTreeWriter {

private readonly maxNodes: number;

private subtreeCtr = 1;

constructor(private readonly outdir: string, private readonly rootFilename: string, options?: FragmentedTreeWriterOptions) {
this.maxNodes = options?.maxNodesPerTree ?? 500_000;
}
Expand All @@ -178,14 +200,59 @@ class FragmentedTreeWriter {
* Write the forest to disk, return the root file name
*/
public writeForest(): string {
for (const tree of this.forest) {
const treeFile: TreeFile = { version: 'tree-0.1', tree: tree.root };
fs.writeFileSync(path.join(this.outdir, tree.filename), JSON.stringify(treeFile), { encoding: 'utf-8' });
const forestFiles = this.allocateSubTreesToForestFiles();

// We can now write the forest files, and the main file.
const mainTree = this.forest.find(t => t.referencingNode === this.mainTreePointer);
if (mainTree) {
const treeFile: TreeFile = { version: 'tree-0.1', tree: mainTree.root };
fs.writeFileSync(path.join(this.outdir, this.rootFilename), JSON.stringify(treeFile), { encoding: 'utf-8' });
}

for (const forestFile of forestFiles) {
fs.writeFileSync(path.join(this.outdir, forestFile.fileName), JSON.stringify(forestFile.file), { encoding: 'utf-8' });
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can we write the files in parallel?

}

return this.rootFilename;
}

/**
* Find all non-main tree and combine them into forest files
*
* This will mutate the pointing nodes as a side effect.
*/
private allocateSubTreesToForestFiles(): IncompleteForestFile[] {
// First, find all non-main trees and allocate them to forests.
const ret = new Array<IncompleteForestFile>();

for (const tree of this.forest) {
if (tree.referencingNode === this.mainTreePointer) {
// Main tree, not interesting for the purposes of allocating subtrees to forests
continue;
}

let targetForest: typeof ret[0];
if (ret.length === 0 || ret[ret.length - 1].nodeCount + tree.nodes > this.maxNodes) {
targetForest = {
fileName: `trees-${ret.length + 1}.json`,
file: { version: 'forest-0.1', forest: { } },
nodeCount: 0,
};
ret.push(targetForest);
} else {
targetForest = ret[ret.length - 1];
}

const treeId = `t${Object.keys(targetForest.file.forest).length}`;
targetForest.file.forest[treeId] = tree.root;
targetForest.nodeCount += tree.nodes;
tree.referencingNode.fileName = targetForest.fileName;
tree.referencingNode.treeId = treeId;
}

return ret;
}

public addNode(construct: IConstruct, parent: IConstruct | undefined, node: Node) {
// NOTE: we could copy the 'node' object to be safe against tampering, but we trust
// the consuming code so we know we don't need to.
Expand All @@ -195,7 +262,7 @@ class FragmentedTreeWriter {
throw new AssumptionError('Can only add exactly one node without a parent');
}

this.addNewTree(node, this.rootFilename);
this.addNewTree(node, this.mainTreePointer);
} else {
// There was a provision in the old code for missing parents, so we're just going to ignore it
// if we can't find a parent.
Expand All @@ -213,10 +280,10 @@ class FragmentedTreeWriter {
/**
* Add a new tree with the given Node as root
*/
private addNewTree(root: Node, filename: string): Tree {
private addNewTree(root: Node, referencingNode: SubTreeReference): Tree {
const tree: Tree = {
root,
filename,
referencingNode: referencingNode,
nodes: nodeCount(root),
};

Expand All @@ -240,13 +307,15 @@ class FragmentedTreeWriter {
throw new AssumptionError(`Could not find parent of ${JSON.stringify(parent)}`);
}

tree = this.addNewTree(parent, `tree-${this.subtreeCtr++}.json`);

setChild(grandParent, {
const subtreeReference: SubTreeReference = {
id: parent.id,
path: parent.path,
fileName: tree.filename,
} satisfies SubTreeReference);
fileName: 'xxx', // Will be replaced later
};

tree = this.addNewTree(parent, subtreeReference);

setChild(grandParent, subtreeReference);

// To be strictly correct we should decrease the original tree's nodeCount here, because
// we may have moved away any number of children as well. We don't do that; the tree
Expand Down Expand Up @@ -286,7 +355,7 @@ class FragmentedTreeWriter {
if (tree) {
return tree;
}
throw new AssumptionError(`Could not find tree for node: ${JSON.stringify(node)}, tried ${tried}, ${Array.from(this.subtreeRoots).map(([k, v]) => `${k.path} => ${v.filename}`)}`);
throw new AssumptionError(`Could not find tree for node: ${JSON.stringify(node)}, tried ${tried}`);
}
}

Expand Down Expand Up @@ -329,17 +398,26 @@ interface Tree {
* The root of this particular tree
*/
root: Node;

/**
* The filename that `root` will be serialized to
* The node that is pointing to this tree
*
* This may be "mainTreePointer", in which case this tree indicates the main tree.
*/
filename: string;
referencingNode: SubTreeReference;

/**
* How many nodes are in this tree already
*/
nodes: number;
}

interface IncompleteForestFile {
fileName: string;
nodeCount: number;
file: ForestFile;
}

export function isSubtreeReference(x: TreeFile['tree']): x is Extract<TreeFile['tree'], { fileName: string }> {
return !!(x as any).fileName;
}
Expand Down
Loading