From 56c0c17c00f36e44165a07b0bf740403954be379 Mon Sep 17 00:00:00 2001
From: Jille Timmermans <jille@hexon.nl>
Date: Mon, 7 Nov 2022 17:04:04 +0100
Subject: [PATCH] Add WaitContext method

---
 condchan.go      | 20 ++++++++++++++++++++
 condchan_test.go | 17 ++++++++++++++++-
 2 files changed, 36 insertions(+), 1 deletion(-)

diff --git a/condchan.go b/condchan.go
index f50fdbf..45a6281 100644
--- a/condchan.go
+++ b/condchan.go
@@ -2,6 +2,7 @@
 package condchan
 
 import (
+	"context"
 	"sync"
 	"sync/atomic"
 	"unsafe"
@@ -67,6 +68,25 @@ func (cc *CondChan) Wait() {
 	cc.L.Lock()
 }
 
+// WaitContext is like Wait, but also returns when the given context is cancelled.
+// WaitContext() will always reacquire cc.L before returning, even if the context expires.
+func (cc *CondChan) WaitContext(ctx context.Context) error {
+	cc.checker.check()
+
+	cc.chL.RLock()
+	ch := cc.ch
+	cc.chL.RUnlock()
+
+	cc.L.Unlock()
+	defer cc.L.Lock()
+	select {
+	case <-ch:
+		return nil
+	case <-ctx.Done():
+		return ctx.Err()
+	}
+}
+
 // Signal wakes one goroutine waiting on cc, if there is any.
 // It is allowed but not required for the caller to hold cc.L during the call.
 func (cc *CondChan) Signal() {
diff --git a/condchan_test.go b/condchan_test.go
index e30cd86..bb9537c 100644
--- a/condchan_test.go
+++ b/condchan_test.go
@@ -1,12 +1,14 @@
 package condchan_test
 
 import (
-	"gitlab.com/jonas.jasas/condchan"
+	"context"
 	"reflect"
 	"runtime"
 	"sync"
 	"testing"
 	"time"
+
+	"gitlab.com/jonas.jasas/condchan"
 )
 
 func TestSelect(t *testing.T) {
@@ -56,6 +58,19 @@ func TestWaitCancel(t *testing.T) {
 	cc.L.Unlock()
 }
 
+func TestWaitContext(t *testing.T) {
+	cc := condchan.New(&sync.Mutex{})
+
+	ctx, cancel := context.WithTimeout(context.Background(), time.Millisecond*10)
+	defer cancel()
+	cc.L.Lock()
+	err := cc.WaitContext(ctx)
+	cc.L.Unlock()
+	if err == nil {
+		t.Fatal("No error returned, despite the context being cancelled")
+	}
+}
+
 func TestSelectBroadcast(t *testing.T) {
 	cc := condchan.New(&sync.Mutex{})
 
-- 
GitLab