feat: cond with timeout control

lang/channel/cond.go

@@ -0,0 +1,103 @@
+package channel
+
+import (
+	"context"
+	"sync/atomic"
+	"time"
+)
+
+var (
+	_ Cond = (*cond)(nil)
+)
+
+type CondOption func(c *cond)
+
+func WithCondTimeout(timeout time.Duration) CondOption {
+	return func(c *cond) {
+		c.timeout = timeout
+	}
+}
+
+type Cond interface {
+	Signal() bool
+	Broadcast() bool
+	Wait(ctx context.Context) bool
+}
+
+func NewCond(opts ...CondOption) Cond {
+	cd := new(cond)
+	for _, opt := range opts {
+		opt(cd)
+	}
+	return cd
+}
+
+type condSignal = chan struct{}
+
+type cond struct {
+	signal  atomic.Value
+	timeout time.Duration
+}
+
+func (c *cond) Signal() bool {
+	sv := c.signal.Load()
+	if sv == nil {
+		return false
+	}
+	signal := sv.(condSignal)
+	select {
+	case signal <- struct{}{}:
+		return true
+	default:
+		return false
+	}
+}
+
+func (c *cond) Broadcast() bool {
+BROADCAST:
+	sv := c.signal.Load()
+	if sv == nil {
+		return false
+	}
+	var signal condSignal = nil
+	if !c.signal.CompareAndSwap(sv, signal) {
+		goto BROADCAST
+	}
+	signal = sv.(condSignal)
+	select {
+	case <-signal:
+		return false
+	default:
+		close(signal)
+		return true
+	}
+}
+
+func (c *cond) Wait(ctx context.Context) bool {
+WAIT:
+	sv := c.signal.Load()
+	var signal condSignal
+	if sv == nil {
+		signal = make(condSignal)
+		if !c.signal.CompareAndSwap(nil, signal) {
+			goto WAIT
+		}
+	} else {
+		signal = sv.(condSignal)
+	}
+	if c.timeout > 0 {
+		var cancel context.CancelFunc
+		ctx, cancel = context.WithTimeout(ctx, c.timeout)
+		defer cancel()
+	}
+	if ctx == nil || ctx.Done() == nil {
+		<-signal
+		return true
+	}
+	select {
+	case <-signal:
+		return true
+	case <-ctx.Done():
+		return false
+	}
+}

lang/channel/cond_test.go

@@ -0,0 +1,85 @@
+package channel
+
+import (
+	"context"
+	"runtime"
+	"sync/atomic"
+	"testing"
+	"time"
+)
+
+func TestCond(t *testing.T) {
+	cd := NewCond()
+	var finished int32
+	emptyCtx := context.Background()
+	cancelCtx, cancelFunc := context.WithCancel(emptyCtx)
+	for i := 0; i < 10; i++ {
+		go func(i int) {
+			if i%2 == 0 {
+				cd.Wait(emptyCtx)
+			} else {
+				cd.Wait(cancelCtx)
+			}
+			atomic.AddInt32(&finished, 1)
+		}(i)
+	}
+	time.Sleep(time.Millisecond * 100)
+	cancelFunc()
+	for atomic.LoadInt32(&finished) != int32(5) {
+		runtime.Gosched()
+	}
+	cd.Signal()
+	for atomic.LoadInt32(&finished) != int32(6) {
+		runtime.Gosched()
+	}
+	cd.Signal()
+	for atomic.LoadInt32(&finished) != int32(7) {
+		runtime.Gosched()
+	}
+	cd.Broadcast()
+	cd.Signal()
+	for atomic.LoadInt32(&finished) != int32(10) {
+		runtime.Gosched()
+	}
+}
+
+func TestCondTimeout(t *testing.T) {
+	cd := NewCond(WithCondTimeout(time.Millisecond * 200))
+	go func() {
+		time.Sleep(time.Millisecond * 500)
+		cd.Broadcast()
+	}()
+	begin := time.Now()
+	cd.Wait(context.Background())
+	cost := time.Since(begin)
+	t.Logf("cost=%dms", cost.Milliseconds())
+}
+
+func BenchmarkChanCond(b *testing.B) {
+	b.ReportAllocs()
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		ch := make(chan struct{})
+		go func() {
+			time.Sleep(time.Millisecond)
+			close(ch)
+		}()
+		select {
+		case <-ch:
+		case <-time.After(10 * time.Millisecond):
+		}
+	}
+}
+
+func BenchmarkCond(b *testing.B) {
+	b.ReportAllocs()
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		cd := NewCond(WithCondTimeout(10 * time.Millisecond))
+		go func() {
+			time.Sleep(time.Millisecond)
+			cd.Signal()
+		}()
+		cd.Wait(context.Background())
+	}
+}

lang/channel/singal.go

@@ -0,0 +1,63 @@
+package channel
+
+import (
+	"context"
+	"time"
+)
+
+var (
+	_ Signal = (*signal)(nil)
+)
+
+type Signal interface {
+	Signal()
+	Wait(ctx context.Context) bool
+}
+
+type SignalOption func(c *signal)
+
+func WithSignalTimeout(timeout time.Duration) SignalOption {
+	return func(s *signal) {
+		s.timeout = timeout
+	}
+}
+
+func NewSignal(opts ...SignalOption) Signal {
+	sg := new(signal)
+	for _, opt := range opts {
+		opt(sg)
+	}
+	sg.trigger = make(chan struct{})
+	return sg
+}
+
+type signal struct {
+	trigger chan struct{}
+	timeout time.Duration
+}
+
+func (s *signal) Signal() {
+	select {
+	case <-s.trigger:
+	default:
+		close(s.trigger)
+	}
+}
+
+func (s *signal) Wait(ctx context.Context) bool {
+	if s.timeout > 0 {
+		var cancel context.CancelFunc
+		ctx, cancel = context.WithTimeout(ctx, s.timeout)
+		defer cancel()
+	}
+	if ctx == nil || ctx.Done() == nil {
+		<-s.trigger
+		return true
+	}
+	select {
+	case <-s.trigger:
+		return true
+	case <-ctx.Done():
+		return false
+	}
+}

lang/channel/singal_test.go

@@ -0,0 +1,47 @@
+package channel
+
+import (
+	"context"
+	"runtime"
+	"sync/atomic"
+	"testing"
+	"time"
+)
+
+func TestSignal(t *testing.T) {
+	sg := NewSignal()
+	var finished int32
+	emptyCtx := context.Background()
+	cancelCtx, cancelFunc := context.WithCancel(emptyCtx)
+	for i := 0; i < 10; i++ {
+		go func(i int) {
+			if i%2 == 0 {
+				sg.Wait(emptyCtx)
+			} else {
+				sg.Wait(cancelCtx)
+			}
+			atomic.AddInt32(&finished, 1)
+		}(i)
+	}
+	time.Sleep(time.Millisecond * 100)
+	cancelFunc()
+	for atomic.LoadInt32(&finished) != int32(5) {
+		runtime.Gosched()
+	}
+	sg.Signal()
+	for atomic.LoadInt32(&finished) != int32(10) {
+		runtime.Gosched()
+	}
+}
+
+func TestSignalTimeout(t *testing.T) {
+	sg := NewSignal(WithSignalTimeout(time.Millisecond * 200))
+	go func() {
+		time.Sleep(time.Millisecond * 500)
+		sg.Signal()
+	}()
+	begin := time.Now()
+	sg.Wait(context.Background())
+	cost := time.Since(begin)
+	t.Logf("cost=%dms", cost.Milliseconds())
+}