1# MicroPython uasyncio module
2# MIT license; Copyright (c) 2019-2020 Damien P. George
3
4from . import core
5
6# Lock class for primitive mutex capability
7class Lock:
8    def __init__(self):
9        # The state can take the following values:
10        # - 0: unlocked
11        # - 1: locked
12        # - <Task>: unlocked but this task has been scheduled to acquire the lock next
13        self.state = 0
14        # Queue of Tasks waiting to acquire this Lock
15        self.waiting = core.TaskQueue()
16
17    def locked(self):
18        return self.state == 1
19
20    def release(self):
21        if self.state != 1:
22            raise RuntimeError("Lock not acquired")
23        if self.waiting.peek():
24            # Task(s) waiting on lock, schedule next Task
25            self.state = self.waiting.pop_head()
26            core._task_queue.push_head(self.state)
27        else:
28            # No Task waiting so unlock
29            self.state = 0
30
31    async def acquire(self):
32        if self.state != 0:
33            # Lock unavailable, put the calling Task on the waiting queue
34            self.waiting.push_head(core.cur_task)
35            # Set calling task's data to the lock's queue so it can be removed if needed
36            core.cur_task.data = self.waiting
37            try:
38                yield
39            except core.CancelledError as er:
40                if self.state == core.cur_task:
41                    # Cancelled while pending on resume, schedule next waiting Task
42                    self.state = 1
43                    self.release()
44                raise er
45        # Lock available, set it as locked
46        self.state = 1
47        return True
48
49    async def __aenter__(self):
50        return await self.acquire()
51
52    async def __aexit__(self, exc_type, exc, tb):
53        return self.release()
54