Add AfterFunc to Deadline

Author: noboruma

deadline/deadline.go

@@ -30,12 +30,15 @@ type Deadline struct {
 	deadline time.Time
 	state    deadlineState
 	pending  uint8
+	cbs      map[int]func()
+	nextCbID int
 }
 
 // New creates new deadline timer.
 func New() *Deadline {
 	return &Deadline{
 		done: make(chan struct{}),
+		cbs:  map[int]func(){},
 	}
 }
 
@@ -49,11 +52,40 @@ func (d *Deadline) timeout() {
 
 	d.state = deadlineExceeded
 	done := d.done
+	for _, cb := range d.cbs {
+		go cb()
+	}
+	clear(d.cbs)
 	d.mu.Unlock()
 
 	close(done)
 }
 
+// AfterFunc attaches a function to the deadline.
+// The functions will be triggered on deadline exceeded.
+// If the deadline is reset, the functions are skipped.
+// Attached functions are only triggered once.
+func (d *Deadline) AfterFunc(cb func()) func() bool {
+	d.mu.Lock()
+	defer d.mu.Unlock()
+	d.cbs[d.nextCbID] = cb
+	usedID := d.nextCbID
+	cancel := func() bool {
+		d.mu.Lock()
+		defer d.mu.Unlock()
+		if _, has := d.cbs[usedID]; has {
+			delete(d.cbs, usedID)
+
+			return true
+		}
+
+		return false
+	}
+	d.nextCbID++
+
+	return cancel
+}
+
 // Set new deadline. Zero value means no deadline.
 func (d *Deadline) Set(setTo time.Time) {
 	d.mu.Lock()
@@ -73,6 +105,7 @@ func (d *Deadline) Set(setTo time.Time) {
 	if setTo.IsZero() {
 		d.pending--
 		d.state = deadlineStopped
+		clear(d.cbs)
 
 		return
 	}
@@ -91,6 +124,10 @@ func (d *Deadline) Set(setTo time.Time) {
 	d.pending--
 	d.state = deadlineExceeded
 	close(d.done)
+	for _, cb := range d.cbs {
+		go cb()
+	}
+	clear(d.cbs)
 }
 
 // Done receives deadline signal.

deadline/deadline_test.go

@@ -5,6 +5,7 @@ package deadline
 
 import (
 	"context"
+	"sync/atomic"
 	"testing"
 	"time"
 
@@ -93,6 +94,68 @@ func TestDeadline(t *testing.T) {
 		expectedCalls := []byte{0}
 		assert.Equal(t, expectedCalls, calls, "Wrong order of deadline signal")
 	})
+
+	t.Run("DeadlineAfterFuncExceedFuture", func(t *testing.T) {
+		now := time.Now()
+
+		d := New()
+		var trigger atomic.Int32
+		d.AfterFunc(func() {
+			trigger.Add(1)
+		})
+		d.Set(now.Add(10 * time.Millisecond))
+		<-time.After(20 * time.Millisecond)
+		d.Set(now.Add(10 * time.Millisecond))
+		<-time.After(20 * time.Millisecond)
+
+		assert.Equal(t, int32(1), trigger.Load(), "Function triggered wrong number of time")
+	})
+
+	t.Run("DeadlineAfterFuncExceedPast", func(t *testing.T) {
+		now := time.Now()
+
+		d := New()
+		var trigger atomic.Int32
+		d.AfterFunc(func() {
+			trigger.Add(1)
+		})
+		d.Set(now.Add(10 * time.Millisecond))
+		d.Set(now.Add(-10 * time.Millisecond))
+		<-time.After(20 * time.Millisecond)
+
+		assert.Equal(t, int32(1), trigger.Load(), "Function triggered wrong number of time")
+	})
+
+	t.Run("DeadlineAfterFuncCancel", func(t *testing.T) {
+		now := time.Now()
+
+		d := New()
+		var trigger atomic.Int32
+		d.AfterFunc(func() {
+			trigger.Add(1)
+		})
+
+		d.Set(now.Add(10 * time.Millisecond))
+		d.Set(time.Time{})
+		<-time.After(20 * time.Millisecond)
+
+		assert.Equal(t, int32(0), trigger.Load(), "Function triggered wrong number of time")
+	})
+
+	t.Run("DeadlineAfterFuncDetach", func(t *testing.T) {
+		now := time.Now()
+
+		d := New()
+		var trigger atomic.Int32
+		detach := d.AfterFunc(func() {
+			trigger.Add(1)
+		})
+		detach()
+		d.Set(now.Add(10 * time.Millisecond))
+		<-time.After(20 * time.Millisecond)
+
+		assert.Equal(t, int32(0), trigger.Load(), "Function triggered wrong number of time")
+	})
 }
 
 func sendOnDone(ctx context.Context, done <-chan struct{}, dest chan byte, val byte) {