Add support for MIX_PROFILE_FLAGS for more configuration around task profiling

Author: josevalim

lib/mix/lib/mix.ex

@@ -360,7 +360,9 @@ defmodule Mix do
     * `MIX_PATH` - appends extra code paths
 
     * `MIX_PROFILE` - a list of comma-separated Mix tasks to profile the time spent on
-      functions by the process running the task
+      functions by the process running the task, such as `MIX_PROFILE=compile,test`.
+      You can also set `MIX_PROFILE_FLAGS` to control the flags given to profiling,
+      see `mix profile.tprof` for all options. By default, it uses `--no-warmup`.
 
     * `MIX_QUIET` - does not print information messages to the terminal
 

lib/mix/lib/mix/cli.ex

@@ -22,7 +22,15 @@ defmodule Mix.CLI do
     if env_variable_activated?("MIX_QUIET"), do: Mix.shell(Mix.Shell.Quiet)
 
     if profile = System.get_env("MIX_PROFILE") do
-      Mix.State.put(:profile, String.split(profile, ","))
+      flags = System.get_env("MIX_PROFILE_FLAGS", "")
+      {opts, args} = Mix.Tasks.Profile.Tprof.parse!(OptionParser.split(flags))
+
+      if args != [] do
+        Mix.raise("Invalid arguments given to MIX_PROFILE_FLAGS: #{inspect(args)}")
+      end
+
+      opts = Keyword.put_new(opts, :warmup, false)
+      Mix.State.put(:profile, {opts, String.split(profile, ",")})
     end
 
     case check_for_shortcuts(args) do

lib/mix/lib/mix/task.ex

@@ -520,16 +520,24 @@ defmodule Mix.Task do
         shell.info(["<- Ran mix ", task, " in ", Integer.to_string(div(time, 1000)), "ms"])
         res
 
-      task in Mix.State.get(:profile, []) ->
+      opts = profile_opts_for(task) ->
         shell = Mix.shell()
         shell.info(["-> Profiling mix ", task_to_string(task, args), project_to_string(proj)])
-        Mix.Tasks.Profile.Tprof.profile(fun, warmup: false, set_on_spawn: false)
+        Mix.Tasks.Profile.Tprof.profile(fun, opts)
 
       true ->
         fun.()
     end
   end
 
+  defp profile_opts_for(task) do
+    {opts, tasks} = Mix.State.get(:profile, {[], []})
+
+    if task in tasks do
+      opts
+    end
+  end
+
   defp project_to_string(nil), do: ""
   defp project_to_string(proj), do: " (inside #{inspect(proj)})"
 

lib/mix/lib/mix/tasks/profile.tprof.ex

@@ -56,6 +56,7 @@ defmodule Mix.Tasks.Profile.Tprof do
     * `--no-deps-check` - does not check dependencies
     * `--no-archives-check` - does not check archives
     * `--no-halt` - does not halt the system after running the command
+    * `--no-set-on-spawn` - does not profile spawned processes
     * `--no-start` - does not start applications after compilation
     * `--no-elixir-version-check` - does not check the Elixir version from mix.exs
 
@@ -159,7 +160,8 @@ defmodule Mix.Tasks.Profile.Tprof do
     archives_check: :boolean,
     warmup: :boolean,
     elixir_version_check: :boolean,
-    parallel_require: :keep
+    parallel_require: :keep,
+    set_on_spawn: :boolean
   ]
 
   @aliases [
@@ -171,7 +173,7 @@ defmodule Mix.Tasks.Profile.Tprof do
 
   @impl true
   def run(args) do
-    {opts, head} = OptionParser.parse_head!(args, aliases: @aliases, strict: @switches)
+    {opts, head} = parse!(args)
     Mix.Task.reenable("profile.tprof")
 
     Mix.Tasks.Run.run(
@@ -183,9 +185,13 @@ defmodule Mix.Tasks.Profile.Tprof do
     )
   end
 
-  defp profile_code(code_string, opts) do
-    opts = Enum.map(opts, &parse_opt/1)
+  @doc false
+  def parse!(args) do
+    {opts, args} = OptionParser.parse!(args, aliases: @aliases, strict: @switches)
+    {Enum.map(opts, &parse_opt/1), args}
+  end
 
+  defp profile_code(code_string, opts) do
     content =
       quote do
         unquote(__MODULE__).profile(