Add Crystal::FdLock to count system fd references

This patch implements a reference counted lock to protect IO objects
that depend on a reusable system fd (IO::FileDescriptor, File and
Socket) to protect them against thread safety issues around close:

- Thread 1 wants to read from fd 123;
- The OS preempts Thread 1;
- Thread 2 closes fd 123;
- Thread 2 opens something else and the OS reuses fd 123;
- The OS resumes Thread 1;
- Thread 1 reads from the reused fd 123!!!

The issue arises for any operation that would mutate the fd: write,
fchown, ftruncate, setsockopt, ... as they risk affecting a reused fd
instead of the expected one.

Author: ysbaddaden

spec/std/crystal/fd_lock_spec.cr

@@ -0,0 +1,137 @@
+require "spec"
+require "../../support/interpreted.cr"
+require "crystal/fd_lock"
+require "wait_group"
+
+describe Crystal::FdLock do
+  describe "#reference" do
+    it "acquires" do
+      lock = Crystal::FdLock.new
+      called = 0
+
+      lock.reference { called += 1 }
+      lock.reference { called += 1 }
+
+      called.should eq(2)
+    end
+
+    it "allows reentrancy (side effect)" do
+      lock = Crystal::FdLock.new
+      called = 0
+
+      lock.reference { called += 1 }
+      lock.reference do
+        lock.reference { called += 1 }
+      end
+
+      called.should eq(2)
+    end
+
+    it "acquires shared reference" do
+      lock = Crystal::FdLock.new
+
+      ready = WaitGroup.new(1)
+      release = Channel(String).new
+
+      spawn do
+        lock.reference do
+          ready.done
+
+          select
+          when release.send("ok")
+          when timeout(1.second)
+            release.send("timeout")
+          end
+        end
+      end
+
+      ready.wait
+      lock.reference { }
+
+      release.receive.should eq("ok")
+    end
+
+    it "raises when closed" do
+      lock = Crystal::FdLock.new
+      lock.try_close? { }
+
+      called = false
+      expect_raises(IO::Error, "Closed") do
+        lock.reference { called = true }
+      end
+
+      called.should be_false
+    end
+  end
+
+  describe "#try_close?" do
+    it "closes" do
+      lock = Crystal::FdLock.new
+      lock.closed?.should be_false
+
+      called = false
+      lock.try_close? { called = true }.should be_true
+      lock.closed?.should be_true
+      called.should be_true
+    end
+
+    it "closes once" do
+      lock = Crystal::FdLock.new
+
+      called = 0
+
+      WaitGroup.wait do |wg|
+        10.times do
+          wg.spawn do
+            lock.try_close? { called += 1 }
+            lock.try_close? { called += 1 }
+          end
+        end
+      end
+
+      called.should eq(1)
+    end
+
+    # FIXME: the interpreter segfaults while running this spec (NULL pointer)
+    pending_interpreted "waits for all references to return" do
+      lock = Crystal::FdLock.new
+
+      ready = WaitGroup.new(10)
+      exceptions = Channel(Exception).new(10)
+
+      WaitGroup.wait do |wg|
+        10.times do
+          wg.spawn do
+            begin
+              lock.reference do
+                ready.done
+                Fiber.yield
+              end
+            rescue ex
+              exceptions.send(ex)
+            end
+          end
+        end
+
+        ready.wait
+
+        called = false
+        lock.try_close? { called = true }.should be_true
+        lock.closed?.should be_true
+        called.should be_true
+      end
+      exceptions.close
+
+      if ex = exceptions.receive?
+        raise ex
+      end
+    end
+  end
+
+  it "#reset" do
+    lock = Crystal::FdLock.new
+    lock.try_close? { }
+    lock.reset
+    lock.try_close? { }.should eq(true)
+  end
+end

src/crystal/fd_lock.cr

@@ -0,0 +1,5 @@
+{% if flag?(:preview_mt) %}
+  require "./fd_lock_mt"
+{% else %}
+  require "./fd_lock_no_mt"
+{% end %}

src/crystal/fd_lock_mt.cr

@@ -0,0 +1,117 @@
+# The general design is influenced by fdMutex in Go (LICENSE: BSD 3-Clause,
+# Copyright Google):
+# https://github.com/golang/go/blob/go1.25.1/src/internal/poll/fd_mutex.go
+
+# :nodoc:
+#
+# Tracks active references over a system file descriptor (fd) and serializes
+# reads and writes.
+#
+# Every access to the fd that may affect its system state or system buffers must
+# acquire a shared lock.
+#
+# The fdlock can be closed at any time, but the actual system close will wait
+# until there are no more references left. This avoids potential races when a
+# thread might try to read a fd that has been closed and has been reused by the
+# OS for example.
+#
+# FIXME: the interpreter segfaults when interpreted code uses this type; for now
+# it uses the thread unsafe alternative (fd_lock_no_mt).
+struct Crystal::FdLock
+  CLOSED = 1_u32 << 0 # the fdlock has been closed
+  REF    = 1_u32 << 1 # the reference counter increment
+  MASK   = ~(REF - 1) # mask for the reference counter
+
+  @m = Atomic(UInt32).new(0_u32)
+  @closing : Fiber?
+
+  # Borrows a reference for the duration of the block. Raises if the fdlock is
+  # closed while trying to borrow.
+  def reference(& : -> F) : F forall F
+    m, success = @m.compare_and_set(0_u32, REF, :acquire, :relaxed)
+    increment_slow(m) unless success
+
+    begin
+      yield
+    ensure
+      m = @m.sub(REF, :release)
+      handle_last_ref(m)
+    end
+  end
+
+  private def increment_slow(m)
+    while true
+      if (m & CLOSED) == CLOSED
+        raise IO::Error.new("Closed")
+      end
+      m, success = @m.compare_and_set(m, m + REF, :acquire, :relaxed)
+      break if success
+    end
+  end
+
+  private def handle_last_ref(m)
+    return unless (m & CLOSED) == CLOSED # is closed?
+    return unless (m & MASK) == REF      # was the last ref?
+
+    # the last ref after close is responsible to resume the closing fiber
+    @closing.not_nil!("BUG: expected a closing fiber to resume.").enqueue
+  end
+
+  # Closes the fdlock. Blocks for as long as there are references.
+  #
+  # The *callback* block must cancel any external waiters (e.g. pending evloop
+  # reads or writes).
+  #
+  # Returns true if the fdlock has been closed: no fiber can acquire a reference
+  # anymore, the calling fiber fully owns the fd and can safely close it.
+  #
+  # Returns false if the fdlock has already been closed: the calling fiber
+  # doesn't own the fd and musn't close it, as there might still be active
+  # references and another fiber will close anyway.
+  def try_close?(&callback : ->) : Bool
+    attempts = 0
+
+    while true
+      m = @m.get(:relaxed)
+
+      if (m & CLOSED) == CLOSED
+        # already closed: abort
+        return false
+      end
+
+      # close + increment ref
+      m, success = @m.compare_and_set(m, (m + REF) | CLOSED, :acquire, :relaxed)
+      break if success
+
+      attempts = Thread.delay(attempts)
+    end
+
+    # set the current fiber as the closing fiber (to be resumed by the last ref)
+    @closing = Fiber.current
+
+    # decrement the last ref
+    m = @m.sub(REF, :release)
+
+    begin
+      yield
+    ensure
+      # wait for the last ref... unless we're the last ref!
+      Fiber.suspend unless (m & MASK) == REF
+    end
+
+    @closing = nil
+    true
+  end
+
+  # Resets the fdlock back to its pristine state so it can be used again.
+  # Assumes the caller owns the fdlock. This is required by
+  # `TCPSocket#initialize`.
+  def reset : Nil
+    @m.lazy_set(0_u32)
+    @closing = nil
+  end
+
+  def closed? : Bool
+    (@m.get(:relaxed) & CLOSED) == CLOSED
+  end
+end

src/crystal/fd_lock_no_mt.cr

@@ -0,0 +1,35 @@
+# :nodoc:
+#
+# Simpler, but thread unsafe, alternative to Crystal::FdLock that only
+# serializes reads and writes and otherwise doesn't count references or waits
+# for references before closing. This is mostly needed for the interpreter that
+# happens to segfault with the thread safe alternative (see fd_lock_mt).
+struct Crystal::FdLock
+  CLOSED = 1_u8 << 0
+
+  @m = 0_u8
+
+  def reference(&)
+    raise IO::Error.new("Closed") if closed?
+    yield
+  end
+
+  def reset : Nil
+    @m = 0_u8
+  end
+
+  def closed?
+    (@m & CLOSED) == CLOSED
+  end
+
+  def try_close?(&)
+    if closed?
+      false
+    else
+      @m |= CLOSED
+
+      yield
+      true
+    end
+  end
+end