From a6f0cfa468ca89b174dddd4b4b8ba3598b35fc9e Mon Sep 17 00:00:00 2001
From: Joyee Cheung <joyeec9h3@gmail.com>
Date: Sat, 7 Dec 2024 18:25:02 +0100
Subject: [PATCH] sea: only assert snapshot main function for main threads

Snapshot main functions are only loaded for main threads in single
executable applications. Update the check to avoid asserting it
in worker threads - this allows worker threads to be spawned in
snapshot main functions bundled into a single executable
application.

PR-URL: https://github.com/nodejs/node/pull/56120
Fixes: https://github.com/nodejs/node/issues/56077
Reviewed-By: Anna Henningsen <anna@addaleax.net>
Reviewed-By: Chengzhong Wu <legendecas@gmail.com>
---
 src/node.cc                                   |  6 +-
 ...-executable-application-snapshot-worker.js | 80 +++++++++++++++++++
 2 files changed, 85 insertions(+), 1 deletion(-)
 create mode 100644 test/sequential/test-single-executable-application-snapshot-worker.js

diff --git a/src/node.cc b/src/node.cc
index dd52fbffac0dee..480681d0b02ff8 100644
--- a/src/node.cc
+++ b/src/node.cc
@@ -320,7 +320,8 @@ MaybeLocal<Value> StartExecution(Environment* env, StartExecutionCallback cb) {
   CHECK(!env->isolate_data()->is_building_snapshot());
 
 #ifndef DISABLE_SINGLE_EXECUTABLE_APPLICATION
-  if (sea::IsSingleExecutable()) {
+  // Snapshot in SEA is only loaded for the main thread.
+  if (sea::IsSingleExecutable() && env->is_main_thread()) {
     sea::SeaResource sea = sea::FindSingleExecutableResource();
     // The SEA preparation blob building process should already enforce this,
     // this check is just here to guard against the unlikely case where
@@ -342,6 +343,9 @@ MaybeLocal<Value> StartExecution(Environment* env, StartExecutionCallback cb) {
   // move the pre-execution part into a different file that can be
   // reused when dealing with user-defined main functions.
   if (!env->snapshot_deserialize_main().IsEmpty()) {
+    // Custom worker snapshot is not supported yet,
+    // so workers can't have deserialize main functions.
+    CHECK(env->is_main_thread());
     return env->RunSnapshotDeserializeMain();
   }
 
diff --git a/test/sequential/test-single-executable-application-snapshot-worker.js b/test/sequential/test-single-executable-application-snapshot-worker.js
new file mode 100644
index 00000000000000..50c77743573a44
--- /dev/null
+++ b/test/sequential/test-single-executable-application-snapshot-worker.js
@@ -0,0 +1,80 @@
+'use strict';
+
+require('../common');
+
+const {
+  generateSEA,
+  skipIfSingleExecutableIsNotSupported,
+} = require('../common/sea');
+
+skipIfSingleExecutableIsNotSupported();
+
+// This tests the snapshot support in single executable applications.
+
+const tmpdir = require('../common/tmpdir');
+const { writeFileSync, existsSync } = require('fs');
+const {
+  spawnSyncAndAssert, spawnSyncAndExitWithoutError,
+} = require('../common/child_process');
+const assert = require('assert');
+
+const configFile = tmpdir.resolve('sea-config.json');
+const seaPrepBlob = tmpdir.resolve('sea-prep.blob');
+const outputFile = tmpdir.resolve(process.platform === 'win32' ? 'sea.exe' : 'sea');
+
+{
+  tmpdir.refresh();
+
+  // FIXME(joyeecheung): currently `worker_threads` cannot be loaded during the
+  // snapshot building process because internal/worker.js is accessing isMainThread at
+  // the top level (and there are maybe more code that access these at the top-level),
+  // and have to be loaded in the deserialized snapshot main function.
+  // Change these states to be accessed on-demand.
+  const code = `
+  const {
+    setDeserializeMainFunction,
+  } = require('v8').startupSnapshot;
+  setDeserializeMainFunction(() => {
+    const { Worker } = require('worker_threads');
+    new Worker("console.log('Hello from Worker')", { eval: true });
+  });
+  `;
+
+  writeFileSync(tmpdir.resolve('snapshot.js'), code, 'utf-8');
+  writeFileSync(configFile, `
+  {
+    "main": "snapshot.js",
+    "output": "sea-prep.blob",
+    "useSnapshot": true
+  }
+  `);
+
+  spawnSyncAndExitWithoutError(
+    process.execPath,
+    ['--experimental-sea-config', 'sea-config.json'],
+    {
+      cwd: tmpdir.path,
+      env: {
+        NODE_DEBUG_NATIVE: 'SEA',
+        ...process.env,
+      },
+    });
+
+  assert(existsSync(seaPrepBlob));
+
+  generateSEA(outputFile, process.execPath, seaPrepBlob);
+
+  spawnSyncAndAssert(
+    outputFile,
+    {
+      env: {
+        NODE_DEBUG_NATIVE: 'SEA',
+        ...process.env,
+      }
+    },
+    {
+      trim: true,
+      stdout: 'Hello from Worker'
+    }
+  );
+}