Merge pull request #327 from PrefectHQ/layout-worker-spawning

Terminate workers when not in use
PrefectHQ · Nov 9, 2023 · b060e7d · b060e7d
2 parents e673dd2 + 40c79ae
commit b060e7d
Show file tree

Hide file tree

Showing 2 changed files with 29 additions and 3 deletions.
diff --git a/src/factories/nodeFlowRun.ts b/src/factories/nodeFlowRun.ts
@@ -18,7 +18,7 @@ export async function flowRunContainerFactory(node: RunGraphNode) {
   const config = await waitForConfig()
   const { element: bar, render: renderBar } = await nodeBarFactory()
   const { element: label, render: renderLabelText } = await nodeLabelFactory()
-  const { element: nodesContainer, render: renderNodes, getSize: getNodesSize } = await nodesContainerFactory()
+  const { element: nodesContainer, render: renderNodes, getSize: getNodesSize, stopWorker: stopNodesWorker } = await nodesContainerFactory()
   const { element: arrowButton, render: renderArrowButtonContainer } = await nodeArrowButtonFactory()
   const { element: border, render: renderBorderContainer } = await borderFactory()
   const { start: startData, stop: stopData } = await dataFactory(node.id, data => {
@@ -103,6 +103,7 @@ export async function flowRunContainerFactory(node: RunGraphNode) {
   async function close(): Promise<void> {
     isOpen = false
     container.removeChild(nodesContainer)
+    stopNodesWorker()
 
     await Promise.all([
       stopData(),

diff --git a/src/factories/nodes.ts b/src/factories/nodes.ts
@@ -13,7 +13,7 @@ import { emitter } from '@/objects/events'
 import { getSelected } from '@/objects/selection'
 import { getHorizontalColumnSize, layout, waitForSettings } from '@/objects/settings'
 import { exhaustive } from '@/utilities/exhaustive'
-import { WorkerLayoutMessage, WorkerMessage, layoutWorkerFactory } from '@/workers/runGraph'
+import { IRunGraphWorker, WorkerLayoutMessage, WorkerMessage, layoutWorkerFactory } from '@/workers/runGraph'
 
 // parentId-childId
 type EdgeKey = `${string}_${string}`
@@ -22,13 +22,14 @@ export type NodesContainer = Awaited<ReturnType<typeof nodesContainerFactory>>
 
 // eslint-disable-next-line @typescript-eslint/explicit-function-return-type
 export async function nodesContainerFactory() {
-  const worker = layoutWorkerFactory(onmessage)
   const nodes = new Map<string, NodeContainerFactory>()
   const edges = new Map<EdgeKey, EdgeFactory>()
   const container = new Container()
   const edgesContainer = new Container()
   const config = await waitForConfig()
 
+  let worker: IRunGraphWorker | null = null
+
   // used for both vertical layouts
   const rows = offsetsFactory({
     gap: () => config.styles.rowGap,
@@ -64,6 +65,8 @@ export async function nodesContainerFactory() {
   })
 
   async function render(data: RunGraphData): Promise<void> {
+    startWorker()
+
     runData = data
     nodesLayout = null
 
@@ -75,6 +78,23 @@ export async function nodesContainerFactory() {
     getLayout(data)
   }
 
+  function startWorker(): void {
+    if (worker) {
+      return
+    }
+
+    worker = layoutWorkerFactory(onmessage)
+  }
+
+  function stopWorker(): void {
+    if (!worker) {
+      return
+    }
+
+    worker.terminate()
+    worker = null
+  }
+
   async function createNodes(data: RunGraphData): Promise<void> {
     const promises: Promise<Container>[] = []
 
@@ -133,6 +153,10 @@ export async function nodesContainerFactory() {
   }
 
   function getLayout(data: RunGraphData): void {
+    if (!worker) {
+      throw new Error('Layout worker not initialized')
+    }
+
     const widths: NodeWidths = new Map()
 
     for (const [nodeId, { element }] of nodes) {
@@ -388,6 +412,7 @@ export async function nodesContainerFactory() {
 
   return {
     element: container,
+    stopWorker,
     getSize,
     render,
   }