From cd532f29c56679a5734e0a0578811d9a8450b2cb Mon Sep 17 00:00:00 2001
From: Jan Dubsky <jan.dubsky@showmax.com>
Date: Sat, 1 Oct 2022 20:40:04 +0200
Subject: [PATCH] [semaphore] Add TryAcquireAll function

---
 semaphore/semaphore.go | 18 ++++++++++++++++++
 1 file changed, 18 insertions(+)

diff --git a/semaphore/semaphore.go b/semaphore/semaphore.go
index 30f632c..c2aeb04 100644
--- a/semaphore/semaphore.go
+++ b/semaphore/semaphore.go
@@ -94,6 +94,24 @@ func (s *Weighted) TryAcquire(n int64) bool {
 	return success
 }
 
+// TryAcquireAll acquires all free weight from the semaphore without blocking.
+// Returns the weight acquired. Returns 0 and leaves the semaphore unchanged if
+// no weight was available.
+func (s *Weighted) TryAcquireAll() int64 {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
+	// When somebody waits for a token there are obviously not enough tokens
+	// for waiters already waiting. Consequently, there are no free tokens.
+	if s.waiters.Len() > 0 {
+		return 0
+	}
+
+	free := s.size - s.cur
+	s.cur = s.size
+	return free
+}
+
 // Release releases the semaphore with a weight of n.
 func (s *Weighted) Release(n int64) {
 	s.mu.Lock()