Reuse timer in Deadline

deadline/deadline.go

@@ -7,98 +7,104 @@ package deadline
 
 import (
 	"context"
+	"math"
 	"sync"
 	"time"
 )
 
+type deadlineState uint8
+
+const (
+	deadlineStopped deadlineState = iota
+	deadlineStarted
+	deadlineExceeded
+)
+
+var _ context.Context = (*Deadline)(nil)
+
 // Deadline signals updatable deadline timer.
 // Also, it implements context.Context.
 type Deadline struct {
-	exceeded chan struct{}
-	stop     chan struct{}
-	stopped  chan bool
-	deadline time.Time
 	mu       sync.RWMutex
+	timer    timer
+	done     chan struct{}
+	deadline time.Time
+	state    deadlineState
+	pending  uint8
 }
 
 // New creates new deadline timer.
 func New() *Deadline {
 	d := &Deadline{
-		exceeded: make(chan struct{}),
-		stop:     make(chan struct{}),
-		stopped:  make(chan bool, 1),
+		done: make(chan struct{}),
 	}
-	d.stopped <- true
+	d.timer = afterFunc(math.MaxInt64, d.timeout)
+	d.timer.Stop()
 	return d
 }
 
+func (d *Deadline) timeout() {
+	d.mu.Lock()
+	if d.pending--; d.pending != 0 || d.state != deadlineStarted {
+		d.mu.Unlock()
+		return
+	}
+
+	d.state = deadlineExceeded
+	d.mu.Unlock()
+
+	close(d.done)
+}
+
 // Set new deadline. Zero value means no deadline.
 func (d *Deadline) Set(t time.Time) {
 	d.mu.Lock()
 	defer d.mu.Unlock()
 
-	d.deadline = t
+	if d.state == deadlineStarted && d.timer.Stop() {
+		d.pending--
+	}
 
-	close(d.stop)
+	d.deadline = t
+	d.pending++
 
-	select {
-	case <-d.exceeded:
-		d.exceeded = make(chan struct{})
-	default:
-		stopped := <-d.stopped
-		if !stopped {
-			d.exceeded = make(chan struct{})
-		}
+	if d.state == deadlineExceeded {
+		d.done = make(chan struct{})
 	}
-	d.stop = make(chan struct{})
-	d.stopped = make(chan bool, 1)
 
 	if t.IsZero() {
-		d.stopped <- true
+		d.pending--
+		d.state = deadlineStopped
 		return
 	}
 
 	if dur := time.Until(t); dur > 0 {
-		exceeded := d.exceeded
-		stopped := d.stopped
-		go func() {
-			timer := time.NewTimer(dur)
-			select {
-			case <-timer.C:
-				close(exceeded)
-				stopped <- false
-			case <-d.stop:
-				if !timer.Stop() {
-					<-timer.C
-				}
-				stopped <- true
-			}
-		}()
+		d.state = deadlineStarted
+		d.timer.Reset(dur)
 		return
 	}
 
-	close(d.exceeded)
-	d.stopped <- false
+	d.pending--
+	d.state = deadlineExceeded
+	close(d.done)
 }
 
 // Done receives deadline signal.
 func (d *Deadline) Done() <-chan struct{} {
 	d.mu.RLock()
 	defer d.mu.RUnlock()
-	return d.exceeded
+	return d.done
 }
 
 // Err returns context.DeadlineExceeded if the deadline is exceeded.
 // Otherwise, it returns nil.
 func (d *Deadline) Err() error {
 	d.mu.RLock()
 	defer d.mu.RUnlock()
-	select {
-	case <-d.exceeded:
+	if d.state == deadlineExceeded {
 		return context.DeadlineExceeded
-	default:
-		return nil
 	}
+	return nil
 }
 
 // Deadline returns current deadline.

deadline/deadline_test.go

@@ -170,3 +170,13 @@ func TestContext(t *testing.T) {
 		}
 	})
 }
+
+func BenchmarkDeadline(b *testing.B) {
+	b.Run("Set", func(b *testing.B) {
+		d := New()
+		t := time.Now().Add(time.Minute)
+		for i := 0; i < b.N; i++ {
+			d.Set(t)
+		}
+	})
+}

deadline/timer.go

@@ -0,0 +1,13 @@
+// SPDX-FileCopyrightText: 2023 The Pion community <https://pion.ly>
+// SPDX-License-Identifier: MIT
+
+package deadline
+
+import (
+	"time"
+)
+
+type timer interface {
+	Stop() bool
+	Reset(time.Duration) bool
+}

deadline/timerfunc_generic.go

@@ -0,0 +1,13 @@
+// SPDX-FileCopyrightText: 2023 The Pion community <https://pion.ly>
+// SPDX-License-Identifier: MIT
+
+//go:build !js
+// +build !js
+
+package deadline
+
+import "time"
+
+func afterFunc(d time.Duration, f func()) timer {
+	return time.AfterFunc(d, f)
+}

deadline/timerfunc_js.go

@@ -0,0 +1,67 @@
+// SPDX-FileCopyrightText: 2023 The Pion community <https://pion.ly>
+// SPDX-License-Identifier: MIT
+
+//go:build js
+// +build js
+
+package deadline
+
+import (
+	"sync"
+	"time"
+)
+
+// jsTimer is a timer utility for wasm with a working Reset function.
+type jsTimer struct {
+	f       func()
+	mu      sync.Mutex
+	timer   *time.Timer
+	gen     uint64
+	started bool
+}
+
+func afterFunc(d time.Duration, f func()) timer {
+	t := &jsTimer{f: f}
+	t.Reset(d)
+	return t
+}
+
+func (t *jsTimer) Stop() bool {
+	t.mu.Lock()
+	defer t.mu.Unlock()
+
+	t.gen++
+	t.timer.Stop()
+
+	started := t.started
+	t.started = false
+	return started
+}
+
+func (t *jsTimer) Reset(d time.Duration) bool {
+	t.mu.Lock()
+	defer t.mu.Unlock()
+
+	if t.timer != nil {
+		t.timer.Stop()
+	}
+
+	t.gen++
+	gen := t.gen
+	t.timer = time.AfterFunc(d, func() {
+		t.mu.Lock()
+		if gen != t.gen {
+			t.mu.Unlock()
+			return
+		}
+
+		t.started = false
+		t.mu.Unlock()
+
+		t.f()
+	})
+
+	started := t.started
+	t.started = true
+	return started
+}