1  # SPDX-FileCopyrightText: 2019-2020 Damien P. George
 2  #
 3  # SPDX-License-Identifier: MIT
 4  #
 5  # MicroPython uasyncio module
 6  # MIT license; Copyright (c) 2019-2020 Damien P. George
 7  #
 8  # This code comes from MicroPython, and has not been run through black or pylint there.
 9  # Altering these files significantly would make merging difficult, so we will not use
10  # pylint or black.
11  # pylint: skip-file
12  # fmt: off
13  """
14  Locks
15  =====
16  """
17  
18  from . import core
19  
20  # Lock class for primitive mutex capability
21  class Lock:
22      """Create a new lock which can be used to coordinate tasks. Locks start in
23      the unlocked state.
24  
25      In addition to the methods below, locks can be used in an ``async with``
26      statement.
27      """
28  
29      def __init__(self):
30          # The state can take the following values:
31          # - 0: unlocked
32          # - 1: locked
33          # - <Task>: unlocked but this task has been scheduled to acquire the lock next
34          self.state = 0
35          # Queue of Tasks waiting to acquire this Lock
36          self.waiting = core.TaskQueue()
37  
38      def locked(self):
39          """Returns ``True`` if the lock is locked, otherwise ``False``."""
40  
41          return self.state == 1
42  
43      def release(self):
44          """Release the lock. If any tasks are waiting on the lock then the next
45          one in the queue is scheduled to run and the lock remains locked. Otherwise,
46          no tasks are waiting and the lock becomes unlocked.
47          """
48  
49          if self.state != 1:
50              raise RuntimeError("Lock not acquired")
51          if self.waiting.peek():
52              # Task(s) waiting on lock, schedule next Task
53              self.state = self.waiting.pop_head()
54              core._task_queue.push_head(self.state)
55          else:
56              # No Task waiting so unlock
57              self.state = 0
58  
59      async def acquire(self):
60          """Wait for the lock to be in the unlocked state and then lock it in an
61          atomic way. Only one task can acquire the lock at any one time.
62  
63          This is a coroutine.
64          """
65  
66          if self.state != 0:
67              # Lock unavailable, put the calling Task on the waiting queue
68              self.waiting.push_head(core.cur_task)
69              # Set calling task's data to the lock's queue so it can be removed if needed
70              core.cur_task.data = self.waiting
71              try:
72                  await core.sleep(0)
73              except core.CancelledError as er:
74                  if self.state == core.cur_task:
75                      # Cancelled while pending on resume, schedule next waiting Task
76                      self.state = 1
77                      self.release()
78                  raise er
79          # Lock available, set it as locked
80          self.state = 1
81          return True
82  
83      async def __aenter__(self):
84          return await self.acquire()
85  
86      async def __aexit__(self, exc_type, exc, tb):
87          return self.release()