From 212e074be2bf8be3639250050ce9867d690f7b2c Mon Sep 17 00:00:00 2001
From: Roger Peppe <rogpeppe@gmail.com>
Date: Mon, 10 Jan 2011 11:51:38 -0800
Subject: [PATCH] time: add AfterFunc to call a function after a given
 duration. The After code is trivially generalisable to provide support for
 this, and it is possible to use AfterFunc to do things that After cannot,
 such as waiting for many events at varied times without an overhead of one
 goroutine per event.

R=rsc, r
CC=golang-dev
https://golang.org/cl/3905041
---
 src/pkg/time/sleep.go      | 31 +++++++++++++++++++++++--------
 src/pkg/time/sleep_test.go | 38 ++++++++++++++++++++++++++++++++++++++
 2 files changed, 61 insertions(+), 8 deletions(-)

diff --git a/src/pkg/time/sleep.go b/src/pkg/time/sleep.go
index 77b7b4a5935..3538775adfb 100644
--- a/src/pkg/time/sleep.go
+++ b/src/pkg/time/sleep.go
@@ -11,11 +11,11 @@ import (
 	"container/heap"
 )
 
-// The event type represents a single After event.
+// The event type represents a single After or AfterFunc event.
 type event struct {
-	t        int64        // The absolute time that the event should fire.
-	c        chan<- int64 // The channel to send on.
-	sleeping bool         // A sleeper is sleeping for this event.
+	t        int64       // The absolute time that the event should fire.
+	f        func(int64) // The function to call when the event fires.
+	sleeping bool        // A sleeper is sleeping for this event.
 }
 
 type eventHeap []*event
@@ -55,15 +55,30 @@ func sleep(t, ns int64) (int64, os.Error) {
 // on the returned channel.
 func After(ns int64) <-chan int64 {
 	c := make(chan int64, 1)
-	t := ns + Nanoseconds()
+	after(ns, func(t int64) { c <- t })
+	return c
+}
+
+// AfterFunc waits at least ns nanoseconds before calling f
+// in its own goroutine.
+func AfterFunc(ns int64, f func()) {
+	after(ns, func(_ int64) {
+		go f()
+	})
+}
+
+// after is the implementation of After and AfterFunc.
+// When the current time is after ns, it calls f with the current time.
+// It assumes that f will not block.
+func after(ns int64, f func(int64)) {
+	t := Nanoseconds() + ns
 	eventMutex.Lock()
 	t0 := events[0].t
-	heap.Push(events, &event{t, c, false})
+	heap.Push(events, &event{t, f, false})
 	if t < t0 {
 		go sleeper()
 	}
 	eventMutex.Unlock()
-	return c
 }
 
 // sleeper continually looks at the earliest event in the queue, marks it
@@ -102,7 +117,7 @@ func sleeper() {
 			e = events[0]
 		}
 		for t >= e.t {
-			e.c <- t
+			e.f(t)
 			heap.Pop(events)
 			e = events[0]
 		}
diff --git a/src/pkg/time/sleep_test.go b/src/pkg/time/sleep_test.go
index e70b84e97c1..9e36288f886 100644
--- a/src/pkg/time/sleep_test.go
+++ b/src/pkg/time/sleep_test.go
@@ -26,6 +26,44 @@ func TestSleep(t *testing.T) {
 	}
 }
 
+// Test the basic function calling behavior. Correct queueing
+// behavior is tested elsewhere, since After and AfterFunc share
+// the same code.
+func TestAfterFunc(t *testing.T) {
+	i := 10
+	c := make(chan bool)
+	var f func()
+	f = func() {
+		i--
+		if i >= 0 {
+			AfterFunc(0, f)
+			Sleep(1e9)
+		} else {
+			c <- true
+		}
+	}
+
+	AfterFunc(0, f)
+	<-c
+}
+
+func BenchmarkAfterFunc(b *testing.B) {
+	i := b.N
+	c := make(chan bool)
+	var f func()
+	f = func() {
+		i--
+		if i >= 0 {
+			AfterFunc(0, f)
+		} else {
+			c <- true
+		}
+	}
+
+	AfterFunc(0, f)
+	<-c
+}
+
 func TestAfter(t *testing.T) {
 	const delay = int64(100e6)
 	start := Nanoseconds()