src: provide a variant of LoadEnvironment taking a callback

This allows embedders to flexibly control how they start JS code
rather than using `third_party_main`.

Backport-PR-URL: #35241
PR-URL: #30467
Reviewed-By: James M Snell <jasnell@gmail.com>
Reviewed-By: Gireesh Punathil <gpunathi@in.ibm.com>

src/api/environment.cc

@@ -409,11 +409,12 @@ NODE_EXTERN std::unique_ptr<InspectorParentHandle> GetInspectorParentHandle(
 }
 
 void LoadEnvironment(Environment* env) {
-  USE(LoadEnvironment(env, {}));
+  USE(LoadEnvironment(env, nullptr, {}));
 }
 
 MaybeLocal<Value> LoadEnvironment(
     Environment* env,
+    StartExecutionCallback cb,
     std::unique_ptr<InspectorParentHandle> inspector_parent_handle) {
   env->InitializeLibuv(per_process::v8_is_profiling);
   env->InitializeDiagnostics();
@@ -428,9 +429,7 @@ MaybeLocal<Value> LoadEnvironment(
   }
 #endif
 
-  // TODO(joyeecheung): Allow embedders to customize the entry
-  // point more directly without using _third_party_main.js
-  return StartExecution(env);
+  return StartExecution(env, cb);
 }
 
 Environment* GetCurrentEnvironment(Local<Context> context) {

src/node.cc

@@ -364,6 +364,7 @@ void MarkBootstrapComplete(const FunctionCallbackInfo<Value>& args) {
       performance::NODE_PERFORMANCE_MILESTONE_BOOTSTRAP_COMPLETE);
 }
 
+static
 MaybeLocal<Value> StartExecution(Environment* env, const char* main_script_id) {
   EscapableHandleScope scope(env->isolate());
   CHECK_NOT_NULL(main_script_id);
@@ -384,17 +385,31 @@ MaybeLocal<Value> StartExecution(Environment* env, const char* main_script_id) {
           ->GetFunction(env->context())
           .ToLocalChecked()};
 
-  InternalCallbackScope callback_scope(
-    env,
-    Object::New(env->isolate()),
-    { 1, 0 },
-    InternalCallbackScope::kSkipAsyncHooks);
-
   return scope.EscapeMaybe(
       ExecuteBootstrapper(env, main_script_id, &parameters, &arguments));
 }
 
-MaybeLocal<Value> StartExecution(Environment* env) {
+MaybeLocal<Value> StartExecution(Environment* env, StartExecutionCallback cb) {
+  InternalCallbackScope callback_scope(
+      env,
+      Object::New(env->isolate()),
+      { 1, 0 },
+      InternalCallbackScope::kSkipAsyncHooks);
+
+  if (cb != nullptr) {
+    EscapableHandleScope scope(env->isolate());
+
+    if (StartExecution(env, "internal/bootstrap/environment").IsEmpty())
+      return {};
+
+    StartExecutionCallbackInfo info = {
+      env->process_object(),
+      env->native_module_require(),
+    };
+
+    return scope.EscapeMaybe(cb(info));
+  }
+
   // To allow people to extend Node in different ways, this hook allows
   // one to drop a file lib/_third_party_main.js into the build
   // directory which will be executed instead of Node's normal loading.

src/node.h

@@ -73,6 +73,7 @@
 #include "node_version.h"  // NODE_MODULE_VERSION
 
 #include <memory>
+#include <functional>
 
 #define NODE_MAKE_VERSION(major, minor, patch)                                \
   ((major) * 0x1000 + (minor) * 0x100 + (patch))
@@ -417,12 +418,20 @@ NODE_EXTERN std::unique_ptr<InspectorParentHandle> GetInspectorParentHandle(
     ThreadId child_thread_id,
     const char* child_url);
 
-// TODO(addaleax): Deprecate this in favour of the MaybeLocal<> overload
-// and provide a more flexible approach than third_party_main.
+struct StartExecutionCallbackInfo {
+  v8::Local<v8::Object> process_object;
+  v8::Local<v8::Function> native_require;
+};
+
+using StartExecutionCallback =
+    std::function<v8::MaybeLocal<v8::Value>(const StartExecutionCallbackInfo&)>;
+
+// TODO(addaleax): Deprecate this in favour of the MaybeLocal<> overload.
 NODE_EXTERN void LoadEnvironment(Environment* env);
 NODE_EXTERN v8::MaybeLocal<v8::Value> LoadEnvironment(
     Environment* env,
-    std::unique_ptr<InspectorParentHandle> inspector_parent_handle);
+    StartExecutionCallback cb,
+    std::unique_ptr<InspectorParentHandle> inspector_parent_handle = {});
 NODE_EXTERN void FreeEnvironment(Environment* env);
 
 // This may return nullptr if context is not associated with a Node instance.

src/node_internals.h

@@ -301,9 +301,10 @@ void DefineZlibConstants(v8::Local<v8::Object> target);
 v8::Isolate* NewIsolate(v8::Isolate::CreateParams* params,
                         uv_loop_t* event_loop,
                         MultiIsolatePlatform* platform);
-v8::MaybeLocal<v8::Value> StartExecution(Environment* env);
+// This overload automatically picks the right 'main_script_id' if no callback
+// was provided by the embedder.
 v8::MaybeLocal<v8::Value> StartExecution(Environment* env,
-                                         const char* main_script_id);
+                                         StartExecutionCallback cb = nullptr);
 v8::MaybeLocal<v8::Object> GetPerContextExports(v8::Local<v8::Context> context);
 v8::MaybeLocal<v8::Value> ExecuteBootstrapper(
     Environment* env,

src/node_worker.cc

@@ -333,6 +333,7 @@ void Worker::Run() {
         CreateEnvMessagePort(env_.get());
         Debug(this, "Created message port for worker %llu", thread_id_.id);
         if (LoadEnvironment(env_.get(),
+                            nullptr,
                             std::move(inspector_parent_handle_))
                 .IsEmpty()) {
           return;

test/cctest/test_environment.cc

@@ -51,6 +51,35 @@ class EnvironmentTest : public EnvironmentTestFixture {
 //   CHECK(result->IsString());
 // }
 
+TEST_F(EnvironmentTest, LoadEnvironmentWithCallback) {
+  const v8::HandleScope handle_scope(isolate_);
+  const Argv argv;
+  Env env {handle_scope, argv};
+
+  v8::Local<v8::Context> context = isolate_->GetCurrentContext();
+  bool called_cb = false;
+  node::LoadEnvironment(*env,
+                        [&](const node::StartExecutionCallbackInfo& info)
+                            -> v8::MaybeLocal<v8::Value> {
+    called_cb = true;
+
+    CHECK(info.process_object->IsObject());
+    CHECK(info.native_require->IsFunction());
+
+    v8::Local<v8::Value> argv0 = info.process_object->Get(
+        context,
+        v8::String::NewFromOneByte(
+            isolate_,
+            reinterpret_cast<const uint8_t*>("argv0"),
+            v8::NewStringType::kNormal).ToLocalChecked()).ToLocalChecked();
+    CHECK(argv0->IsString());
+
+    return info.process_object;
+  });
+
+  CHECK(called_cb);
+}
+
 TEST_F(EnvironmentTest, AtExitWithEnvironment) {
   const v8::HandleScope handle_scope(isolate_);
   const Argv argv;
@@ -188,6 +217,35 @@ static void at_exit_js(void* arg) {
   called_at_exit_js = true;
 }
 
+TEST_F(EnvironmentTest, SetImmediateCleanup) {
+  int called = 0;
+  int called_unref = 0;
+
+  {
+    const v8::HandleScope handle_scope(isolate_);
+    const Argv argv;
+    Env env {handle_scope, argv};
+
+    node::LoadEnvironment(*env,
+                          [&](const node::StartExecutionCallbackInfo& info)
+                              -> v8::MaybeLocal<v8::Value> {
+      return v8::Object::New(isolate_);
+    });
+
+    (*env)->SetImmediate([&](node::Environment* env_arg) {
+      EXPECT_EQ(env_arg, *env);
+      called++;
+    });
+    (*env)->SetUnrefImmediate([&](node::Environment* env_arg) {
+      EXPECT_EQ(env_arg, *env);
+      called_unref++;
+    });
+  }
+
+  EXPECT_EQ(called, 1);
+  EXPECT_EQ(called_unref, 0);
+}
+
 static char hello[] = "hello";
 
 TEST_F(EnvironmentTest, BufferWithFreeCallbackIsDetached) {