Add $callers tracking to Task (#8510)

This allows us to track which process invoked the Task
besides its ancestors in the supervision tree.

Closes #7995

Author: josevalim

lib/elixir/lib/task.ex

@@ -271,9 +271,8 @@ defmodule Task do
   @spec start_link(module, atom, [term]) :: {:ok, pid}
   def start_link(module, function_name, args)
       when is_atom(module) and is_atom(function_name) and is_list(args) do
-    self()
-    |> get_owner()
-    |> Task.Supervised.start_link({module, function_name, args})
+    mfa = {module, function_name, args}
+    Task.Supervised.start_link(get_owner(self()), get_callers(self()), mfa)
   end
 
   @doc """
@@ -300,9 +299,8 @@ defmodule Task do
   @spec start(module, atom, [term]) :: {:ok, pid}
   def start(module, function_name, args)
       when is_atom(module) and is_atom(function_name) and is_list(args) do
-    self()
-    |> get_owner()
-    |> Task.Supervised.start({module, function_name, args})
+    mfa = {module, function_name, args}
+    Task.Supervised.start(get_owner(self()), get_callers(self()), mfa)
   end
 
   @doc """
@@ -391,7 +389,7 @@ defmodule Task do
       when is_atom(module) and is_atom(function_name) and is_list(args) do
     mfa = {module, function_name, args}
     owner = self()
-    {:ok, pid} = Task.Supervised.start_link(get_owner(owner), :nomonitor, mfa)
+    {:ok, pid} = Task.Supervised.start_link(get_owner(owner), get_callers(owner), :nomonitor, mfa)
     ref = Process.monitor(pid)
     send(pid, {owner, ref})
     %Task{pid: pid, ref: ref, owner: owner}
@@ -498,8 +496,8 @@ defmodule Task do
   end
 
   defp build_stream(enumerable, fun, options) do
-    &Task.Supervised.stream(enumerable, &1, &2, fun, options, fn owner, mfa ->
-      {:ok, pid} = Task.Supervised.start_link(get_owner(owner), :nomonitor, mfa)
+    &Task.Supervised.stream(enumerable, &1, &2, fun, options, fn [owner | _] = callers, mfa ->
+      {:ok, pid} = Task.Supervised.start_link(get_owner(owner), callers, :nomonitor, mfa)
       {:ok, :link, pid}
     end)
   end
@@ -517,6 +515,13 @@ defmodule Task do
     {node(), self_or_name, pid}
   end
 
+  defp get_callers(owner) do
+    case :erlang.get(:"$callers") do
+      [_ | _] = list -> [owner | list]
+      _ -> [owner]
+    end
+  end
+
   @doc """
   Awaits a task reply and returns it.
 

lib/elixir/lib/task/supervised.ex

@@ -2,20 +2,21 @@ defmodule Task.Supervised do
   @moduledoc false
   @ref_timeout 5000
 
-  def start(owner, fun) do
-    {:ok, :proc_lib.spawn(__MODULE__, :noreply, [owner, fun])}
+  def start(owner, callers, fun) do
+    {:ok, :proc_lib.spawn(__MODULE__, :noreply, [owner, callers, fun])}
   end
 
-  def start_link(owner, fun) do
-    {:ok, :proc_lib.spawn_link(__MODULE__, :noreply, [owner, fun])}
+  def start_link(owner, callers, fun) do
+    {:ok, :proc_lib.spawn_link(__MODULE__, :noreply, [owner, callers, fun])}
   end
 
-  def start_link(owner, monitor, fun) do
-    {:ok, :proc_lib.spawn_link(__MODULE__, :reply, [owner, monitor, fun])}
+  def start_link(owner, callers, monitor, fun) do
+    {:ok, :proc_lib.spawn_link(__MODULE__, :reply, [owner, callers, monitor, fun])}
   end
 
-  def reply({_, _, owner_pid} = owner, monitor, mfa) do
+  def reply({_, _, owner_pid} = owner, callers, monitor, mfa) do
     initial_call(mfa)
+    put_callers(callers)
 
     case monitor do
       :monitor ->
@@ -31,7 +32,7 @@ defmodule Task.Supervised do
     receive do
       {^owner_pid, ref} ->
         _ = if mref, do: Process.demonitor(mref, [:flush])
-        send(owner_pid, {ref, do_apply(owner, mfa)})
+        send(owner_pid, {ref, invoke_mfa(owner, mfa)})
 
       {:DOWN, ^mref, _, _, reason} when is_reference(mref) ->
         exit({:shutdown, reason})
@@ -60,9 +61,14 @@ defmodule Task.Supervised do
     end
   end
 
-  def noreply(owner, mfa) do
+  def noreply(owner, callers, mfa) do
     initial_call(mfa)
-    do_apply(owner, mfa)
+    put_callers(callers)
+    invoke_mfa(owner, mfa)
+  end
+
+  defp put_callers(callers) do
+    Process.put(:"$callers", callers)
   end
 
   defp initial_call(mfa) do
@@ -79,7 +85,7 @@ defmodule Task.Supervised do
     {mod, fun, length(args)}
   end
 
-  defp do_apply(owner, {module, fun, args} = mfa) do
+  defp invoke_mfa(owner, {module, fun, args} = mfa) do
     try do
       apply(module, fun, args)
     catch
@@ -147,6 +153,7 @@ defmodule Task.Supervised do
     timeout = Keyword.get(options, :timeout, 5000)
     on_timeout = Keyword.get(options, :on_timeout, :exit)
     parent = self()
+    callers = get_callers()
 
     {:trap_exit, trap_exit?} = Process.info(self(), :trap_exit)
 
@@ -156,7 +163,10 @@ defmodule Task.Supervised do
     spawn_opts = [:link, :monitor]
 
     {monitor_pid, monitor_ref} =
-      Process.spawn(fn -> stream_monitor(parent, mfa, spawn, trap_exit?, timeout) end, spawn_opts)
+      Process.spawn(
+        fn -> stream_monitor(callers, mfa, spawn, trap_exit?, timeout) end,
+        spawn_opts
+      )
 
     # Now that we have the pid of the "monitor" process and the reference of the
     # monitor we use to monitor such process, we can inform the monitor process
@@ -183,6 +193,13 @@ defmodule Task.Supervised do
     )
   end
 
+  defp get_callers do
+    case :erlang.get(:"$callers") do
+      [_ | _] = list -> [self() | list]
+      _ -> [self()]
+    end
+  end
+
   defp stream_reduce({:halt, acc}, _max, _spawned, _delivered, _waiting, next, config) do
     %{monitor_pid: monitor_pid, monitor_ref: monitor_ref, timeout: timeout} = config
     stream_close(monitor_pid, monitor_ref, timeout)
@@ -426,9 +443,8 @@ defmodule Task.Supervised do
     end
   end
 
-  defp stream_monitor(parent_pid, mfa, spawn, trap_exit?, timeout) do
+  defp stream_monitor([parent_pid | _] = callers, mfa, spawn, trap_exit?, timeout) do
     Process.flag(:trap_exit, trap_exit?)
-
     parent_ref = Process.monitor(parent_pid)
 
     # Let's wait for the parent process to tell this process the monitor ref
@@ -437,7 +453,7 @@ defmodule Task.Supervised do
     receive do
       {^parent_pid, monitor_ref} ->
         config = %{
-          parent_pid: parent_pid,
+          callers: callers,
           parent_ref: parent_ref,
           mfa: mfa,
           spawn: spawn,
@@ -454,7 +470,7 @@ defmodule Task.Supervised do
 
   defp stream_monitor_loop(running_tasks, config) do
     %{
-      parent_pid: parent_pid,
+      callers: [parent_pid | _] = callers,
       mfa: mfa,
       spawn: spawn,
       monitor_ref: monitor_ref,
@@ -465,7 +481,7 @@ defmodule Task.Supervised do
       # The parent process is telling us to spawn a new task to process
       # "value". We spawn it and notify the parent about its pid.
       {:spawn, position, value} ->
-        case spawn.(parent_pid, normalize_mfa_with_arg(mfa, value)) do
+        case spawn.(callers, normalize_mfa_with_arg(mfa, value)) do
           {:ok, type, pid} ->
             ref = Process.monitor(pid)
 

lib/elixir/lib/task/supervisor.ex

@@ -395,7 +395,7 @@ defmodule Task.Supervisor do
   def start_child(supervisor, fun, options \\ []) do
     restart = options[:restart]
     shutdown = options[:shutdown]
-    args = [get_owner(self()), {:erlang, :apply, [fun, []]}]
+    args = [get_owner(self()), get_callers(self()), {:erlang, :apply, [fun, []]}]
     start_child_with_spec(supervisor, args, restart, shutdown)
   end
 
@@ -411,7 +411,7 @@ defmodule Task.Supervisor do
       when is_atom(fun) and is_list(args) do
     restart = options[:restart]
     shutdown = options[:shutdown]
-    args = [get_owner(self()), {module, fun, args}]
+    args = [get_owner(self()), get_callers(self()), {module, fun, args}]
     start_child_with_spec(supervisor, args, restart, shutdown)
   end
 
@@ -432,9 +432,16 @@ defmodule Task.Supervisor do
     {node(), self_or_name, pid}
   end
 
+  defp get_callers(owner) do
+    case :erlang.get(:"$callers") do
+      [_ | _] = list -> [owner | list]
+      _ -> [owner]
+    end
+  end
+
   defp async(supervisor, link_type, module, fun, args, options) do
     owner = self()
-    args = [get_owner(owner), :monitor, {module, fun, args}]
+    args = [get_owner(owner), get_callers(owner), :monitor, {module, fun, args}]
     shutdown = options[:shutdown]
 
     case start_child_with_spec(supervisor, args, :temporary, shutdown) do
@@ -456,8 +463,8 @@ defmodule Task.Supervisor do
   defp build_stream(supervisor, link_type, enumerable, fun, options) do
     shutdown = options[:shutdown]
 
-    &Task.Supervised.stream(enumerable, &1, &2, fun, options, fn owner, mfa ->
-      args = [get_owner(owner), :monitor, mfa]
+    &Task.Supervised.stream(enumerable, &1, &2, fun, options, fn [owner | _] = callers, mfa ->
+      args = [get_owner(owner), callers, :monitor, mfa]
 
       case start_child_with_spec(supervisor, args, :temporary, shutdown) do
         {:ok, pid} ->

lib/elixir/test/elixir/task/supervisor_test.exs

@@ -105,13 +105,30 @@ defmodule Task.SupervisorTest do
 
     test "raises when :max_children is reached" do
       {:ok, sup} = Task.Supervisor.start_link(max_children: 1)
-
       Task.Supervisor.async(sup, fn -> Process.sleep(:infinity) end)
 
       assert_raise RuntimeError, ~r/reached the maximum number of tasks/, fn ->
         Task.Supervisor.async(sup, fn -> :ok end)
       end
     end
+
+    test "with $callers", config do
+      sup = config[:supervisor]
+      grandparent = self()
+
+      Task.Supervisor.async(sup, fn ->
+        parent = self()
+        assert Process.get(:"$callers") == [grandparent]
+        assert Process.get(:"$ancestors") == [sup, grandparent]
+
+        Task.Supervisor.async(sup, fn ->
+          assert Process.get(:"$callers") == [parent, grandparent]
+          assert Process.get(:"$ancestors") == [sup, grandparent]
+        end)
+        |> Task.await()
+      end)
+      |> Task.await()
+    end
   end
 
   test "async/3", config do
@@ -262,6 +279,25 @@ defmodule Task.SupervisorTest do
     assert_receive :ready
   end
 
+  test "start_child/1 with $callers", config do
+    sup = config[:supervisor]
+    grandparent = self()
+
+    Task.Supervisor.start_child(sup, fn ->
+      parent = self()
+      assert Process.get(:"$callers") == [grandparent]
+      assert Process.get(:"$ancestors") == [sup, grandparent]
+
+      Task.Supervisor.start_child(sup, fn ->
+        assert Process.get(:"$callers") == [parent, grandparent]
+        assert Process.get(:"$ancestors") == [sup, grandparent]
+        send(grandparent, :done)
+      end)
+    end)
+
+    assert_receive :done
+  end
+
   test "terminate_child/2", config do
     args = [self(), :done]
 
@@ -369,6 +405,27 @@ defmodule Task.SupervisorTest do
       assert MapSet.new(links) == MapSet.new([unused_supervisor, supervisor])
       refute_received _
     end
+
+    test "with $callers", config do
+      sup = config[:supervisor]
+      grandparent = self()
+
+      Task.Supervisor.async_stream(sup, [1], fn 1 ->
+        parent = self()
+        assert Process.get(:"$callers") == [grandparent]
+        assert Process.get(:"$ancestors") == [sup, grandparent]
+
+        Task.Supervisor.async_stream(sup, [1], fn 1 ->
+          assert Process.get(:"$callers") == [parent, grandparent]
+          assert Process.get(:"$ancestors") == [sup, grandparent]
+          send(grandparent, :done)
+        end)
+        |> Stream.run()
+      end)
+      |> Stream.run()
+
+      assert_receive :done
+    end
   end
 
   describe "async_stream_nolink" do