Ensure mix escript.install works with deps.partition, closes #14885

Author: josevalim

lib/mix/lib/mix/cli.ex

@@ -8,18 +8,18 @@ defmodule Mix.CLI do
   @doc """
   Runs Mix according to the command line arguments.
   """
-  def main(args \\ System.argv()) do
+  def main(args \\ System.argv(), mix_exs \\ System.get_env("MIX_EXS") || "mix.exs") do
     if env_variable_activated?("MIX_DEBUG") do
       IO.puts("-> Running mix CLI")
-      {time, res} = :timer.tc(&main/2, [args, true])
+      {time, res} = :timer.tc(&main/3, [args, mix_exs, true])
       IO.puts(["<- Ran mix CLI in ", Integer.to_string(div(time, 1000)), "ms"])
       res
     else
-      main(args, false)
+      main(args, mix_exs, false)
     end
   end
 
-  defp main(args, debug?) do
+  defp main(args, mix_exs, debug?) do
     Mix.start()
 
     if debug?, do: Mix.debug(true)
@@ -46,28 +46,26 @@ defmodule Mix.CLI do
         display_version()
 
       nil ->
-        proceed(args)
+        proceed(args, mix_exs)
     end
   end
 
-  defp proceed(args) do
+  defp proceed(args, mix_exs) do
     load_dot_config()
-    load_mix_exs(args)
+    load_mix_exs(args, mix_exs)
     project = Mix.Project.get()
     {task, args} = get_task(args, project)
     ensure_hex(task)
     maybe_change_env_and_target(task, project)
     run_task(task, args)
   end
 
-  defp load_mix_exs(args) do
-    file = System.get_env("MIX_EXS") || "mix.exs"
-
-    if File.regular?(file) do
+  defp load_mix_exs(args, mix_exs) do
+    if is_binary(mix_exs) and File.regular?(mix_exs) do
       Mix.ProjectStack.post_config(state_loader: {:cli, List.first(args)})
       old_undefined = Code.get_compiler_option(:no_warn_undefined)
       Code.put_compiler_option(:no_warn_undefined, :all)
-      Code.compile_file(file)
+      Code.compile_file(mix_exs)
       Code.put_compiler_option(:no_warn_undefined, old_undefined)
     end
   end

lib/mix/lib/mix/tasks/deps.partition.ex

@@ -7,8 +7,6 @@ defmodule Mix.Tasks.Deps.Partition do
 
   ## Server
 
-  @deps_partition_install_mix_exs ~c"deps.partition.mix.exs"
-
   def server(deps, count, force?) do
     {:ok, socket} = :gen_tcp.listen(0, [:binary, packet: :line, active: false])
 
@@ -29,41 +27,18 @@ defmodule Mix.Tasks.Deps.Partition do
     ansi_flag = if IO.ANSI.enabled?(), do: ~c"--color", else: ~c"--no-color"
     force_flag = if force?, do: ~c"--force", else: ~c"--no-force"
 
-    env_vars =
-      if Mix.install?() do
-        blob =
-          Mix.Project.config()
-          |> :erlang.term_to_binary()
-          |> :binary.bin_to_list()
-          |> Enum.join(",")
-
-        # We replicate the initialization logic from Mix.install/2 as part of mix.exs
-        File.write!(@deps_partition_install_mix_exs, """
-        config = <<#{blob}>>
-        project = :erlang.binary_to_term(config)
-
-        if compile_config = project[:compile_config] do
-          Application.put_all_env(compile_config, persistent: true)
-        end
-
-        Mix.ProjectStack.push(Mix.InstallProject, project, "nofile")
-        """)
-
-        [{~c"MIX_EXS", @deps_partition_install_mix_exs}]
-      else
-        []
-      end
-
     args = [
       ansi_flag,
       ~c"-e",
-      ~c"Mix.CLI.main",
-      ~c"deps.partition",
+      ~c"Mix.Tasks.Deps.Partition.client",
+      ~c"--",
       force_flag,
       ~c"--port",
       Integer.to_charlist(port),
       ~c"--host",
-      ~c"127.0.0.1"
+      ~c"127.0.0.1",
+      ~c"--config",
+      Mix.ProjectStack.peek() |> :erlang.term_to_binary() |> Base.url_encode64()
     ]
 
     options = [
@@ -77,7 +52,6 @@ defmodule Mix.Tasks.Deps.Partition do
         {~c"MIX_OS_CONCURRENCY_LOCK", ~c"false"},
         {~c"MIX_ENV", Atom.to_charlist(Mix.env())},
         {~c"MIX_TARGET", Atom.to_charlist(Mix.target())}
-        | env_vars
       ]
     ]
 
@@ -241,16 +215,36 @@ defmodule Mix.Tasks.Deps.Partition do
 
   ## Client
 
-  @switches [port: :integer, host: :string, force: :boolean, index: :string]
+  @switches [port: :integer, host: :string, force: :boolean, index: :string, config: :string]
 
-  @impl true
-  def run(args) do
+  def client do
     # If stdin closes, we shutdown the VM
     spawn(fn ->
       _ = IO.gets("")
       System.halt(0)
     end)
 
+    args = System.argv()
+    {opts, []} = OptionParser.parse!(args, strict: @switches)
+    peek = Keyword.fetch!(opts, :config) |> Base.url_decode64!() |> :erlang.binary_to_term()
+
+    # This is specific to Mix.install/2 and how it handles compile-time config
+    if compile_config = peek.config[:compile_config] do
+      Application.put_all_env(compile_config, persistent: true)
+    end
+
+    partition_args =
+      opts
+      |> Keyword.take([:host, :port, :index, :force])
+      |> OptionParser.to_argv()
+
+    Mix.start()
+    Mix.ProjectStack.push(peek.name, peek.config, peek.file)
+    Mix.CLI.main(["deps.partition" | partition_args], nil)
+  end
+
+  @impl true
+  def run(args) do
     {opts, []} = OptionParser.parse!(args, strict: @switches)
     host = Keyword.fetch!(opts, :host)
     port = Keyword.fetch!(opts, :port)

lib/mix/test/mix/tasks/escript_test.exs

@@ -341,48 +341,67 @@ defmodule Mix.Tasks.EscriptTest do
     end)
   end
 
-  test "escript.install from Git" do
-    in_fixture("git_repo", fn ->
-      File.mkdir_p!("config")
-
-      File.write!("config/config.exs", """
-      import Config
-      config :git_repo, :escript_config, true
-      """)
-
-      File.write!("lib/git_repo.ex", """
-      require Application
-      true = Application.compile_env!(:git_repo, :escript_config)
-
-      defmodule GitRepo.Escript do
-        def main(_argv) do
-          IO.puts("TEST")
+  for count <- [1, 2] do
+    test "escript.install from Git (with MIX_OS_DEPS_COMPILE_PARTITION_COUNT=#{count})" do
+      System.put_env("MIX_OS_DEPS_COMPILE_PARTITION_COUNT", "#{unquote(count)}")
+
+      in_fixture("deps_status", fn ->
+        File.mkdir_p!("config")
+        File.mkdir_p!("lib")
+
+        File.write!("config/config.exs", """
+        import Config
+        config :source_repo, :escript_config, true
+        """)
+
+        File.write!("lib/source_repo.ex", """
+        require Application
+        true = Application.compile_env!(:source_repo, :escript_config)
+
+        defmodule SourceRepo.Escript do
+          def main(_argv) do
+            IO.puts("TEST")
+          end
         end
-      end
-      """)
-
-      File.write!("mix.exs", """
-      defmodule GitRepo.MixProject do
-        use Mix.Project
-
-        def project do
-          [app: :git_repo, version: "0.1.0", escript: [main_module: GitRepo.Escript]]
+        """)
+
+        File.write!("mix.exs", """
+        defmodule SourceRepo.MixProject do
+          use Mix.Project
+
+          def project do
+            [
+              app: :source_repo,
+              version: "0.1.0",
+              escript: [main_module: SourceRepo.Escript],
+              deps: [
+                {:git_repo, path: "#{MixTest.Case.fixture_path("git_repo")}"},
+                {:git_rebar, path: "#{MixTest.Case.fixture_path("git_rebar")}", manager: :rebar3},
+                {:ok, path: "deps/ok"},
+              ]
+            ]
+          end
         end
-      end
-      """)
+        """)
 
-      System.cmd("git", ~w[add .])
-      System.cmd("git", ~w[commit -m "ok"])
+        System.cmd("git", ~w[init])
+        System.cmd("git", ~w[add .])
+        System.cmd("git", ~w[commit -m "ok"])
 
-      send(self(), {:mix_shell_input, :yes?, true})
-      Mix.Tasks.Escript.Install.run(["git", File.cwd!()])
-      assert_received {:mix_shell, :info, ["Generated escript git_repo with MIX_ENV=prod"]}
+        ExUnit.CaptureIO.capture_io(fn ->
+          send(self(), {:mix_shell_input, :yes?, true})
+          Mix.Tasks.Escript.Install.run(["git", File.cwd!()])
+        end)
 
-      escript_path = Path.join([tmp_path(".mix"), "escripts", "git_repo"])
-      assert System.cmd("escript", [escript_path]) == {"TEST\n", 0}
-    end)
-  after
-    purge([GitRepo.Escript, GitRepo.MixProject, Mix.Local.Installer.MixProject])
+        assert_received {:mix_shell, :info, ["Generated escript source_repo with MIX_ENV=prod"]}
+
+        escript_path = Path.join([tmp_path(".mix"), "escripts", "source_repo"])
+        assert System.cmd("escript", [escript_path]) == {"TEST\n", 0}
+      end)
+    after
+      System.delete_env("MIX_OS_DEPS_COMPILE_PARTITION_COUNT")
+      purge([SourceRepo.Escript, SourceRepo.MixProject, Mix.Local.Installer.MixProject])
+    end
   end
 
   defp push_project_with_config(module, config \\ []) do