From 7cb28186f2b963246e87adefcb85489fe3f02a6d Mon Sep 17 00:00:00 2001
From: eldritch horrors <pennae@lix.systems>
Date: Tue, 11 Feb 2025 01:06:15 +0100
Subject: [PATCH] libutil: add promise-compatible Sync specialization

we can't use Sync with stl mutex types since yielding from one coroutine
to another while holding a lock is a deadlock hazard. if we don't expose
blocking lock operations from Sync though and emulate blocking with more
promises we can reuse most of the Sync code instead of rewriting it all.
hiding them completely would make it impossible to migrate current users
to promise-based code though, so we use our new NeverAsync lint instead.

Change-Id: I7eebf18e29473963040e7448da27c52012f1f134
---
 lix/libutil/sync.hh | 94 ++++++++++++++++++++++++++++++++++++++++++++-
 1 file changed, 93 insertions(+), 1 deletion(-)

diff --git a/lix/libutil/sync.hh b/lix/libutil/sync.hh
index 2c5424f2a..a6d9f4c0b 100644
--- a/lix/libutil/sync.hh
+++ b/lix/libutil/sync.hh
@@ -1,13 +1,20 @@
 #pragma once
 ///@file
 
+#include "lix/libutil/types.hh"
 #include <cstdlib>
+#include <kj/async.h>
+#include <kj/common.h>
+#include <list>
 #include <mutex>
 #include <condition_variable>
 #include <cassert>
+#include <optional>
 
 namespace nix {
 
+struct AsyncMutex;
+
 /**
  * This template class ensures synchronized access to a value of type
  * T. It is used as follows:
@@ -39,7 +46,7 @@ public:
 
     class Lock
     {
-    private:
+    protected:
         // Non-owning pointer. This would be an
         // optional<reference_wrapper<Sync>> if it didn't break gdb accessing
         // Lock values (as of 2024-06-15, gdb 14.2)
@@ -47,6 +54,7 @@ public:
         std::unique_lock<M> lk;
         friend Sync;
         Lock(Sync &s) : s(&s), lk(s.mutex) { }
+        Lock(Sync &s, std::unique_lock<M> lk) : s(&s), lk(std::move(lk)) { }
 
         inline void checkLockingInvariants()
         {
@@ -140,6 +148,90 @@ public:
      * Lock this Sync and return a RAII guard object.
      */
     Lock lock() { return Lock(*this); }
+
+    std::optional<Lock> tryLock()
+    {
+        if (std::unique_lock lk(mutex, std::try_to_lock_t{}); lk.owns_lock()) {
+            return Lock{*this, std::move(lk)};
+        } else {
+            return std::nullopt;
+        }
+    }
+};
+
+template<class T>
+class Sync<T, AsyncMutex> : private Sync<T, std::mutex>
+{
+private:
+    using base_type = Sync<T, std::mutex>;
+
+    std::mutex waitMutex;
+    std::list<kj::Own<kj::CrossThreadPromiseFulfiller<void>>> waiters;
+
+public:
+    Sync() = default;
+    Sync(T && data) : base_type(std::move(data)) {}
+
+    class Lock : private base_type::Lock
+    {
+        friend Sync;
+
+        Lock(base_type::Lock lk) : base_type::Lock(std::move(lk)) {}
+
+    public:
+        Lock(Lock &&) = default;
+        Lock & operator=(Lock &&) = default;
+
+        ~Lock()
+        {
+            if (this->lk.owns_lock()) {
+                this->lk.unlock();
+                auto * s = static_cast<Sync *>(this->s);
+                std::lock_guard wlk(s->waitMutex);
+                // wake them all. it's too hard to ensure liveness with promises
+                // that can be cancelled, and contention isn't usually that big.
+                for (auto & f : s->waiters) {
+                    f->fulfill();
+                }
+                s->waiters.clear();
+            }
+        }
+
+        using base_type::Lock::operator->, base_type::Lock::operator*;
+    };
+
+    auto lockSync(NeverAsync = {})
+    {
+        return base_type::lock();
+    }
+
+    kj::Promise<Lock> lock()
+    {
+        if (auto lk = tryLock()) {
+            co_return std::move(*lk);
+        }
+
+        while (true) {
+            auto pfp = kj::newPromiseAndCrossThreadFulfiller<void>();
+            {
+                std::lock_guard wlk(waitMutex);
+                waiters.push_back(std::move(pfp.fulfiller));
+            }
+            if (auto lk = tryLock()) {
+                co_return std::move(*lk);
+            }
+            co_await pfp.promise;
+        }
+    }
+
+    std::optional<Lock> tryLock()
+    {
+        if (auto lk = base_type::tryLock()) {
+            return Lock(std::move(*lk));
+        } else {
+            return std::nullopt;
+        }
+    }
 };
 
 }