From 2f23bc271f2075de740c6507619caa53303eacc1 Mon Sep 17 00:00:00 2001
From: Edward <crazybber@outlook.com>
Date: Sun, 3 May 2020 10:53:23 +0800
Subject: [PATCH] add timeout pattern

---
 gomore/deadline/deadline.go      | 49 ++++++++++++++++++++++++
 gomore/deadline/deadline_test.go | 65 ++++++++++++++++++++++++++++++++
 2 files changed, 114 insertions(+)
 create mode 100644 gomore/deadline/deadline.go
 create mode 100644 gomore/deadline/deadline_test.go

diff --git a/gomore/deadline/deadline.go b/gomore/deadline/deadline.go
new file mode 100644
index 0000000..c97fd86
--- /dev/null
+++ b/gomore/deadline/deadline.go
@@ -0,0 +1,49 @@
+// Package deadline implements the deadline (also known as "timeout") resiliency pattern for Go.
+package deadline
+
+import (
+	"errors"
+	"time"
+)
+
+// ErrTimedOut is the error returned from Run when the deadline expires.
+var ErrTimedOut = errors.New("timed out waiting for function to finish")
+
+// Deadline implements the deadline/timeout resiliency pattern.
+type Deadline struct {
+	timeout time.Duration
+}
+
+// New constructs a new Deadline with the given timeout.
+func New(timeout time.Duration) *Deadline {
+	return &Deadline{
+		timeout: timeout,
+	}
+}
+
+// Run runs the given function, passing it a stopper channel. If the deadline passes before
+// the function finishes executing, Run returns ErrTimeOut to the caller and closes the stopper
+// channel so that the work function can attempt to exit gracefully. It does not (and cannot)
+// simply kill the running function, so if it doesn't respect the stopper channel then it may
+// keep running after the deadline passes. If the function finishes before the deadline, then
+// the return value of the function is returned from Run.
+func (d *Deadline) Run(work func(<-chan struct{}) error) error {
+	result := make(chan error)
+	stopper := make(chan struct{})
+
+	go func() {
+		value := work(stopper)
+		select {
+		case result <- value:
+		case <-stopper:
+		}
+	}()
+
+	select {
+	case ret := <-result:
+		return ret
+	case <-time.After(d.timeout):
+		close(stopper)
+		return ErrTimedOut
+	}
+}
diff --git a/gomore/deadline/deadline_test.go b/gomore/deadline/deadline_test.go
new file mode 100644
index 0000000..6939f52
--- /dev/null
+++ b/gomore/deadline/deadline_test.go
@@ -0,0 +1,65 @@
+package deadline
+
+import (
+	"errors"
+	"testing"
+	"time"
+)
+
+func takesFiveMillis(stopper <-chan struct{}) error {
+	time.Sleep(5 * time.Millisecond)
+	return nil
+}
+
+func takesTwentyMillis(stopper <-chan struct{}) error {
+	time.Sleep(20 * time.Millisecond)
+	return nil
+}
+
+func returnsError(stopper <-chan struct{}) error {
+	return errors.New("foo")
+}
+
+func TestDeadline(t *testing.T) {
+	dl := New(10 * time.Millisecond)
+
+	if err := dl.Run(takesFiveMillis); err != nil {
+		t.Error(err)
+	}
+
+	if err := dl.Run(takesTwentyMillis); err != ErrTimedOut {
+		t.Error(err)
+	}
+
+	if err := dl.Run(returnsError); err.Error() != "foo" {
+		t.Error(err)
+	}
+
+	done := make(chan struct{})
+	err := dl.Run(func(stopper <-chan struct{}) error {
+		<-stopper
+		close(done)
+		return nil
+	})
+	if err != ErrTimedOut {
+		t.Error(err)
+	}
+	<-done
+}
+
+func ExampleDeadline() {
+	dl := New(1 * time.Second)
+
+	err := dl.Run(func(stopper <-chan struct{}) error {
+		// do something possibly slow
+		// check stopper function and give up if timed out
+		return nil
+	})
+
+	switch err {
+	case ErrTimedOut:
+		// execution took too long, oops
+	default:
+		// some other error
+	}
+}