Repository: python-trio/trio (Python)

Issue:
Allow `CapacityLimiter` to have `total_tokens=0`
Sometimes, there are use-cases where you want to entirely pause processing of some tasks, and then later scale back up.  `CapacityLimiter` is basically perfect for this kind of thing, except that it currently raises an error if you pass `total_tokens=0`:

https://github.com/python-trio/trio/blob/d988edd368f094855bd1c36442302113f45f2a1c/src/trio/_sync.py#L254-L255

It's already possible to get the desired effect just by deliberately "leaking" one of the tokens, but it would be much nicer to simply allow zero as the intended capacity, by editing the numbers in the setter above.


Discussion:
Might need to make sure this is tested too, since zero is rather susceptible to edge cases. Looks safe though.
@Zac-HD , I'm interested, could you kindly assign me this task
Hi @kachida any progress? If not, would you mind un-assigning yourself?

You should be able to re-assign yourself whenever you have the time and want to fix this issue.



Patch:
diff --git a/newsfragments/3321.feature.rst b/newsfragments/3321.feature.rst
new file mode 100644
index 0000000000..e62649bbd2
--- /dev/null
+++ b/newsfragments/3321.feature.rst
@@ -0,0 +1,1 @@
+Allow `trio.CapacityLimiter` to have zero total_tokens.
diff --git a/src/trio/_sync.py b/src/trio/_sync.py
index d026f4bc37..dbdee07984 100644
--- a/src/trio/_sync.py
+++ b/src/trio/_sync.py
@@ -283,8 +283,8 @@ def total_tokens(self) -> int | float:
     def total_tokens(self, new_total_tokens: int | float) -> None:  # noqa: PYI041
         if not isinstance(new_total_tokens, int) and new_total_tokens != math.inf:
             raise TypeError("total_tokens must be an int or math.inf")
-        if new_total_tokens < 1:
-            raise ValueError("total_tokens must be >= 1")
+        if new_total_tokens < 0:
+            raise ValueError("total_tokens must be >= 0")
         self._total_tokens = new_total_tokens
         self._wake_waiters()
 


Test patch:
diff --git a/src/trio/_tests/test_sync.py b/src/trio/_tests/test_sync.py
index 6096510d3f..836e8597e8 100644
--- a/src/trio/_tests/test_sync.py
+++ b/src/trio/_tests/test_sync.py
@@ -55,9 +55,10 @@ async def child() -> None:
 
 
 async def test_CapacityLimiter() -> None:
+    assert CapacityLimiter(0).total_tokens == 0
     with pytest.raises(TypeError):
         CapacityLimiter(1.0)
-    with pytest.raises(ValueError, match=r"^total_tokens must be >= 1$"):
+    with pytest.raises(ValueError, match=r"^total_tokens must be >= 0$"):
         CapacityLimiter(-1)
     c = CapacityLimiter(2)
     repr(c)  # smoke test
@@ -145,10 +146,10 @@ async def test_CapacityLimiter_change_total_tokens() -> None:
     with pytest.raises(TypeError):
         c.total_tokens = 1.0
 
-    with pytest.raises(ValueError, match=r"^total_tokens must be >= 1$"):
-        c.total_tokens = 0
+    with pytest.raises(ValueError, match=r"^total_tokens must be >= 0$"):
+        c.total_tokens = -1
 
-    with pytest.raises(ValueError, match=r"^total_tokens must be >= 1$"):
+    with pytest.raises(ValueError, match=r"^total_tokens must be >= 0$"):
         c.total_tokens = -10
 
     assert c.total_tokens == 2
@@ -190,6 +191,83 @@ async def test_CapacityLimiter_memleak_548() -> None:
     assert len(limiter._pending_borrowers) == 0
 
 
+async def test_CapacityLimiter_zero_limit_tokens() -> None:
+    c = CapacityLimiter(5)
+
+    assert c.total_tokens == 5
+
+    async with _core.open_nursery() as nursery:
+        c.total_tokens = 0
+
+        for i in range(5):
+            nursery.start_soon(c.acquire_on_behalf_of, i)
+            await wait_all_tasks_blocked()
+
+        assert set(c.statistics().borrowers) == set()
+        assert c.statistics().tasks_waiting == 5
+
+        c.total_tokens = 5
+
+        assert set(c.statistics().borrowers) == {0, 1, 2, 3, 4}
+
+        nursery.start_soon(c.acquire_on_behalf_of, 5)
+        await wait_all_tasks_blocked()
+
+        assert c.statistics().tasks_waiting == 1
+
+        for i in range(5):
+            c.release_on_behalf_of(i)
+
+        assert c.statistics().tasks_waiting == 0
+        c.release_on_behalf_of(5)
+
+        # making sure that zero limit capacity limiter doesn't let any tasks through
+
+        c.total_tokens = 0
+
+        with pytest.raises(_core.WouldBlock):
+            c.acquire_nowait()
+
+        nursery.start_soon(c.acquire_on_behalf_of, 6)
+        await wait_all_tasks_blocked()
+
+        assert c.statistics().tasks_waiting == 1
+        assert c.statistics().borrowers == []
+
+        c.total_tokens = 1
+        assert c.statistics().tasks_waiting == 0
+        assert c.statistics().borrowers == [6]
+        c.release_on_behalf_of(6)
+
+        await c.acquire_on_behalf_of(0)  # total_tokens is 1
+
+        nursery.start_soon(c.acquire_on_behalf_of, 1)
+        await wait_all_tasks_blocked()
+        c.total_tokens = 0
+
+        assert c.statistics().borrowers == [0]
+
+        c.release_on_behalf_of(0)
+        await wait_all_tasks_blocked()
+        assert c.statistics().borrowers == []
+        assert c.statistics().tasks_waiting == 1
+
+        c.total_tokens = 1
+        await wait_all_tasks_blocked()
+        assert c.statistics().borrowers == [1]
+        assert c.statistics().tasks_waiting == 0
+
+        c.release_on_behalf_of(1)
+
+        c.total_tokens = 0
+
+        nursery.cancel_scope.cancel()
+
+    assert c.total_tokens == 0
+    assert c.statistics().borrowers == []
+    assert c._pending_borrowers == {}
+
+
 async def test_Semaphore() -> None:
     with pytest.raises(TypeError):
         Semaphore(1.0)  # type: ignore[arg-type]