From c8db042ff0f6f2289539dc0c3cd65499c8fbf388 Mon Sep 17 00:00:00 2001
From: ajatprabha <ajat.prabha.leo@gmail.com>
Date: Sat, 22 Oct 2022 19:33:42 +0530
Subject: [PATCH] expose semaphore lock count

---
 x/rate/semaphore.go      |  4 ++++
 x/rate/semaphore_test.go | 43 ++++++++++++++++++++++++++++++++++++++++
 2 files changed, 47 insertions(+)

diff --git a/x/rate/semaphore.go b/x/rate/semaphore.go
index 5a422a5..32e64b5 100644
--- a/x/rate/semaphore.go
+++ b/x/rate/semaphore.go
@@ -104,6 +104,10 @@ func (s *Semaphore) Release(ctx context.Context) error {
 	return nil
 }
 
+func (s *Semaphore) LockCount(ctx context.Context) (int64, error) {
+	return s.rc.ZCard(ctx, semaphoreKey(s.scope)).Result()
+}
+
 // Close closes the connection to redis.
 func (s *Semaphore) Close() error {
 	return s.rc.Close()
diff --git a/x/rate/semaphore_test.go b/x/rate/semaphore_test.go
index 391cf0f..7929a0b 100644
--- a/x/rate/semaphore_test.go
+++ b/x/rate/semaphore_test.go
@@ -406,3 +406,46 @@ type badConnOpt struct {
 func (b badConnOpt) MakeRedisClient() interface{} {
 	return nil
 }
+
+func TestSemaphore_LockCount(t *testing.T) {
+	desc := "ShouldReturnSetCardinalityAsTheLockCount"
+
+	sema := NewSemaphore(getRedisConnOpt(t), "task-9", 3)
+	defer sema.Close()
+
+	ctx, cancel := asynqcontext.New(&base.TaskMessage{
+		ID:    "task-9-id",
+		Queue: "task-9",
+	}, time.Now().Add(time.Second))
+	defer cancel()
+
+	b, err := sema.Acquire(ctx)
+	if !b {
+		t.Errorf("%s;\nSemaphore.Acquire() failed, want true", desc)
+	}
+	if err != nil {
+		t.Errorf("%s;\nSemaphore.Acquire() got error %v want nil", desc, err)
+	}
+
+	n, err := sema.LockCount(ctx)
+	if err != nil {
+		t.Errorf("%s;\nSemaphore.LockCount() got error %v want nil", desc, err)
+	}
+
+	if n != 1 {
+		t.Errorf("%s;\nSemaphore.LockCount() got %d want 1", desc, n)
+	}
+
+	if err := sema.Release(ctx); err != nil {
+		t.Errorf("%s;\nSemaphore.Release() got error %v", desc, err)
+	}
+
+	n, err = sema.LockCount(ctx)
+	if err != nil {
+		t.Errorf("%s;\nSemaphore.LockCount() got error %v want nil", desc, err)
+	}
+
+	if n != 0 {
+		t.Errorf("%s;\nSemaphore.LockCount() got %d want 0", desc, n)
+	}
+}