From a1f2689652ad4d2a0bdcb283598505ce6edba751 Mon Sep 17 00:00:00 2001
From: Cameron <cameron@moodycamel.com>
Date: Wed, 11 May 2016 13:41:33 -0400
Subject: [PATCH] Added support for timeouts to lightweight semaphore

---
 common/sema.h | 97 ++++++++++++++++++++++++++++++++++++++++++++++++---
 1 file changed, 92 insertions(+), 5 deletions(-)

diff --git a/common/sema.h b/common/sema.h
index 62d17ba..1994353 100644
--- a/common/sema.h
+++ b/common/sema.h
@@ -8,6 +8,7 @@
 
 #include <atomic>
 #include <cassert>
+#include <cstdint>
 
 
 #if defined(_WIN32)
@@ -15,9 +16,8 @@
 // Semaphore (Windows)
 //---------------------------------------------------------
 
+#define NOMINMAX
 #include <windows.h>
-#undef min
-#undef max
 
 class Semaphore
 {
@@ -44,6 +44,16 @@ class Semaphore
         WaitForSingleObject(m_hSema, INFINITE);
     }
 
+    bool try_wait()
+    {
+        return WaitForSingleObject(m_hSema, 0) != RC_WAIT_TIMEOUT;
+    }
+
+    bool timed_wait(std::uint64_t usecs)
+    {
+        return WaitForSingleObject(m_hSema, (unsigned long)(usecs / 1000)) != RC_WAIT_TIMEOUT;
+    }
+
     void signal(int count = 1)
     {
         ReleaseSemaphore(m_hSema, count, NULL);
@@ -84,6 +94,23 @@ class Semaphore
         semaphore_wait(m_sema);
     }
 
+    bool try_wait()
+    {
+        return timed_wait(0);
+    }
+    
+    bool timed_wait(std::uint64_t timeout_usecs)
+    {
+        mach_timespec_t ts;
+        ts.tv_sec = timeout_usecs / 1000000;
+        ts.tv_nsec = (timeout_usecs % 1000000) * 1000;
+
+        // added in OSX 10.10: https://developer.apple.com/library/prerelease/mac/documentation/General/Reference/APIDiffsMacOSX10_10SeedDiff/modules/Darwin.html
+        kern_return_t rc = semaphore_timedwait(m_sema, ts);
+
+        return rc != KERN_OPERATION_TIMED_OUT;
+    }
+
     void signal()
     {
         semaphore_signal(m_sema);
@@ -105,6 +132,7 @@ class Semaphore
 //---------------------------------------------------------
 
 #include <semaphore.h>
+#include <cerrno>
 
 class Semaphore
 {
@@ -137,6 +165,42 @@ class Semaphore
         while (rc == -1 && errno == EINTR);
     }
 
+    bool try_wait()
+    {
+        int rc;
+        do
+        {
+            rc = sem_trywait(&m_sema);
+        }
+        while (rc == -1 && errno == EINTR);
+        return !(rc == -1 && errno == EAGAIN);
+    }
+
+    bool timed_wait(std::uint64_t usecs)
+    {
+        struct timespec ts;
+        const int usecs_in_1_sec = 1000000;
+        const int nsecs_in_1_sec = 1000000000;
+        clock_gettime(CLOCK_REALTIME, &ts);
+        ts.tv_sec += usecs / usecs_in_1_sec;
+        ts.tv_nsec += (usecs % usecs_in_1_sec) * 1000;
+        // sem_timedwait bombs if you have more than 1e9 in tv_nsec
+        // so we have to clean things up before passing it in
+        if (ts.tv_nsec > nsecs_in_1_sec)
+        {
+            ts.tv_nsec -= nsecs_in_1_sec;
+            ++ts.tv_sec;
+        }
+
+        int rc;
+        do
+        {
+            rc = sem_timedwait(&m_sema, &ts);
+        }
+        while (rc == -1 && errno == EINTR);
+        return !(rc == -1 && errno == ETIMEDOUT);
+    }
+
     void signal()
     {
         sem_post(&m_sema);
@@ -168,7 +232,7 @@ class LightweightSemaphore
     std::atomic<int> m_count;
     Semaphore m_sema;
 
-    void waitWithPartialSpinning()
+    bool waitWithPartialSpinning(std::int64_t timeout_usecs = -1)
     {
         int oldCount;
         // Is there a better way to set the initial spin count?
@@ -179,13 +243,31 @@ class LightweightSemaphore
         {
             oldCount = m_count.load(std::memory_order_relaxed);
             if ((oldCount > 0) && m_count.compare_exchange_strong(oldCount, oldCount - 1, std::memory_order_acquire))
-                return;
+                return true;
             std::atomic_signal_fence(std::memory_order_acquire);     // Prevent the compiler from collapsing the loop.
         }
         oldCount = m_count.fetch_sub(1, std::memory_order_acquire);
-        if (oldCount <= 0)
+        if (oldCount > 0)
+            return true;
+        if (timeout_usecs < 0)
         {
             m_sema.wait();
+            return true;
+        }
+        if (m_sema.timed_wait((std::uint64_t)timeout_usecs))
+            return true;
+        // At this point, we've timed out waiting for the semaphore, but the
+        // count is still decremented indicating we may still be waiting on
+        // it. So we have to re-adjust the count, but only if the semaphore
+        // wasn't signaled enough times for us too since then. If it was, we
+        // need to release the semaphore too.
+        while (true)
+        {
+            oldCount = m_count.load(std::memory_order_acquire);
+            if (oldCount >= 0 && m_sema.try_wait())
+                return true;
+            if (oldCount < 0 && m_count.compare_exchange_strong(oldCount, oldCount + 1, std::memory_order_relaxed))
+                return false;
         }
     }
 
@@ -207,6 +289,11 @@ class LightweightSemaphore
             waitWithPartialSpinning();
     }
 
+    bool wait(std::int64_t timeout_usecs)
+    {
+        return tryWait() || waitWithPartialSpinning(timeout_usecs);
+    }
+
     void signal(int count = 1)
     {
         int oldCount = m_count.fetch_add(count, std::memory_order_release);