From 2dd64c68d9fb3ba8ead5cb0ab70b816016adca7c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jonatan=20K=C5=82osko?= <jonatanklosko@gmail.com>
Date: Thu, 26 Sep 2024 20:29:30 +0700
Subject: [PATCH 01/20] Wrap compilation in a cross-os-process lock

---
 lib/mix/lib/mix/lock.ex                 | 238 ++++++++++++++++++++++++
 lib/mix/lib/mix/state.ex                |  74 --------
 lib/mix/lib/mix/tasks/compile.all.ex    |   7 +
 lib/mix/lib/mix/tasks/compile.elixir.ex |   5 +-
 lib/mix/lib/mix/tasks/compile.erlang.ex |  13 +-
 lib/mix/lib/mix/tasks/compile.ex        |  13 ++
 lib/mix/lib/mix/tasks/deps.compile.ex   |  17 +-
 lib/mix/test/mix/lock_test.exs          | 161 ++++++++++++++++
 lib/mix/test/mix/state_test.exs         | 101 ----------
 9 files changed, 439 insertions(+), 190 deletions(-)
 create mode 100644 lib/mix/lib/mix/lock.ex
 create mode 100644 lib/mix/test/mix/lock_test.exs
 delete mode 100644 lib/mix/test/mix/state_test.exs

diff --git a/lib/mix/lib/mix/lock.ex b/lib/mix/lib/mix/lock.ex
new file mode 100644
index 00000000000..dbf7af52a3d
--- /dev/null
+++ b/lib/mix/lib/mix/lock.ex
@@ -0,0 +1,238 @@
+defmodule Mix.Lock do
+  @moduledoc false
+
+  # Lock implementation working across multiple OS processes.
+  #
+  # The lock is implemented using TCP sockets and hard links.
+  #
+  # A process holds the lock if it owns a TCP socket, whose port is
+  # written in the lock_0 file. We need to create such lock files
+  # atomically, so the process first writes its port to a port_P
+  # file and then attempts to create a hard link to it at lock_0.
+  #
+  # An inherent problem with lock files is that the lock owner may
+  # terminate abruptly, leaving a "stale" file. Other processes can
+  # detect a stale file by reading the port written in that file,
+  # trying to connect to thart port and failing. In order for another
+  # process to link to the same path, the file needs to be replaced.
+  # However, we need to guarantee that only a single process can
+  # remove or replace the file, otherwise a concurrent process may
+  # end up removing a newly linked file.
+  #
+  # To address this problem we employ a chained locking procedure.
+  # Specifically, we attempt to link our port to lock_0, if that
+  # fails, we try to connect to the lock_0 port. If we manage to
+  # connect, it means the lock is taken, so we wait for it to close
+  # and start over. If we fail to connect, it means the lock is stale,
+  # so we want to replace it. In order to do that, we try to obtain
+  # lock_1. Again, we try to link and connect. Eventually, we should
+  # successfully link to lock_N. At that point we can clean up all
+  # the files, so we perform these steps:
+  #
+  #   * move our port_P to lock_0
+  #   * remove all the other port_P files
+  #   * remove all lock_1+ files
+  #
+  # It is important to perform these steps in this order, to avoid
+  # race conditions. By moving to lock_0, we make sure that all new
+  # processes trying to lock will connect to our port. By removing
+  # all port_P files we make sure that currently paused processes
+  # that are about to link port_P at lock_N will fail to link, since
+  # the port_P file will no longer exist (once lock_N is removed).
+  #
+  # Finally, note that we do not remove the lock file in `unlock/1`.
+  # If we did that, another process could try to connect and fail
+  # because the file would not exist, in such case the process would
+  # assume the file is stale and needs to be replaced, therefore
+  # possibly replacing another process who successfully links at the
+  # empty spot. This means we effectively always leave a stale file,
+  # however, in order to shortcut the port check for future processes,
+  # we atomically replace the file content with port 0, to indicate
+  # the file is stale.
+  #
+  # The main caveat of using ephemeral TCP ports is that they are not
+  # unique. This creates a theoretical scenario where the lock holder
+  # terminates abruptly and leaves its port in lock_0, then the port
+  # is assigned to a unrelated process (unaware of the locking). To
+  # handle this scenario, when we connect to a lock_N port, we expect
+  # it to immediately send us `@probe_data`. If this does not happen
+  # within `@probe_timeout_ms`, we assume the port is taken by an
+  # unrelated process and the lock file is stale. Note that it is ok
+  # to use a long timeout, because this scenario is very unlikely.
+  # Theoretically, if an actual lock owner is not able to send the
+  # probe data within the timeout, the lock will fail, however with
+  # a high enough timeout, this should not be a problem in practice.
+
+  @loopback {127, 0, 0, 1}
+  @listen_opts [:binary, ip: @loopback, packet: :raw, nodelay: true, backlog: 128, active: false]
+  @connect_opts [:binary, packet: :raw, nodelay: true, active: false]
+  @probe_data <<"elixirlock">>
+  @probe_timeout_ms 5_000
+
+  @doc """
+  Acquires a lock identified by the given key.
+
+  This function blocks until the lock is acquired by this process.
+
+  This function can also be called if this process already has the
+  lock. In such case the function is executed immediately.
+  """
+  @spec lock(iodata(), (-> term())) :: :ok
+  def lock(key, fun) do
+    key = key |> :erlang.md5() |> Base.url_encode64(padding: false)
+    path = Path.join([System.tmp_dir!(), "mix_lock", key])
+
+    pdict_key = {__MODULE__, path}
+    acquire_lock? = Process.get(pdict_key) == nil
+
+    try do
+      if acquire_lock? do
+        lock = lock(path)
+        Process.put(pdict_key, lock)
+      end
+
+      fun.()
+    after
+      if acquire_lock? do
+        lock = Process.get(pdict_key)
+        unlock(lock)
+        Process.delete(pdict_key)
+      end
+    end
+  end
+
+  defp lock(path) do
+    :ok = File.mkdir_p(path)
+
+    {:ok, socket} = :gen_tcp.listen(0, @listen_opts)
+    {:ok, port} = :inet.port(socket)
+
+    spawn_link(fn -> accept_loop(socket) end)
+
+    try_lock(path, socket, port)
+  end
+
+  defp try_lock(path, socket, port) do
+    port_path = Path.join(path, "port_#{port}")
+
+    :ok = File.write(port_path, <<port::unsigned-integer-32>>, [:raw])
+
+    case grab_lock(path, port_path, 0) do
+      {:ok, 0} ->
+        # We grabbed lock_0, so all good
+        %{socket: socket, path: path}
+
+      {:ok, _n} ->
+        # We grabbed lock_1+, so we need to replace lock_0 and clean up
+        take_over(path, port_path)
+        %{socket: socket, path: path}
+
+      {:taken, probe_socket} ->
+        # Another process has the lock, wait for close and start over
+        await_close(probe_socket)
+        try_lock(path, socket, port)
+
+      :invalidated ->
+        try_lock(path, socket, port)
+    end
+  end
+
+  defp grab_lock(path, port_path, n) do
+    lock_path = Path.join(path, "lock_#{n}")
+
+    case File.ln(port_path, lock_path) do
+      :ok ->
+        {:ok, n}
+
+      {:error, :eexist} ->
+        case probe(lock_path) do
+          {:ok, probe_socket} ->
+            {:taken, probe_socket}
+
+          :error ->
+            grab_lock(path, port_path, n + 1)
+        end
+
+      {:error, :enoent} ->
+        :invalidated
+    end
+  end
+
+  defp accept_loop(listen_socket) do
+    case :gen_tcp.accept(listen_socket) do
+      {:ok, socket} ->
+        _ = :gen_tcp.send(socket, @probe_data)
+        accept_loop(listen_socket)
+
+      {:error, reason} when reason in [:closed, :einval] ->
+        :ok
+    end
+  end
+
+  defp probe(port_path) do
+    with {:ok, <<port::unsigned-integer-32>>} when port > 0 <- File.read(port_path),
+         {:ok, socket} <- connect(port) do
+      case :gen_tcp.recv(socket, 0, @probe_timeout_ms) do
+        {:ok, @probe_data} ->
+          {:ok, socket}
+
+        {:error, _reason} ->
+          :gen_tcp.close(socket)
+          :error
+      end
+    else
+      _other -> :error
+    end
+  end
+
+  defp connect(port) do
+    # On Windows connecting to an unbound port takes a few seconds to
+    # fail, so instead we shortcut the check by attempting a listen,
+    # which succeeds or fails immediately
+    case :gen_tcp.listen(port, [reuseaddr: true] ++ @listen_opts) do
+      {:ok, socket} ->
+        :gen_tcp.close(socket)
+        # The port is free, so connecting would fail
+        {:error, :econnrefused}
+
+      {:error, _reason} ->
+        :gen_tcp.connect(@loopback, port, @connect_opts)
+    end
+  end
+
+  defp take_over(path, port_path) do
+    lock_path = Path.join(path, "lock_0")
+
+    # We linked to lock_N successfully, so port_path should exist
+    :ok = File.rename(port_path, lock_path)
+
+    {:ok, names} = File.ls(path)
+
+    for "port_" <> _ = name <- names do
+      _ = File.rm(Path.join(path, name))
+    end
+
+    for "lock_" <> _ = name <- names, name != "lock_0" do
+      _ = File.rm(Path.join(path, name))
+    end
+  end
+
+  defp await_close(socket) do
+    {:error, _reason} = :gen_tcp.recv(socket, 0)
+  end
+
+  defp unlock(lock) do
+    port_path = Path.join(lock.path, "port_0")
+    lock_path = Path.join(lock.path, "lock_0")
+
+    with :ok <- File.write(port_path, <<0::unsigned-integer-32>>, [:raw]) do
+      _ = File.rename(port_path, lock_path)
+    end
+
+    # Closing the socket will cause the accepting process to finish
+    # and all accepted sockets (tied to that process) will get closed
+    :gen_tcp.close(lock.socket)
+
+    :ok
+  end
+end
diff --git a/lib/mix/lib/mix/state.ex b/lib/mix/lib/mix/state.ex
index 41e6744b08a..5c99f13f9b6 100644
--- a/lib/mix/lib/mix/state.ex
+++ b/lib/mix/lib/mix/state.ex
@@ -9,15 +9,6 @@ defmodule Mix.State do
     GenServer.start_link(__MODULE__, :ok, name: @name)
   end
 
-  def lock(key, fun) do
-    try do
-      GenServer.call(@name, {:lock, key}, @timeout)
-      fun.()
-    after
-      GenServer.call(@name, {:unlock, key}, @timeout)
-    end
-  end
-
   def builtin_apps do
     GenServer.call(@name, :builtin_apps, @timeout)
   end
@@ -83,8 +74,6 @@ defmodule Mix.State do
     )
 
     state = %{
-      key_to_waiting: %{},
-      pid_to_key: %{},
       builtin_apps: :code.get_path()
     }
 
@@ -114,69 +103,6 @@ defmodule Mix.State do
     end
   end
 
-  @impl true
-  def handle_call({:lock, key}, {pid, _} = from, state) do
-    %{key_to_waiting: key_to_waiting, pid_to_key: pid_to_key} = state
-
-    key_to_waiting =
-      case key_to_waiting do
-        %{^key => {locked, waiting}} ->
-          Map.put(key_to_waiting, key, {locked, :queue.in(from, waiting)})
-
-        %{} ->
-          go!(from)
-          Map.put(key_to_waiting, key, {pid, :queue.new()})
-      end
-
-    ref = Process.monitor(pid)
-    pid_to_key = Map.put(pid_to_key, pid, {key, ref})
-    {:noreply, %{state | key_to_waiting: key_to_waiting, pid_to_key: pid_to_key}}
-  end
-
-  @impl true
-  def handle_call({:unlock, key}, {pid, _}, state) do
-    %{key_to_waiting: key_to_waiting, pid_to_key: pid_to_key} = state
-    {{^key, ref}, pid_to_key} = Map.pop(pid_to_key, pid)
-    Process.demonitor(ref, [:flush])
-    key_to_waiting = unlock(key_to_waiting, pid_to_key, key)
-    {:reply, :ok, %{state | key_to_waiting: key_to_waiting, pid_to_key: pid_to_key}}
-  end
-
-  @impl true
-  def handle_info({:DOWN, ref, _type, pid, _reason}, state) do
-    %{key_to_waiting: key_to_waiting, pid_to_key: pid_to_key} = state
-    {{key, ^ref}, pid_to_key} = Map.pop(pid_to_key, pid)
-
-    key_to_waiting =
-      case key_to_waiting do
-        %{^key => {^pid, _}} ->
-          unlock(key_to_waiting, pid_to_key, key)
-
-        %{^key => {locked, waiting}} ->
-          waiting = :queue.delete_with(fn {qpid, _qref} -> qpid == pid end, waiting)
-          Map.put(key_to_waiting, key, {locked, waiting})
-      end
-
-    {:noreply, %{state | key_to_waiting: key_to_waiting, pid_to_key: pid_to_key}}
-  end
-
-  defp unlock(key_to_waiting, pid_to_key, key) do
-    %{^key => {_locked, waiting}} = key_to_waiting
-
-    case :queue.out(waiting) do
-      {{:value, {pid, _} = from}, waiting} ->
-        # Assert that we still know this PID
-        _ = Map.fetch!(pid_to_key, pid)
-        go!(from)
-        Map.put(key_to_waiting, key, {pid, waiting})
-
-      {:empty, _waiting} ->
-        Map.delete(key_to_waiting, key)
-    end
-  end
-
-  defp go!(from), do: GenServer.reply(from, :ok)
-
   # ../elixir/ebin -> elixir
   # ../ssl-9.6/ebin -> ssl
   defp app_from_code_path(path) do
diff --git a/lib/mix/lib/mix/tasks/compile.all.ex b/lib/mix/lib/mix/tasks/compile.all.ex
index 04930d769ab..47b29dffea9 100644
--- a/lib/mix/lib/mix/tasks/compile.all.ex
+++ b/lib/mix/lib/mix/tasks/compile.all.ex
@@ -12,6 +12,13 @@ defmodule Mix.Tasks.Compile.All do
   @impl true
   def run(args) do
     Mix.Project.get!()
+
+    Mix.Lock.lock(Mix.Tasks.Compile.compile_lock_key(), fn ->
+      do_run(args)
+    end)
+  end
+
+  defp do_run(args) do
     config = Mix.Project.config()
 
     # Compute the app cache if it is stale and we are
diff --git a/lib/mix/lib/mix/tasks/compile.elixir.ex b/lib/mix/lib/mix/tasks/compile.elixir.ex
index 2b2d67d7739..b10433a5e02 100644
--- a/lib/mix/lib/mix/tasks/compile.elixir.ex
+++ b/lib/mix/lib/mix/tasks/compile.elixir.ex
@@ -123,11 +123,10 @@ defmodule Mix.Tasks.Compile.Elixir do
       |> profile_opts()
 
     # Having compilations racing with other is most undesired,
-    # so we wrap the compiler in a lock. Ideally we would use
-    # flock in the future.
+    # so we wrap the compiler in a lock.
 
     with_logger_app(project, fn ->
-      Mix.State.lock(__MODULE__, fn ->
+      Mix.Lock.lock(Mix.Tasks.Compile.compile_lock_key(), fn ->
         Mix.Compilers.Elixir.compile(
           manifest,
           srcs,
diff --git a/lib/mix/lib/mix/tasks/compile.erlang.ex b/lib/mix/lib/mix/tasks/compile.erlang.ex
index 301331b99ac..3477fb894b2 100644
--- a/lib/mix/lib/mix/tasks/compile.erlang.ex
+++ b/lib/mix/lib/mix/tasks/compile.erlang.ex
@@ -46,11 +46,14 @@ defmodule Mix.Tasks.Compile.Erlang do
   @impl true
   def run(args) do
     {opts, _, _} = OptionParser.parse(args, switches: @switches)
-    project = Mix.Project.config()
-    source_paths = project[:erlc_paths]
-    Mix.Compilers.Erlang.assert_valid_erlc_paths(source_paths)
-    files = Mix.Utils.extract_files(source_paths, [:erl])
-    do_run(files, opts, project, source_paths)
+
+    Mix.Lock.lock(Mix.Tasks.Compile.compile_lock_key(), fn ->
+      project = Mix.Project.config()
+      source_paths = project[:erlc_paths]
+      Mix.Compilers.Erlang.assert_valid_erlc_paths(source_paths)
+      files = Mix.Utils.extract_files(source_paths, [:erl])
+      do_run(files, opts, project, source_paths)
+    end)
   end
 
   defp do_run([], _, _, _), do: {:noop, []}
diff --git a/lib/mix/lib/mix/tasks/compile.ex b/lib/mix/lib/mix/tasks/compile.ex
index a893fc19227..29145342f60 100644
--- a/lib/mix/lib/mix/tasks/compile.ex
+++ b/lib/mix/lib/mix/tasks/compile.ex
@@ -243,6 +243,19 @@ defmodule Mix.Tasks.Compile do
     end)
   end
 
+  @doc false
+  def compile_lock_key() do
+    # To avoid duplicated compilation, we wrap compilation tasks, such
+    # as compile.all, deps.compile, compile.elixir, compile.erlang in
+    # a lock. Note that compile.all covers compile.elixir, but the
+    # latter can still be invoked directly, hence we put the lock over
+    # the individual tasks.
+
+    config = Mix.Project.config()
+    build_path = Mix.Project.build_path(config)
+    ["compile", build_path]
+  end
+
   ## Consolidation handling
 
   defp reconsolidate_protocols?(:ok), do: true
diff --git a/lib/mix/lib/mix/tasks/deps.compile.ex b/lib/mix/lib/mix/tasks/deps.compile.ex
index 0865990c6e9..ce0f4dadb9a 100644
--- a/lib/mix/lib/mix/tasks/deps.compile.ex
+++ b/lib/mix/lib/mix/tasks/deps.compile.ex
@@ -53,15 +53,18 @@ defmodule Mix.Tasks.Deps.Compile do
     end
 
     Mix.Project.get!()
-    deps = Mix.Dep.load_and_cache()
 
-    case OptionParser.parse(args, switches: @switches) do
-      {opts, [], _} ->
-        compile(filter_available_and_local_deps(deps), opts)
+    Mix.Lock.lock(Mix.Tasks.Compile.compile_lock_key(), fn ->
+      deps = Mix.Dep.load_and_cache()
 
-      {opts, tail, _} ->
-        compile(Mix.Dep.filter_by_name(tail, deps, opts), opts)
-    end
+      case OptionParser.parse(args, switches: @switches) do
+        {opts, [], _} ->
+          compile(filter_available_and_local_deps(deps), opts)
+
+        {opts, tail, _} ->
+          compile(Mix.Dep.filter_by_name(tail, deps, opts), opts)
+      end
+    end)
   end
 
   @doc false
diff --git a/lib/mix/test/mix/lock_test.exs b/lib/mix/test/mix/lock_test.exs
new file mode 100644
index 00000000000..5e2d2754fcc
--- /dev/null
+++ b/lib/mix/test/mix/lock_test.exs
@@ -0,0 +1,161 @@
+Code.require_file("../test_helper.exs", __DIR__)
+
+defmodule Mix.LockTest do
+  use ExUnit.Case, async: true
+
+  @lock_key Atom.to_string(__MODULE__)
+
+  test "executes functions" do
+    assert Mix.Lock.lock(@lock_key, fn -> :it_works! end) == :it_works!
+    assert Mix.Lock.lock(@lock_key, fn -> :still_works! end) == :still_works!
+  end
+
+  test "releases lock on error" do
+    assert_raise RuntimeError, fn ->
+      Mix.Lock.lock(@lock_key, fn -> raise "oops" end)
+    end
+
+    assert Mix.Lock.lock(@lock_key, fn -> :still_works! end) == :still_works!
+  end
+
+  test "releases lock on exit" do
+    {_pid, ref} =
+      spawn_monitor(fn ->
+        Mix.Lock.lock(@lock_key, fn -> Process.exit(self(), :kill) end)
+      end)
+
+    assert_receive {:DOWN, ^ref, _, _, _}
+    assert Mix.Lock.lock(@lock_key, fn -> :still_works! end) == :still_works!
+  end
+
+  test "blocks until released" do
+    parent = self()
+
+    task =
+      Task.async(fn ->
+        Mix.Lock.lock(@lock_key, fn ->
+          send(parent, :locked)
+          assert_receive :will_lock
+          :it_works!
+        end)
+      end)
+
+    assert_receive :locked
+    send(task.pid, :will_lock)
+    assert Mix.Lock.lock(@lock_key, fn -> :still_works! end) == :still_works!
+    assert Task.await(task) == :it_works!
+  end
+
+  @tag :capture_log
+  test "blocks until released on error" do
+    parent = self()
+
+    {pid, ref} =
+      spawn_monitor(fn ->
+        Mix.Lock.lock(@lock_key, fn ->
+          send(parent, :locked)
+          assert_receive :will_lock
+          raise "oops"
+        end)
+      end)
+
+    assert_receive :locked
+    send(pid, :will_lock)
+    assert Mix.Lock.lock(@lock_key, fn -> :still_works! end) == :still_works!
+    assert_receive {:DOWN, ^ref, _, _, _}
+  end
+
+  test "blocks until released on exit" do
+    parent = self()
+
+    {pid, ref} =
+      spawn_monitor(fn ->
+        Mix.Lock.lock(@lock_key, fn ->
+          send(parent, :locked)
+          assert_receive :will_not_lock
+        end)
+      end)
+
+    assert_receive :locked
+    Process.exit(pid, :kill)
+    assert Mix.Lock.lock(@lock_key, fn -> :still_works! end) == :still_works!
+    assert_receive {:DOWN, ^ref, _, _, _}
+  end
+
+  test "scheduls and releases on exit" do
+    assert Mix.Lock.lock(@lock_key, fn ->
+             {pid, ref} =
+               spawn_monitor(fn ->
+                 Mix.Lock.lock(@lock_key, fn ->
+                   raise "this will never be invoked"
+                 end)
+               end)
+
+             Process.exit(pid, :kill)
+             assert_receive {:DOWN, ^ref, _, _, :killed}
+             :it_works!
+           end) == :it_works!
+
+    assert Mix.Lock.lock(@lock_key, fn -> :still_works! end) == :still_works!
+  end
+
+  @tag :tmp_dir
+  test "property test with file access", %{tmp_dir: tmp_dir} do
+    # Spawn N concurrent processes incrementing number in a file
+    n = 10
+    number_path = Path.join(tmp_dir, "number.txt")
+
+    File.write!(number_path, "0")
+
+    refs =
+      for _ <- 1..n do
+        spawn_monitor(fn ->
+          Mix.Lock.lock(@lock_key, fn ->
+            number = number_path |> File.read!() |> String.to_integer()
+            new_number = number + 1
+            File.write!(number_path, Integer.to_string(new_number))
+
+            assert File.read!(number_path) == Integer.to_string(new_number)
+
+            # Terminate without unlocking in random cases
+            case Enum.random(1..2) do
+              1 -> Process.exit(self(), :kill)
+              2 -> :ok
+            end
+          end)
+        end)
+        |> elem(1)
+      end
+
+    await_monitors(refs)
+
+    assert File.read!(number_path) == Integer.to_string(n)
+  end
+
+  test "lock can be acquired multiple times by the same process" do
+    {_pid, ref} =
+      spawn_monitor(fn ->
+        Mix.Lock.lock(@lock_key, fn ->
+          Mix.Lock.lock(@lock_key, fn ->
+            Process.exit(self(), :kill)
+          end)
+        end)
+      end)
+
+    assert_receive {:DOWN, ^ref, _, _, _}
+
+    assert Mix.Lock.lock(@lock_key, fn ->
+             Mix.Lock.lock(@lock_key, fn ->
+               :still_works!
+             end)
+           end) == :still_works!
+  end
+
+  defp await_monitors([]), do: :ok
+
+  defp await_monitors(refs) do
+    receive do
+      {:DOWN, ref, _, _, _} -> await_monitors(refs -- [ref])
+    end
+  end
+end
diff --git a/lib/mix/test/mix/state_test.exs b/lib/mix/test/mix/state_test.exs
deleted file mode 100644
index aaf131bec82..00000000000
--- a/lib/mix/test/mix/state_test.exs
+++ /dev/null
@@ -1,101 +0,0 @@
-Code.require_file("../test_helper.exs", __DIR__)
-
-defmodule Mix.StateTest do
-  use ExUnit.Case, async: true
-
-  describe "lock" do
-    test "executes functions" do
-      assert Mix.State.lock(:key, fn -> :it_works! end) == :it_works!
-      assert Mix.State.lock(:key, fn -> :still_works! end) == :still_works!
-    end
-
-    test "releases lock on error" do
-      assert_raise RuntimeError, fn ->
-        Mix.State.lock(:key, fn -> raise "oops" end)
-      end
-
-      assert Mix.State.lock(:key, fn -> :still_works! end) == :still_works!
-    end
-
-    test "releases lock on exit" do
-      {_pid, ref} =
-        spawn_monitor(fn ->
-          Mix.State.lock(:key, fn -> Process.exit(self(), :kill) end)
-        end)
-
-      assert_receive {:DOWN, ^ref, _, _, _}
-      assert Mix.State.lock(:key, fn -> :still_works! end) == :still_works!
-    end
-
-    test "blocks until released" do
-      parent = self()
-
-      task =
-        Task.async(fn ->
-          Mix.State.lock(:key, fn ->
-            send(parent, :locked)
-            assert_receive :will_lock
-            :it_works!
-          end)
-        end)
-
-      assert_receive :locked
-      send(task.pid, :will_lock)
-      assert Mix.State.lock(:key, fn -> :still_works! end) == :still_works!
-      assert Task.await(task) == :it_works!
-    end
-
-    @tag :capture_log
-    test "blocks until released on error" do
-      parent = self()
-
-      {pid, ref} =
-        spawn_monitor(fn ->
-          Mix.State.lock(:key, fn ->
-            send(parent, :locked)
-            assert_receive :will_lock
-            raise "oops"
-          end)
-        end)
-
-      assert_receive :locked
-      send(pid, :will_lock)
-      assert Mix.State.lock(:key, fn -> :still_works! end) == :still_works!
-      assert_receive {:DOWN, ^ref, _, _, _}
-    end
-
-    test "blocks until released on exit" do
-      parent = self()
-
-      {pid, ref} =
-        spawn_monitor(fn ->
-          Mix.State.lock(:key, fn ->
-            send(parent, :locked)
-            assert_receive :will_not_lock
-          end)
-        end)
-
-      assert_receive :locked
-      Process.exit(pid, :kill)
-      assert Mix.State.lock(:key, fn -> :still_works! end) == :still_works!
-      assert_receive {:DOWN, ^ref, _, _, _}
-    end
-
-    test "scheduls and releases on exit" do
-      assert Mix.State.lock(:key, fn ->
-               {pid, ref} =
-                 spawn_monitor(fn ->
-                   Mix.State.lock(:key, fn ->
-                     raise "this will never be invoked"
-                   end)
-                 end)
-
-               Process.exit(pid, :kill)
-               assert_receive {:DOWN, ^ref, _, _, :killed}
-               :it_works!
-             end) == :it_works!
-
-      assert Mix.State.lock(:key, fn -> :still_works! end) == :still_works!
-    end
-  end
-end

From 745b6101ef29278e4d649dc6d725e87eebee2e79 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jonatan=20K=C5=82osko?= <jonatanklosko@gmail.com>
Date: Fri, 27 Sep 2024 09:16:56 +0200
Subject: [PATCH 02/20] Apply suggestions from code review

Co-authored-by: Andrea Leopardi <an.leopardi@gmail.com>
---
 lib/mix/lib/mix/lock.ex | 20 ++++++++++----------
 1 file changed, 10 insertions(+), 10 deletions(-)

diff --git a/lib/mix/lib/mix/lock.ex b/lib/mix/lib/mix/lock.ex
index dbf7af52a3d..8370d8a24f6 100644
--- a/lib/mix/lib/mix/lock.ex
+++ b/lib/mix/lib/mix/lock.ex
@@ -13,7 +13,7 @@ defmodule Mix.Lock do
   # An inherent problem with lock files is that the lock owner may
   # terminate abruptly, leaving a "stale" file. Other processes can
   # detect a stale file by reading the port written in that file,
-  # trying to connect to thart port and failing. In order for another
+  # trying to connect to that port and failing. In order for another
   # process to link to the same path, the file needs to be replaced.
   # However, we need to guarantee that only a single process can
   # remove or replace the file, otherwise a concurrent process may
@@ -66,13 +66,14 @@ defmodule Mix.Lock do
   @loopback {127, 0, 0, 1}
   @listen_opts [:binary, ip: @loopback, packet: :raw, nodelay: true, backlog: 128, active: false]
   @connect_opts [:binary, packet: :raw, nodelay: true, active: false]
-  @probe_data <<"elixirlock">>
+  @probe_data "elixirlock"
   @probe_timeout_ms 5_000
 
   @doc """
   Acquires a lock identified by the given key.
 
-  This function blocks until the lock is acquired by this process.
+  This function blocks until the lock is acquired by this process,
+  and then executes `fun`, returning its return value.
 
   This function can also be called if this process already has the
   lock. In such case the function is executed immediately.
@@ -85,19 +86,18 @@ defmodule Mix.Lock do
     pdict_key = {__MODULE__, path}
     acquire_lock? = Process.get(pdict_key) == nil
 
-    try do
-      if acquire_lock? do
+    if acquire_lock? do
+      try do
         lock = lock(path)
         Process.put(pdict_key, lock)
-      end
-
-      fun.()
-    after
-      if acquire_lock? do
+        fun.()
+      after
         lock = Process.get(pdict_key)
         unlock(lock)
         Process.delete(pdict_key)
       end
+    else
+      fun.()
     end
   end
 

From 308ceaefee7674c2646c9e3b6869f03540d31500 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jonatan=20K=C5=82osko?= <jonatanklosko@gmail.com>
Date: Fri, 27 Sep 2024 14:27:16 +0700
Subject: [PATCH 03/20] Use assertive file operations

---
 lib/mix/lib/mix/lock.ex | 17 ++++++++---------
 1 file changed, 8 insertions(+), 9 deletions(-)

diff --git a/lib/mix/lib/mix/lock.ex b/lib/mix/lib/mix/lock.ex
index 8370d8a24f6..3a0ddac1d99 100644
--- a/lib/mix/lib/mix/lock.ex
+++ b/lib/mix/lib/mix/lock.ex
@@ -102,7 +102,7 @@ defmodule Mix.Lock do
   end
 
   defp lock(path) do
-    :ok = File.mkdir_p(path)
+    File.mkdir_p!(path)
 
     {:ok, socket} = :gen_tcp.listen(0, @listen_opts)
     {:ok, port} = :inet.port(socket)
@@ -115,7 +115,7 @@ defmodule Mix.Lock do
   defp try_lock(path, socket, port) do
     port_path = Path.join(path, "port_#{port}")
 
-    :ok = File.write(port_path, <<port::unsigned-integer-32>>, [:raw])
+    File.write!(port_path, <<port::unsigned-integer-32>>, [:raw])
 
     case grab_lock(path, port_path, 0) do
       {:ok, 0} ->
@@ -204,16 +204,16 @@ defmodule Mix.Lock do
     lock_path = Path.join(path, "lock_0")
 
     # We linked to lock_N successfully, so port_path should exist
-    :ok = File.rename(port_path, lock_path)
+    File.rename!(port_path, lock_path)
 
-    {:ok, names} = File.ls(path)
+    names = File.ls!(path)
 
     for "port_" <> _ = name <- names do
-      _ = File.rm(Path.join(path, name))
+      File.rm!(Path.join(path, name))
     end
 
     for "lock_" <> _ = name <- names, name != "lock_0" do
-      _ = File.rm(Path.join(path, name))
+      File.rm!(Path.join(path, name))
     end
   end
 
@@ -225,9 +225,8 @@ defmodule Mix.Lock do
     port_path = Path.join(lock.path, "port_0")
     lock_path = Path.join(lock.path, "lock_0")
 
-    with :ok <- File.write(port_path, <<0::unsigned-integer-32>>, [:raw]) do
-      _ = File.rename(port_path, lock_path)
-    end
+    File.write!(port_path, <<0::unsigned-integer-32>>, [:raw])
+    File.rename!(port_path, lock_path)
 
     # Closing the socket will cause the accepting process to finish
     # and all accepted sockets (tied to that process) will get closed

From 48b6da6f83959f5bb0394f029f8c7108851ffa40 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jonatan=20K=C5=82osko?= <jonatanklosko@gmail.com>
Date: Fri, 27 Sep 2024 15:00:22 +0700
Subject: [PATCH 04/20] Improve error if opening the socket fails

---
 lib/mix/lib/mix/lock.ex | 20 ++++++++++++--------
 1 file changed, 12 insertions(+), 8 deletions(-)

diff --git a/lib/mix/lib/mix/lock.ex b/lib/mix/lib/mix/lock.ex
index 3a0ddac1d99..c5fee7467e0 100644
--- a/lib/mix/lib/mix/lock.ex
+++ b/lib/mix/lib/mix/lock.ex
@@ -87,9 +87,10 @@ defmodule Mix.Lock do
     acquire_lock? = Process.get(pdict_key) == nil
 
     if acquire_lock? do
+      lock = lock(path)
+      Process.put(pdict_key, lock)
+
       try do
-        lock = lock(path)
-        Process.put(pdict_key, lock)
         fun.()
       after
         lock = Process.get(pdict_key)
@@ -104,12 +105,15 @@ defmodule Mix.Lock do
   defp lock(path) do
     File.mkdir_p!(path)
 
-    {:ok, socket} = :gen_tcp.listen(0, @listen_opts)
-    {:ok, port} = :inet.port(socket)
-
-    spawn_link(fn -> accept_loop(socket) end)
-
-    try_lock(path, socket, port)
+    with {:ok, socket} <- :gen_tcp.listen(0, @listen_opts),
+         {:ok, port} <- :inet.port(socket) do
+      spawn_link(fn -> accept_loop(socket) end)
+      try_lock(path, socket, port)
+    else
+      {:error, reason} ->
+        raise Mix.Error,
+              "failed to open a TCP socket while acquiring a lock, reason: #{inspect(reason)}"
+    end
   end
 
   defp try_lock(path, socket, port) do

From 7ca97ce92b375451887df6449009b715219f935c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jonatan=20K=C5=82osko?= <jonatanklosko@gmail.com>
Date: Fri, 27 Sep 2024 15:13:27 +0700
Subject: [PATCH 05/20] Improve error paths

---
 lib/mix/lib/mix/lock.ex | 14 ++++++++++----
 1 file changed, 10 insertions(+), 4 deletions(-)

diff --git a/lib/mix/lib/mix/lock.ex b/lib/mix/lib/mix/lock.ex
index c5fee7467e0..8c84e3f3988 100644
--- a/lib/mix/lib/mix/lock.ex
+++ b/lib/mix/lib/mix/lock.ex
@@ -108,7 +108,15 @@ defmodule Mix.Lock do
     with {:ok, socket} <- :gen_tcp.listen(0, @listen_opts),
          {:ok, port} <- :inet.port(socket) do
       spawn_link(fn -> accept_loop(socket) end)
-      try_lock(path, socket, port)
+
+      try do
+        try_lock(path, socket, port)
+      rescue
+        exception ->
+          # Close the socket to make sure we don't block the lock
+          :gen_tcp.close(socket)
+          reraise exception, __STACKTRACE__
+      end
     else
       {:error, reason} ->
         raise Mix.Error,
@@ -231,11 +239,9 @@ defmodule Mix.Lock do
 
     File.write!(port_path, <<0::unsigned-integer-32>>, [:raw])
     File.rename!(port_path, lock_path)
-
+  after
     # Closing the socket will cause the accepting process to finish
     # and all accepted sockets (tied to that process) will get closed
     :gen_tcp.close(lock.socket)
-
-    :ok
   end
 end

From 76a455e98c63336fc0be111b1165b3146feaf81b Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jonatan=20K=C5=82osko?= <jonatanklosko@gmail.com>
Date: Fri, 27 Sep 2024 15:16:03 +0700
Subject: [PATCH 06/20] Simplify pdict

---
 lib/mix/lib/mix/lock.ex | 11 +++++------
 1 file changed, 5 insertions(+), 6 deletions(-)

diff --git a/lib/mix/lib/mix/lock.ex b/lib/mix/lib/mix/lock.ex
index 8c84e3f3988..7c55313391d 100644
--- a/lib/mix/lib/mix/lock.ex
+++ b/lib/mix/lib/mix/lock.ex
@@ -84,21 +84,20 @@ defmodule Mix.Lock do
     path = Path.join([System.tmp_dir!(), "mix_lock", key])
 
     pdict_key = {__MODULE__, path}
-    acquire_lock? = Process.get(pdict_key) == nil
+    has_lock? = Process.get(pdict_key)
 
-    if acquire_lock? do
+    if has_lock? do
+      fun.()
+    else
       lock = lock(path)
-      Process.put(pdict_key, lock)
+      Process.put(pdict_key, true)
 
       try do
         fun.()
       after
-        lock = Process.get(pdict_key)
         unlock(lock)
         Process.delete(pdict_key)
       end
-    else
-      fun.()
     end
   end
 

From 90dd9d52a75e8387ee769a39ea758c5ea2c84cb7 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jonatan=20K=C5=82osko?= <jonatanklosko@gmail.com>
Date: Fri, 27 Sep 2024 15:17:52 +0700
Subject: [PATCH 07/20] Up

---
 lib/mix/lib/mix/lock.ex | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/lib/mix/lib/mix/lock.ex b/lib/mix/lib/mix/lock.ex
index 7c55313391d..a0a81acf720 100644
--- a/lib/mix/lib/mix/lock.ex
+++ b/lib/mix/lib/mix/lock.ex
@@ -95,8 +95,8 @@ defmodule Mix.Lock do
       try do
         fun.()
       after
-        unlock(lock)
         Process.delete(pdict_key)
+        unlock(lock)
       end
     end
   end

From d870aa7476eb9de120332bf9a09abf38c2c79c8b Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jonatan=20K=C5=82osko?= <jonatanklosko@gmail.com>
Date: Fri, 27 Sep 2024 16:48:55 +0700
Subject: [PATCH 08/20] Up

---
 lib/mix/lib/mix/lock.ex                 | 65 +++++++++++++++++++------
 lib/mix/lib/mix/tasks/compile.all.ex    |  2 +-
 lib/mix/lib/mix/tasks/compile.elixir.ex |  2 +-
 lib/mix/lib/mix/tasks/compile.erlang.ex |  2 +-
 lib/mix/lib/mix/tasks/deps.compile.ex   |  2 +-
 lib/mix/test/mix/lock_test.exs          | 40 +++++++--------
 6 files changed, 73 insertions(+), 40 deletions(-)

diff --git a/lib/mix/lib/mix/lock.ex b/lib/mix/lib/mix/lock.ex
index a0a81acf720..e069a2136f9 100644
--- a/lib/mix/lib/mix/lock.ex
+++ b/lib/mix/lib/mix/lock.ex
@@ -78,8 +78,8 @@ defmodule Mix.Lock do
   This function can also be called if this process already has the
   lock. In such case the function is executed immediately.
   """
-  @spec lock(iodata(), (-> term())) :: :ok
-  def lock(key, fun) do
+  @spec with_lock(iodata(), (-> term())) :: :ok
+  def with_lock(key, fun) do
     key = key |> :erlang.md5() |> Base.url_encode64(padding: false)
     path = Path.join([System.tmp_dir!(), "mix_lock", key])
 
@@ -95,6 +95,8 @@ defmodule Mix.Lock do
       try do
         fun.()
       after
+        # Unlocking will always close the socket, but it may raise,
+        # so we remove key from the dictionary first
         Process.delete(pdict_key)
         unlock(lock)
       end
@@ -104,25 +106,38 @@ defmodule Mix.Lock do
   defp lock(path) do
     File.mkdir_p!(path)
 
-    with {:ok, socket} <- :gen_tcp.listen(0, @listen_opts),
-         {:ok, port} <- :inet.port(socket) do
-      spawn_link(fn -> accept_loop(socket) end)
+    case listen() do
+      {:ok, socket, port} ->
+        spawn_link(fn -> accept_loop(socket) end)
+
+        try do
+          try_lock(path, socket, port)
+        rescue
+          exception ->
+            # Close the socket to make sure we don't block the lock
+            :gen_tcp.close(socket)
+            reraise exception, __STACKTRACE__
+        end
 
-      try do
-        try_lock(path, socket, port)
-      rescue
-        exception ->
-          # Close the socket to make sure we don't block the lock
-          :gen_tcp.close(socket)
-          reraise exception, __STACKTRACE__
-      end
-    else
       {:error, reason} ->
         raise Mix.Error,
               "failed to open a TCP socket while acquiring a lock, reason: #{inspect(reason)}"
     end
   end
 
+  defp listen() do
+    with {:ok, socket} <- :gen_tcp.listen(0, @listen_opts) do
+      case :inet.port(socket) do
+        {:ok, port} ->
+          {:ok, socket, port}
+
+        {:error, reason} ->
+          :socket.close(socket)
+          {:error, reason}
+      end
+    end
+  end
+
   defp try_lock(path, socket, port) do
     port_path = Path.join(path, "port_#{port}")
 
@@ -134,7 +149,10 @@ defmodule Mix.Lock do
         %{socket: socket, path: path}
 
       {:ok, _n} ->
-        # We grabbed lock_1+, so we need to replace lock_0 and clean up
+        # We grabbed lock_1+, so we need to replace lock_0 and clean
+        # up. This must happen in a precise order, so if anything
+        # fails, we keep the files as is and the next process that
+        # grabs the lock will do the cleanup
         take_over(path, port_path)
         %{socket: socket, path: path}
 
@@ -166,6 +184,13 @@ defmodule Mix.Lock do
 
       {:error, :enoent} ->
         :invalidated
+
+      {:error, reason} ->
+        raise File.LinkError,
+          reason: reason,
+          action: "create hard link",
+          existing: port_path,
+          new: lock_path
     end
   end
 
@@ -229,7 +254,15 @@ defmodule Mix.Lock do
   end
 
   defp await_close(socket) do
-    {:error, _reason} = :gen_tcp.recv(socket, 0)
+    case :gen_tcp.recv(socket, 0) do
+      {:error, :closed} ->
+        :ok
+
+      {:error, _other} ->
+        # In case of an unexpected error, we close the socket ourselves
+        # to retry
+        :gen_tcp.close(socket)
+    end
   end
 
   defp unlock(lock) do
diff --git a/lib/mix/lib/mix/tasks/compile.all.ex b/lib/mix/lib/mix/tasks/compile.all.ex
index 47b29dffea9..85c2f2813fe 100644
--- a/lib/mix/lib/mix/tasks/compile.all.ex
+++ b/lib/mix/lib/mix/tasks/compile.all.ex
@@ -13,7 +13,7 @@ defmodule Mix.Tasks.Compile.All do
   def run(args) do
     Mix.Project.get!()
 
-    Mix.Lock.lock(Mix.Tasks.Compile.compile_lock_key(), fn ->
+    Mix.Lock.with_lock(Mix.Tasks.Compile.compile_lock_key(), fn ->
       do_run(args)
     end)
   end
diff --git a/lib/mix/lib/mix/tasks/compile.elixir.ex b/lib/mix/lib/mix/tasks/compile.elixir.ex
index b10433a5e02..54dfd792ba4 100644
--- a/lib/mix/lib/mix/tasks/compile.elixir.ex
+++ b/lib/mix/lib/mix/tasks/compile.elixir.ex
@@ -126,7 +126,7 @@ defmodule Mix.Tasks.Compile.Elixir do
     # so we wrap the compiler in a lock.
 
     with_logger_app(project, fn ->
-      Mix.Lock.lock(Mix.Tasks.Compile.compile_lock_key(), fn ->
+      Mix.Lock.with_lock(Mix.Tasks.Compile.compile_lock_key(), fn ->
         Mix.Compilers.Elixir.compile(
           manifest,
           srcs,
diff --git a/lib/mix/lib/mix/tasks/compile.erlang.ex b/lib/mix/lib/mix/tasks/compile.erlang.ex
index 3477fb894b2..7364e9b692f 100644
--- a/lib/mix/lib/mix/tasks/compile.erlang.ex
+++ b/lib/mix/lib/mix/tasks/compile.erlang.ex
@@ -47,7 +47,7 @@ defmodule Mix.Tasks.Compile.Erlang do
   def run(args) do
     {opts, _, _} = OptionParser.parse(args, switches: @switches)
 
-    Mix.Lock.lock(Mix.Tasks.Compile.compile_lock_key(), fn ->
+    Mix.Lock.with_lock(Mix.Tasks.Compile.compile_lock_key(), fn ->
       project = Mix.Project.config()
       source_paths = project[:erlc_paths]
       Mix.Compilers.Erlang.assert_valid_erlc_paths(source_paths)
diff --git a/lib/mix/lib/mix/tasks/deps.compile.ex b/lib/mix/lib/mix/tasks/deps.compile.ex
index ce0f4dadb9a..b579b33e732 100644
--- a/lib/mix/lib/mix/tasks/deps.compile.ex
+++ b/lib/mix/lib/mix/tasks/deps.compile.ex
@@ -54,7 +54,7 @@ defmodule Mix.Tasks.Deps.Compile do
 
     Mix.Project.get!()
 
-    Mix.Lock.lock(Mix.Tasks.Compile.compile_lock_key(), fn ->
+    Mix.Lock.with_lock(Mix.Tasks.Compile.compile_lock_key(), fn ->
       deps = Mix.Dep.load_and_cache()
 
       case OptionParser.parse(args, switches: @switches) do
diff --git a/lib/mix/test/mix/lock_test.exs b/lib/mix/test/mix/lock_test.exs
index 5e2d2754fcc..0b02d60813b 100644
--- a/lib/mix/test/mix/lock_test.exs
+++ b/lib/mix/test/mix/lock_test.exs
@@ -6,26 +6,26 @@ defmodule Mix.LockTest do
   @lock_key Atom.to_string(__MODULE__)
 
   test "executes functions" do
-    assert Mix.Lock.lock(@lock_key, fn -> :it_works! end) == :it_works!
-    assert Mix.Lock.lock(@lock_key, fn -> :still_works! end) == :still_works!
+    assert Mix.Lock.with_lock(@lock_key, fn -> :it_works! end) == :it_works!
+    assert Mix.Lock.with_lock(@lock_key, fn -> :still_works! end) == :still_works!
   end
 
   test "releases lock on error" do
     assert_raise RuntimeError, fn ->
-      Mix.Lock.lock(@lock_key, fn -> raise "oops" end)
+      Mix.Lock.with_lock(@lock_key, fn -> raise "oops" end)
     end
 
-    assert Mix.Lock.lock(@lock_key, fn -> :still_works! end) == :still_works!
+    assert Mix.Lock.with_lock(@lock_key, fn -> :still_works! end) == :still_works!
   end
 
   test "releases lock on exit" do
     {_pid, ref} =
       spawn_monitor(fn ->
-        Mix.Lock.lock(@lock_key, fn -> Process.exit(self(), :kill) end)
+        Mix.Lock.with_lock(@lock_key, fn -> Process.exit(self(), :kill) end)
       end)
 
     assert_receive {:DOWN, ^ref, _, _, _}
-    assert Mix.Lock.lock(@lock_key, fn -> :still_works! end) == :still_works!
+    assert Mix.Lock.with_lock(@lock_key, fn -> :still_works! end) == :still_works!
   end
 
   test "blocks until released" do
@@ -33,7 +33,7 @@ defmodule Mix.LockTest do
 
     task =
       Task.async(fn ->
-        Mix.Lock.lock(@lock_key, fn ->
+        Mix.Lock.with_lock(@lock_key, fn ->
           send(parent, :locked)
           assert_receive :will_lock
           :it_works!
@@ -42,7 +42,7 @@ defmodule Mix.LockTest do
 
     assert_receive :locked
     send(task.pid, :will_lock)
-    assert Mix.Lock.lock(@lock_key, fn -> :still_works! end) == :still_works!
+    assert Mix.Lock.with_lock(@lock_key, fn -> :still_works! end) == :still_works!
     assert Task.await(task) == :it_works!
   end
 
@@ -52,7 +52,7 @@ defmodule Mix.LockTest do
 
     {pid, ref} =
       spawn_monitor(fn ->
-        Mix.Lock.lock(@lock_key, fn ->
+        Mix.Lock.with_lock(@lock_key, fn ->
           send(parent, :locked)
           assert_receive :will_lock
           raise "oops"
@@ -61,7 +61,7 @@ defmodule Mix.LockTest do
 
     assert_receive :locked
     send(pid, :will_lock)
-    assert Mix.Lock.lock(@lock_key, fn -> :still_works! end) == :still_works!
+    assert Mix.Lock.with_lock(@lock_key, fn -> :still_works! end) == :still_works!
     assert_receive {:DOWN, ^ref, _, _, _}
   end
 
@@ -70,7 +70,7 @@ defmodule Mix.LockTest do
 
     {pid, ref} =
       spawn_monitor(fn ->
-        Mix.Lock.lock(@lock_key, fn ->
+        Mix.Lock.with_lock(@lock_key, fn ->
           send(parent, :locked)
           assert_receive :will_not_lock
         end)
@@ -78,15 +78,15 @@ defmodule Mix.LockTest do
 
     assert_receive :locked
     Process.exit(pid, :kill)
-    assert Mix.Lock.lock(@lock_key, fn -> :still_works! end) == :still_works!
+    assert Mix.Lock.with_lock(@lock_key, fn -> :still_works! end) == :still_works!
     assert_receive {:DOWN, ^ref, _, _, _}
   end
 
   test "scheduls and releases on exit" do
-    assert Mix.Lock.lock(@lock_key, fn ->
+    assert Mix.Lock.with_lock(@lock_key, fn ->
              {pid, ref} =
                spawn_monitor(fn ->
-                 Mix.Lock.lock(@lock_key, fn ->
+                 Mix.Lock.with_lock(@lock_key, fn ->
                    raise "this will never be invoked"
                  end)
                end)
@@ -96,7 +96,7 @@ defmodule Mix.LockTest do
              :it_works!
            end) == :it_works!
 
-    assert Mix.Lock.lock(@lock_key, fn -> :still_works! end) == :still_works!
+    assert Mix.Lock.with_lock(@lock_key, fn -> :still_works! end) == :still_works!
   end
 
   @tag :tmp_dir
@@ -110,7 +110,7 @@ defmodule Mix.LockTest do
     refs =
       for _ <- 1..n do
         spawn_monitor(fn ->
-          Mix.Lock.lock(@lock_key, fn ->
+          Mix.Lock.with_lock(@lock_key, fn ->
             number = number_path |> File.read!() |> String.to_integer()
             new_number = number + 1
             File.write!(number_path, Integer.to_string(new_number))
@@ -135,8 +135,8 @@ defmodule Mix.LockTest do
   test "lock can be acquired multiple times by the same process" do
     {_pid, ref} =
       spawn_monitor(fn ->
-        Mix.Lock.lock(@lock_key, fn ->
-          Mix.Lock.lock(@lock_key, fn ->
+        Mix.Lock.with_lock(@lock_key, fn ->
+          Mix.Lock.with_lock(@lock_key, fn ->
             Process.exit(self(), :kill)
           end)
         end)
@@ -144,8 +144,8 @@ defmodule Mix.LockTest do
 
     assert_receive {:DOWN, ^ref, _, _, _}
 
-    assert Mix.Lock.lock(@lock_key, fn ->
-             Mix.Lock.lock(@lock_key, fn ->
+    assert Mix.Lock.with_lock(@lock_key, fn ->
+             Mix.Lock.with_lock(@lock_key, fn ->
                :still_works!
              end)
            end) == :still_works!

From 4f90a175a822597f5e7ab4012336cc009cc63864 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jonatan=20K=C5=82osko?= <jonatanklosko@gmail.com>
Date: Fri, 27 Sep 2024 18:03:20 +0700
Subject: [PATCH 09/20] Up

---
 lib/mix/lib/mix/lock.ex | 36 +++++++++++++++++++++++++-----------
 1 file changed, 25 insertions(+), 11 deletions(-)

diff --git a/lib/mix/lib/mix/lock.ex b/lib/mix/lib/mix/lock.ex
index e069a2136f9..83cd0d3325a 100644
--- a/lib/mix/lib/mix/lock.ex
+++ b/lib/mix/lib/mix/lock.ex
@@ -178,7 +178,7 @@ defmodule Mix.Lock do
           {:ok, probe_socket} ->
             {:taken, probe_socket}
 
-          :error ->
+          {:error, _reason} ->
             grab_lock(path, port_path, n + 1)
         end
 
@@ -206,18 +206,17 @@ defmodule Mix.Lock do
   end
 
   defp probe(port_path) do
-    with {:ok, <<port::unsigned-integer-32>>} when port > 0 <- File.read(port_path),
+    with {:ok, port} <- fetch_probe_port(port_path),
          {:ok, socket} <- connect(port) do
-      case :gen_tcp.recv(socket, 0, @probe_timeout_ms) do
-        {:ok, @probe_data} ->
-          {:ok, socket}
+      await_probe_data(socket)
+    end
+  end
 
-        {:error, _reason} ->
-          :gen_tcp.close(socket)
-          :error
-      end
-    else
-      _other -> :error
+  defp fetch_probe_port(port_path) do
+    case File.read(port_path) do
+      {:ok, <<0::unsigned-integer-32>>} -> {:error, :ignore}
+      {:ok, <<port::unsigned-integer-32>>} -> {:ok, port}
+      {:error, reason} -> {:error, reason}
     end
   end
 
@@ -236,6 +235,21 @@ defmodule Mix.Lock do
     end
   end
 
+  defp await_probe_data(socket) do
+    case :gen_tcp.recv(socket, 0, @probe_timeout_ms) do
+      {:ok, @probe_data} ->
+        {:ok, socket}
+
+      {:ok, _data} ->
+        :gen_tcp.close(socket)
+        {:error, :unexpected_port_owner}
+
+      {:error, reason} ->
+        :gen_tcp.close(socket)
+        {:error, reason}
+    end
+  end
+
   defp take_over(path, port_path) do
     lock_path = Path.join(path, "lock_0")
 

From 854e1b5d42b9681b4f9e69a35615279683c0782c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jonatan=20K=C5=82osko?= <jonatanklosko@gmail.com>
Date: Fri, 27 Sep 2024 18:23:28 +0700
Subject: [PATCH 10/20] Handle interrupts

---
 lib/mix/lib/mix/lock.ex | 9 +++++++++
 1 file changed, 9 insertions(+)

diff --git a/lib/mix/lib/mix/lock.ex b/lib/mix/lib/mix/lock.ex
index 83cd0d3325a..5fbcbc43974 100644
--- a/lib/mix/lib/mix/lock.ex
+++ b/lib/mix/lib/mix/lock.ex
@@ -200,6 +200,9 @@ defmodule Mix.Lock do
         _ = :gen_tcp.send(socket, @probe_data)
         accept_loop(listen_socket)
 
+      {:error, :eintr} ->
+        accept_loop(listen_socket)
+
       {:error, reason} when reason in [:closed, :einval] ->
         :ok
     end
@@ -244,6 +247,9 @@ defmodule Mix.Lock do
         :gen_tcp.close(socket)
         {:error, :unexpected_port_owner}
 
+      {:error, :eintr} ->
+        await_probe_data(socket)
+
       {:error, reason} ->
         :gen_tcp.close(socket)
         {:error, reason}
@@ -272,6 +278,9 @@ defmodule Mix.Lock do
       {:error, :closed} ->
         :ok
 
+      {:error, :eintr} ->
+        await_close(socket)
+
       {:error, _other} ->
         # In case of an unexpected error, we close the socket ourselves
         # to retry

From b9b56fccacedfdeba7c47d8b4812c4b184f20ef1 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jonatan=20K=C5=82osko?= <jonatanklosko@gmail.com>
Date: Fri, 27 Sep 2024 14:31:01 +0200
Subject: [PATCH 11/20] Apply suggestions from code review

Co-authored-by: Andrea Leopardi <an.leopardi@gmail.com>
---
 lib/mix/lib/mix/lock.ex | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/lib/mix/lib/mix/lock.ex b/lib/mix/lib/mix/lock.ex
index 5fbcbc43974..2d6c3d90b9c 100644
--- a/lib/mix/lib/mix/lock.ex
+++ b/lib/mix/lib/mix/lock.ex
@@ -125,14 +125,14 @@ defmodule Mix.Lock do
     end
   end
 
-  defp listen() do
+  defp listen do
     with {:ok, socket} <- :gen_tcp.listen(0, @listen_opts) do
       case :inet.port(socket) do
         {:ok, port} ->
           {:ok, socket, port}
 
         {:error, reason} ->
-          :socket.close(socket)
+          :gen_tcp.close(socket)
           {:error, reason}
       end
     end
@@ -200,6 +200,7 @@ defmodule Mix.Lock do
         _ = :gen_tcp.send(socket, @probe_data)
         accept_loop(listen_socket)
 
+      # eintr is "Interrupted system call".
       {:error, :eintr} ->
         accept_loop(listen_socket)
 

From a11d1c3dda8ca12ae660859cae0f470a01f5d5f4 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jonatan=20K=C5=82osko?= <jonatanklosko@gmail.com>
Date: Fri, 27 Sep 2024 19:31:47 +0700
Subject: [PATCH 12/20] Use Mix.raise/1

---
 lib/mix/lib/mix/lock.ex | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/lib/mix/lib/mix/lock.ex b/lib/mix/lib/mix/lock.ex
index 2d6c3d90b9c..4238f975c03 100644
--- a/lib/mix/lib/mix/lock.ex
+++ b/lib/mix/lib/mix/lock.ex
@@ -120,8 +120,9 @@ defmodule Mix.Lock do
         end
 
       {:error, reason} ->
-        raise Mix.Error,
-              "failed to open a TCP socket while acquiring a lock, reason: #{inspect(reason)}"
+        Mix.raise(
+          "failed to open a TCP socket while acquiring a lock, reason: #{inspect(reason)}"
+        )
     end
   end
 

From 0ab81d75736aaa812da27aa23a6ef57f5f8ed576 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jonatan=20K=C5=82osko?= <jonatanklosko@gmail.com>
Date: Fri, 27 Sep 2024 14:40:18 +0200
Subject: [PATCH 13/20] Update lib/mix/lib/mix/lock.ex

Co-authored-by: Andrea Leopardi <an.leopardi@gmail.com>
---
 lib/mix/lib/mix/lock.ex | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/lib/mix/lib/mix/lock.ex b/lib/mix/lib/mix/lock.ex
index 4238f975c03..d8eead8118e 100644
--- a/lib/mix/lib/mix/lock.ex
+++ b/lib/mix/lib/mix/lock.ex
@@ -66,6 +66,8 @@ defmodule Mix.Lock do
   @loopback {127, 0, 0, 1}
   @listen_opts [:binary, ip: @loopback, packet: :raw, nodelay: true, backlog: 128, active: false]
   @connect_opts [:binary, packet: :raw, nodelay: true, active: false]
+  # The probe data needs to be small enough that it will for sure be
+  # sent in a single packet by the socket.
   @probe_data "elixirlock"
   @probe_timeout_ms 5_000
 

From 624d56e2b7de17ec318a31430501b6843e15cada Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jonatan=20K=C5=82osko?= <jonatanklosko@gmail.com>
Date: Fri, 27 Sep 2024 22:26:35 +0700
Subject: [PATCH 14/20] Up

---
 lib/mix/lib/mix/lock.ex | 9 +++++----
 1 file changed, 5 insertions(+), 4 deletions(-)

diff --git a/lib/mix/lib/mix/lock.ex b/lib/mix/lib/mix/lock.ex
index d8eead8118e..f5a75801dc2 100644
--- a/lib/mix/lib/mix/lock.ex
+++ b/lib/mix/lib/mix/lock.ex
@@ -152,10 +152,7 @@ defmodule Mix.Lock do
         %{socket: socket, path: path}
 
       {:ok, _n} ->
-        # We grabbed lock_1+, so we need to replace lock_0 and clean
-        # up. This must happen in a precise order, so if anything
-        # fails, we keep the files as is and the next process that
-        # grabs the lock will do the cleanup
+        # We grabbed lock_1+, so we need to replace lock_0 and clean up
         take_over(path, port_path)
         %{socket: socket, path: path}
 
@@ -261,6 +258,10 @@ defmodule Mix.Lock do
   end
 
   defp take_over(path, port_path) do
+    # The operations here must happen in precise order, so if anything
+    # fails, we keep the files as is and the next process that grabs
+    # the lock will do the cleanup
+
     lock_path = Path.join(path, "lock_0")
 
     # We linked to lock_N successfully, so port_path should exist

From 313868c2ab5d05844db5fe14af1a40f0a7c68f6f Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jonatan=20K=C5=82osko?= <jonatanklosko@gmail.com>
Date: Mon, 30 Sep 2024 12:46:49 +0200
Subject: [PATCH 15/20] Update lib/mix/lib/mix/lock.ex

Co-authored-by: Jean Klingler <sabiwara@proton.me>
---
 lib/mix/lib/mix/lock.ex | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/lib/mix/lib/mix/lock.ex b/lib/mix/lib/mix/lock.ex
index f5a75801dc2..cfe3adb0357 100644
--- a/lib/mix/lib/mix/lock.ex
+++ b/lib/mix/lib/mix/lock.ex
@@ -80,7 +80,7 @@ defmodule Mix.Lock do
   This function can also be called if this process already has the
   lock. In such case the function is executed immediately.
   """
-  @spec with_lock(iodata(), (-> term())) :: :ok
+  @spec with_lock(iodata(), (-> term())) :: term()
   def with_lock(key, fun) do
     key = key |> :erlang.md5() |> Base.url_encode64(padding: false)
     path = Path.join([System.tmp_dir!(), "mix_lock", key])

From b76f79589f6b8d888574405bfdcb8f32acaa9734 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jonatan=20K=C5=82osko?= <jonatanklosko@gmail.com>
Date: Mon, 30 Sep 2024 17:48:51 +0700
Subject: [PATCH 16/20] Specify message size on receive

---
 lib/mix/lib/mix/lock.ex | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

diff --git a/lib/mix/lib/mix/lock.ex b/lib/mix/lib/mix/lock.ex
index cfe3adb0357..b3998bad4c2 100644
--- a/lib/mix/lib/mix/lock.ex
+++ b/lib/mix/lib/mix/lock.ex
@@ -66,9 +66,8 @@ defmodule Mix.Lock do
   @loopback {127, 0, 0, 1}
   @listen_opts [:binary, ip: @loopback, packet: :raw, nodelay: true, backlog: 128, active: false]
   @connect_opts [:binary, packet: :raw, nodelay: true, active: false]
-  # The probe data needs to be small enough that it will for sure be
-  # sent in a single packet by the socket.
   @probe_data "elixirlock"
+  @probe_data_size byte_size(@probe_data)
   @probe_timeout_ms 5_000
 
   @doc """
@@ -240,7 +239,7 @@ defmodule Mix.Lock do
   end
 
   defp await_probe_data(socket) do
-    case :gen_tcp.recv(socket, 0, @probe_timeout_ms) do
+    case :gen_tcp.recv(socket, @probe_data_size, @probe_timeout_ms) do
       {:ok, @probe_data} ->
         {:ok, socket}
 

From 18fc73320decd0470651cd96f87f5aac14f8df04 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jonatan=20K=C5=82osko?= <jonatanklosko@gmail.com>
Date: Fri, 4 Oct 2024 21:40:47 +0800
Subject: [PATCH 17/20] Add comment

---
 lib/mix/lib/mix/lock.ex | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/lib/mix/lib/mix/lock.ex b/lib/mix/lib/mix/lock.ex
index b3998bad4c2..378c4173731 100644
--- a/lib/mix/lib/mix/lock.ex
+++ b/lib/mix/lib/mix/lock.ex
@@ -226,7 +226,9 @@ defmodule Mix.Lock do
   defp connect(port) do
     # On Windows connecting to an unbound port takes a few seconds to
     # fail, so instead we shortcut the check by attempting a listen,
-    # which succeeds or fails immediately
+    # which succeeds or fails immediately. Note that `reuseaddr` here
+    # ensures that if the listening socket closed recently, we can
+    # immediately reclaim the same port.
     case :gen_tcp.listen(port, [reuseaddr: true] ++ @listen_opts) do
       {:ok, socket} ->
         :gen_tcp.close(socket)

From 4ddd5b22e69b47730804d990ef485cac9818634e Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jonatan=20K=C5=82osko?= <jonatanklosko@gmail.com>
Date: Tue, 8 Oct 2024 02:23:59 +0800
Subject: [PATCH 18/20] Move to Mix.Task.Compiler.with_lock/2

---
 lib/mix/lib/mix/task.compiler.ex        | 14 ++++++++++++++
 lib/mix/lib/mix/tasks/compile.all.ex    | 10 +++++-----
 lib/mix/lib/mix/tasks/compile.elixir.ex |  2 +-
 lib/mix/lib/mix/tasks/compile.erlang.ex |  5 +++--
 lib/mix/lib/mix/tasks/compile.ex        | 13 -------------
 lib/mix/lib/mix/tasks/deps.compile.ex   |  4 +++-
 6 files changed, 26 insertions(+), 22 deletions(-)

diff --git a/lib/mix/lib/mix/task.compiler.ex b/lib/mix/lib/mix/task.compiler.ex
index cd0856d4202..6392bd1d691 100644
--- a/lib/mix/lib/mix/task.compiler.ex
+++ b/lib/mix/lib/mix/task.compiler.ex
@@ -162,4 +162,18 @@ defmodule Mix.Task.Compiler do
         {:noop, []}
     end
   end
+
+  @doc false
+  def with_lock(config, fun) do
+    # To avoid duplicated compilation, we wrap compilation tasks, such
+    # as compile.all, deps.compile, compile.elixir, compile.erlang in
+    # a lock. Note that compile.all covers compile.elixir, but the
+    # latter can still be invoked directly, hence we put the lock over
+    # the individual tasks.
+
+    build_path = Mix.Project.build_path(config)
+    lock_key = ["compile", build_path]
+
+    Mix.Lock.with_lock(lock_key, fun)
+  end
 end
diff --git a/lib/mix/lib/mix/tasks/compile.all.ex b/lib/mix/lib/mix/tasks/compile.all.ex
index 85c2f2813fe..e16d0cde81d 100644
--- a/lib/mix/lib/mix/tasks/compile.all.ex
+++ b/lib/mix/lib/mix/tasks/compile.all.ex
@@ -13,14 +13,14 @@ defmodule Mix.Tasks.Compile.All do
   def run(args) do
     Mix.Project.get!()
 
-    Mix.Lock.with_lock(Mix.Tasks.Compile.compile_lock_key(), fn ->
-      do_run(args)
+    config = Mix.Project.config()
+
+    Mix.Task.Compiler.with_lock(config, fn ->
+      do_run(config, args)
     end)
   end
 
-  defp do_run(args) do
-    config = Mix.Project.config()
-
+  defp do_run(config, args) do
     # Compute the app cache if it is stale and we are
     # not compiling from a dependency.
     app_cache =
diff --git a/lib/mix/lib/mix/tasks/compile.elixir.ex b/lib/mix/lib/mix/tasks/compile.elixir.ex
index 54dfd792ba4..9c15ab246fe 100644
--- a/lib/mix/lib/mix/tasks/compile.elixir.ex
+++ b/lib/mix/lib/mix/tasks/compile.elixir.ex
@@ -126,7 +126,7 @@ defmodule Mix.Tasks.Compile.Elixir do
     # so we wrap the compiler in a lock.
 
     with_logger_app(project, fn ->
-      Mix.Lock.with_lock(Mix.Tasks.Compile.compile_lock_key(), fn ->
+      Mix.Task.Compiler.with_lock(project, fn ->
         Mix.Compilers.Elixir.compile(
           manifest,
           srcs,
diff --git a/lib/mix/lib/mix/tasks/compile.erlang.ex b/lib/mix/lib/mix/tasks/compile.erlang.ex
index 7364e9b692f..3b88bf7ef58 100644
--- a/lib/mix/lib/mix/tasks/compile.erlang.ex
+++ b/lib/mix/lib/mix/tasks/compile.erlang.ex
@@ -47,8 +47,9 @@ defmodule Mix.Tasks.Compile.Erlang do
   def run(args) do
     {opts, _, _} = OptionParser.parse(args, switches: @switches)
 
-    Mix.Lock.with_lock(Mix.Tasks.Compile.compile_lock_key(), fn ->
-      project = Mix.Project.config()
+    project = Mix.Project.config()
+
+    Mix.Task.Compiler.with_lock(project, fn ->
       source_paths = project[:erlc_paths]
       Mix.Compilers.Erlang.assert_valid_erlc_paths(source_paths)
       files = Mix.Utils.extract_files(source_paths, [:erl])
diff --git a/lib/mix/lib/mix/tasks/compile.ex b/lib/mix/lib/mix/tasks/compile.ex
index 29145342f60..a893fc19227 100644
--- a/lib/mix/lib/mix/tasks/compile.ex
+++ b/lib/mix/lib/mix/tasks/compile.ex
@@ -243,19 +243,6 @@ defmodule Mix.Tasks.Compile do
     end)
   end
 
-  @doc false
-  def compile_lock_key() do
-    # To avoid duplicated compilation, we wrap compilation tasks, such
-    # as compile.all, deps.compile, compile.elixir, compile.erlang in
-    # a lock. Note that compile.all covers compile.elixir, but the
-    # latter can still be invoked directly, hence we put the lock over
-    # the individual tasks.
-
-    config = Mix.Project.config()
-    build_path = Mix.Project.build_path(config)
-    ["compile", build_path]
-  end
-
   ## Consolidation handling
 
   defp reconsolidate_protocols?(:ok), do: true
diff --git a/lib/mix/lib/mix/tasks/deps.compile.ex b/lib/mix/lib/mix/tasks/deps.compile.ex
index b579b33e732..246ecd843b4 100644
--- a/lib/mix/lib/mix/tasks/deps.compile.ex
+++ b/lib/mix/lib/mix/tasks/deps.compile.ex
@@ -54,7 +54,9 @@ defmodule Mix.Tasks.Deps.Compile do
 
     Mix.Project.get!()
 
-    Mix.Lock.with_lock(Mix.Tasks.Compile.compile_lock_key(), fn ->
+    config = Mix.Project.config()
+
+    Mix.Task.Compiler.with_lock(config, fn ->
       deps = Mix.Dep.load_and_cache()
 
       case OptionParser.parse(args, switches: @switches) do

From d9a6eea91998d2f921918b89d2bddf62da8da227 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jonatan=20K=C5=82osko?= <jonatanklosko@gmail.com>
Date: Tue, 8 Oct 2024 14:03:30 +0800
Subject: [PATCH 19/20] Show message when the lock is held by a different
 process

---
 lib/mix/lib/mix/lock.ex          | 46 +++++++++++++++++++++-----------
 lib/mix/lib/mix/task.compiler.ex |  7 +++--
 lib/mix/test/mix/lock_test.exs   | 26 ++++++++++++++++++
 3 files changed, 61 insertions(+), 18 deletions(-)

diff --git a/lib/mix/lib/mix/lock.ex b/lib/mix/lib/mix/lock.ex
index 378c4173731..855bb891633 100644
--- a/lib/mix/lib/mix/lock.ex
+++ b/lib/mix/lib/mix/lock.ex
@@ -78,9 +78,20 @@ defmodule Mix.Lock do
 
   This function can also be called if this process already has the
   lock. In such case the function is executed immediately.
+
+  ## Options
+
+    * `:on_taken` - a one-arity function called if the lock is held
+      by a different process. The operating system PID of that process
+      is given as the first argument (as a string). This function may
+      be called multiple times, if the lock owner changes, until it
+      is successfully acquired by this process.
+
   """
-  @spec with_lock(iodata(), (-> term())) :: term()
-  def with_lock(key, fun) do
+  @spec with_lock(iodata(), (-> term()), keyword()) :: term()
+  def with_lock(key, fun, opts \\ []) do
+    opts = Keyword.validate!(opts, [:on_taken])
+
     key = key |> :erlang.md5() |> Base.url_encode64(padding: false)
     path = Path.join([System.tmp_dir!(), "mix_lock", key])
 
@@ -90,7 +101,7 @@ defmodule Mix.Lock do
     if has_lock? do
       fun.()
     else
-      lock = lock(path)
+      lock = lock(path, opts[:on_taken])
       Process.put(pdict_key, true)
 
       try do
@@ -104,7 +115,7 @@ defmodule Mix.Lock do
     end
   end
 
-  defp lock(path) do
+  defp lock(path, on_taken) do
     File.mkdir_p!(path)
 
     case listen() do
@@ -112,7 +123,7 @@ defmodule Mix.Lock do
         spawn_link(fn -> accept_loop(socket) end)
 
         try do
-          try_lock(path, socket, port)
+          try_lock(path, socket, port, on_taken)
         rescue
           exception ->
             # Close the socket to make sure we don't block the lock
@@ -140,10 +151,11 @@ defmodule Mix.Lock do
     end
   end
 
-  defp try_lock(path, socket, port) do
+  defp try_lock(path, socket, port, on_taken) do
     port_path = Path.join(path, "port_#{port}")
+    os_pid = System.pid()
 
-    File.write!(port_path, <<port::unsigned-integer-32>>, [:raw])
+    File.write!(port_path, <<port::unsigned-integer-32, os_pid::binary>>, [:raw])
 
     case grab_lock(path, port_path, 0) do
       {:ok, 0} ->
@@ -155,13 +167,14 @@ defmodule Mix.Lock do
         take_over(path, port_path)
         %{socket: socket, path: path}
 
-      {:taken, probe_socket} ->
+      {:taken, probe_socket, os_pid} ->
         # Another process has the lock, wait for close and start over
+        if on_taken, do: on_taken.(os_pid)
         await_close(probe_socket)
-        try_lock(path, socket, port)
+        try_lock(path, socket, port, on_taken)
 
       :invalidated ->
-        try_lock(path, socket, port)
+        try_lock(path, socket, port, on_taken)
     end
   end
 
@@ -174,8 +187,8 @@ defmodule Mix.Lock do
 
       {:error, :eexist} ->
         case probe(lock_path) do
-          {:ok, probe_socket} ->
-            {:taken, probe_socket}
+          {:ok, probe_socket, os_pid} ->
+            {:taken, probe_socket, os_pid}
 
           {:error, _reason} ->
             grab_lock(path, port_path, n + 1)
@@ -209,16 +222,17 @@ defmodule Mix.Lock do
   end
 
   defp probe(port_path) do
-    with {:ok, port} <- fetch_probe_port(port_path),
-         {:ok, socket} <- connect(port) do
-      await_probe_data(socket)
+    with {:ok, port, os_pid} <- fetch_probe_port(port_path),
+         {:ok, socket} <- connect(port),
+         {:ok, socket} <- await_probe_data(socket) do
+      {:ok, socket, os_pid}
     end
   end
 
   defp fetch_probe_port(port_path) do
     case File.read(port_path) do
       {:ok, <<0::unsigned-integer-32>>} -> {:error, :ignore}
-      {:ok, <<port::unsigned-integer-32>>} -> {:ok, port}
+      {:ok, <<port::unsigned-integer-32, os_pid::binary>>} -> {:ok, port, os_pid}
       {:error, reason} -> {:error, reason}
     end
   end
diff --git a/lib/mix/lib/mix/task.compiler.ex b/lib/mix/lib/mix/task.compiler.ex
index 6392bd1d691..b8ddd275112 100644
--- a/lib/mix/lib/mix/task.compiler.ex
+++ b/lib/mix/lib/mix/task.compiler.ex
@@ -172,8 +172,11 @@ defmodule Mix.Task.Compiler do
     # the individual tasks.
 
     build_path = Mix.Project.build_path(config)
-    lock_key = ["compile", build_path]
 
-    Mix.Lock.with_lock(lock_key, fun)
+    on_taken = fn os_pid ->
+      Mix.shell().info("Waiting for lock on the build directory (held by process #{os_pid})")
+    end
+
+    Mix.Lock.with_lock(build_path, fun, on_taken: on_taken)
   end
 end
diff --git a/lib/mix/test/mix/lock_test.exs b/lib/mix/test/mix/lock_test.exs
index 0b02d60813b..1c53b0217ff 100644
--- a/lib/mix/test/mix/lock_test.exs
+++ b/lib/mix/test/mix/lock_test.exs
@@ -151,6 +151,32 @@ defmodule Mix.LockTest do
            end) == :still_works!
   end
 
+  test "calls :on_taken when the lock is held by a different process" do
+    parent = self()
+
+    {pid, ref} =
+      spawn_monitor(fn ->
+        Mix.Lock.with_lock(@lock_key, fn ->
+          send(parent, :locked)
+          assert_receive :will_lock
+        end)
+      end)
+
+    assert_receive :locked
+
+    on_taken = fn os_pid ->
+      send(pid, :will_lock)
+      send(self(), {:on_taken_called, os_pid})
+    end
+
+    assert Mix.Lock.with_lock(@lock_key, fn -> :it_works! end, on_taken: on_taken) == :it_works!
+
+    os_pid = System.pid()
+    assert_receive {:on_taken_called, ^os_pid}
+
+    assert_receive {:DOWN, ^ref, _, _, _}
+  end
+
   defp await_monitors([]), do: :ok
 
   defp await_monitors(refs) do

From 162c42bb87895e2ecfaca55feafea3f2b9139611 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jonatan=20K=C5=82osko?= <jonatanklosko@gmail.com>
Date: Tue, 8 Oct 2024 14:10:10 +0800
Subject: [PATCH 20/20] Mix.Task.Compiler.with_lock ->
 Mix.Project.with_build_lock

---
 lib/mix/lib/mix/project.ex              | 17 +++++++++++++++++
 lib/mix/lib/mix/task.compiler.ex        | 17 -----------------
 lib/mix/lib/mix/tasks/compile.all.ex    |  2 +-
 lib/mix/lib/mix/tasks/compile.elixir.ex |  2 +-
 lib/mix/lib/mix/tasks/compile.erlang.ex |  2 +-
 lib/mix/lib/mix/tasks/deps.compile.ex   |  2 +-
 6 files changed, 21 insertions(+), 21 deletions(-)

diff --git a/lib/mix/lib/mix/project.ex b/lib/mix/lib/mix/project.ex
index 6d9988d8ee9..d7ae332754d 100644
--- a/lib/mix/lib/mix/project.ex
+++ b/lib/mix/lib/mix/project.ex
@@ -882,6 +882,23 @@ defmodule Mix.Project do
     end
   end
 
+  @doc false
+  def with_build_lock(config \\ config(), fun) do
+    # To avoid duplicated compilation, we wrap compilation tasks, such
+    # as compile.all, deps.compile, compile.elixir, compile.erlang in
+    # a lock. Note that compile.all covers compile.elixir, but the
+    # latter can still be invoked directly, so we put the lock over
+    # each individual task.
+
+    build_path = build_path(config)
+
+    on_taken = fn os_pid ->
+      Mix.shell().info("Waiting for lock on the build directory (held by process #{os_pid})")
+    end
+
+    Mix.Lock.with_lock(build_path, fun, on_taken: on_taken)
+  end
+
   # Loads mix.exs in the current directory or loads the project from the
   # mixfile cache and pushes the project onto the project stack.
   defp load_project(app, post_config) do
diff --git a/lib/mix/lib/mix/task.compiler.ex b/lib/mix/lib/mix/task.compiler.ex
index b8ddd275112..cd0856d4202 100644
--- a/lib/mix/lib/mix/task.compiler.ex
+++ b/lib/mix/lib/mix/task.compiler.ex
@@ -162,21 +162,4 @@ defmodule Mix.Task.Compiler do
         {:noop, []}
     end
   end
-
-  @doc false
-  def with_lock(config, fun) do
-    # To avoid duplicated compilation, we wrap compilation tasks, such
-    # as compile.all, deps.compile, compile.elixir, compile.erlang in
-    # a lock. Note that compile.all covers compile.elixir, but the
-    # latter can still be invoked directly, hence we put the lock over
-    # the individual tasks.
-
-    build_path = Mix.Project.build_path(config)
-
-    on_taken = fn os_pid ->
-      Mix.shell().info("Waiting for lock on the build directory (held by process #{os_pid})")
-    end
-
-    Mix.Lock.with_lock(build_path, fun, on_taken: on_taken)
-  end
 end
diff --git a/lib/mix/lib/mix/tasks/compile.all.ex b/lib/mix/lib/mix/tasks/compile.all.ex
index e16d0cde81d..c4b73477bda 100644
--- a/lib/mix/lib/mix/tasks/compile.all.ex
+++ b/lib/mix/lib/mix/tasks/compile.all.ex
@@ -15,7 +15,7 @@ defmodule Mix.Tasks.Compile.All do
 
     config = Mix.Project.config()
 
-    Mix.Task.Compiler.with_lock(config, fn ->
+    Mix.Project.with_build_lock(config, fn ->
       do_run(config, args)
     end)
   end
diff --git a/lib/mix/lib/mix/tasks/compile.elixir.ex b/lib/mix/lib/mix/tasks/compile.elixir.ex
index 9c15ab246fe..7ae75848361 100644
--- a/lib/mix/lib/mix/tasks/compile.elixir.ex
+++ b/lib/mix/lib/mix/tasks/compile.elixir.ex
@@ -126,7 +126,7 @@ defmodule Mix.Tasks.Compile.Elixir do
     # so we wrap the compiler in a lock.
 
     with_logger_app(project, fn ->
-      Mix.Task.Compiler.with_lock(project, fn ->
+      Mix.Project.with_build_lock(project, fn ->
         Mix.Compilers.Elixir.compile(
           manifest,
           srcs,
diff --git a/lib/mix/lib/mix/tasks/compile.erlang.ex b/lib/mix/lib/mix/tasks/compile.erlang.ex
index 3b88bf7ef58..4667e15a71b 100644
--- a/lib/mix/lib/mix/tasks/compile.erlang.ex
+++ b/lib/mix/lib/mix/tasks/compile.erlang.ex
@@ -49,7 +49,7 @@ defmodule Mix.Tasks.Compile.Erlang do
 
     project = Mix.Project.config()
 
-    Mix.Task.Compiler.with_lock(project, fn ->
+    Mix.Project.with_build_lock(project, fn ->
       source_paths = project[:erlc_paths]
       Mix.Compilers.Erlang.assert_valid_erlc_paths(source_paths)
       files = Mix.Utils.extract_files(source_paths, [:erl])
diff --git a/lib/mix/lib/mix/tasks/deps.compile.ex b/lib/mix/lib/mix/tasks/deps.compile.ex
index 246ecd843b4..89bed61743e 100644
--- a/lib/mix/lib/mix/tasks/deps.compile.ex
+++ b/lib/mix/lib/mix/tasks/deps.compile.ex
@@ -56,7 +56,7 @@ defmodule Mix.Tasks.Deps.Compile do
 
     config = Mix.Project.config()
 
-    Mix.Task.Compiler.with_lock(config, fn ->
+    Mix.Project.with_build_lock(config, fn ->
       deps = Mix.Dep.load_and_cache()
 
       case OptionParser.parse(args, switches: @switches) do