Locking

Author: chrilves

raffiot/__internal.py

@@ -10,8 +10,10 @@
 
 from typing_extensions import final
 from functools import total_ordering
+import threading
+from queue import Queue
 
-__all__ = ["IOTag", "ContTag", "ResultTag", "Scheduled"]
+__all__ = ["IOTag", "ContTag", "ResultTag", "Scheduled", "Lock", "Semaphore"]
 
 
 @final
@@ -43,6 +45,7 @@ class IOTag(Enum):
     WAIT = 24  # FIBERS
     SLEEP_UNTIL = 25  # EPOCH IN SECONDS
     REC = 26  # FUN
+    LOCK = 27  # LOCK
 
 
 @final
@@ -86,3 +89,60 @@ def __lt__(self, other):
         if self.__schedule == other.__schedule:
             return hash(self.__fiber) < hash(other.__fiber)
         return self.__schedule < other.__schedule
+
+
+@final
+class Lock:
+
+    __slots__ = ["lock", "fiber", "__nb_taken", "waiting"]
+
+    def __init__(self):
+        self.lock = threading.Lock()
+        self.fiber = None
+        self.__nb_taken = 0
+        self.waiting = Queue()
+
+    def acquire(self, fiber) -> bool:
+        if self.fiber is None:
+            self.__nb_taken = 1
+            return True
+        if self.fiber is fiber:
+            self.fiber = fiber
+            self.__nb_taken += 1
+            return True
+        return False
+
+    def release(self):
+        with self.lock:
+            self.__nb_taken -= 1
+            if self.__nb_taken == 0:
+                if self.waiting.empty():
+                    self.fiber = None
+                    return
+                self.fiber = self.waiting.get()
+                self.__nb_taken = 1
+                self.fiber._Fiber__monitor._Monitor__resume(self.fiber)
+
+@final
+class Semaphore:
+
+    __slots__ = ["lock", "tokens", "waiting"]
+
+    def __init__(self, tokens: int):
+        self.lock = threading.Lock()
+        self.tokens = tokens
+        self.waiting = Queue()
+
+    def acquire(self, fiber) -> bool:
+        if self.tokens > 0:
+            self.tokens -= 1
+            return True
+        return False
+
+    def release(self):
+        with self.lock:
+            if self.waiting.empty():
+                self.tokens += 1
+                return
+            fiber = self.waiting.get()
+            fiber._Fiber__monitor._Monitor__resume(fiber)

raffiot/io.py

@@ -351,6 +351,8 @@ def __str__(self) -> str:
             return f"SleepUntil({self.__fields})"
         if self.__tag == IOTag.REC:
             return f"Rec({self.__fields})"
+        if self.__tag == IOTag.LOCK:
+            return f"Lock({self.__fields})"
 
     def __repr__(self):
         return str(self)
@@ -1078,6 +1080,24 @@ def callback(r):
                             arg_tag = ResultTag.PANIC
                             arg_value = exception
                             break
+                    if tag == IOTag.LOCK:
+                        lock = io._IO__fields
+                        try:
+                            with lock.lock:
+                                if lock.acquire(self):
+                                    arg_tag = ResultTag.OK
+                                    arg_value = None
+                                    break
+                                else:
+                                    self.__io = IO(IOTag.PURE, None)
+                                    self.__context = context
+                                    self.__cont = cont
+                                    lock.waiting.put(self)
+                                    return
+                        except Exception as exception:
+                            arg_tag = ResultTag.PANIC
+                            arg_value = exception
+                            break
                     arg_tag = ResultTag.PANIC
                     arg_value = _MatchError(f"{io} should be an IO")
                     break

raffiot/resource.py

@@ -15,6 +15,7 @@
 from raffiot import io, _MatchError
 from raffiot.io import IO
 from raffiot.result import Result, Ok, Error, Panic
+import __internal
 
 R = TypeVar("R")
 E = TypeVar("E")
@@ -593,3 +594,12 @@ def sleep(seconds: float) -> Resource[R, E, None]:
     :return:
     """
     return lift_io(io.sleep(seconds))
+
+
+def lock() -> Resource[None, None, None]:
+    new_lock = __internal.Lock()
+    return Resource(IO(__internal.IOTag.LOCK, new_lock).then(io.pure((None, io.defer(new_lock.release)))))
+
+def semaphore(tokens: int) -> Resource[None, None, None]:
+    new_semaphore = __internal.Semaphore(tokens)
+    return Resource(IO(__internal.IOTag.LOCK, new_semaphore).then(io.pure((None, io.defer(new_semaphore.release)))))