plugin/cancel: add context cancelation plugin (#2711)

* plugin/cancel: add context cancelation plugin

Per review comments on #2704, move this into a plugin that gets called.
Add the most minimal plugin, tests and documenation.

Signed-off-by: Miek Gieben <miek@miek.nl>

* plugin/cache: add timeout option

review feedback: add option to set custom timeout.

Signed-off-by: Miek Gieben <miek@miek.nl>

* spelling

Signed-off-by: Miek Gieben <miek@miek.nl>

plugin/cancel/OWNERS

@@ -0,0 +1,4 @@
+reviewers:
+  - miekg
+approvers:
+  - miekg

plugin/cancel/README.md

@@ -0,0 +1,45 @@
+# cancel
+
+## Name
+
+*cancel* - a plugin that cancels a request's context after 5001 milliseconds.
+
+## Description
+
+The *cancel* plugin creates a canceling context for each request. It adds a timeout that gets
+triggered after 5001 milliseconds.
+
+The 5001 number is chosen because the default timeout for DNS clients is 5 seconds, after that they
+give up.
+
+A plugin interested in the cancellation status should call `plugin.Done()` on the context. If the
+context was canceled due to a timeout the plugin should not write anything back to the client and
+return a value indicating CoreDNS should not either; a zero return value should suffice for that.
+
+~~~ txt
+cancel [TIMEOUT]
+~~~
+
+* **TIMEOUT** allows setting a custom timeout. The default timeout is 5001 milliseconds (`5001 ms`)
+
+## Examples
+
+~~~ corefile
+. {
+    cancel
+    whoami
+}
+~~~
+
+Or with a custom timeout:
+
+~~~ corefile
+. {
+    cancel 1s
+    whoami
+}
+~~~
+
+## Also See
+
+The Go documentation for the context package.

plugin/cancel/cancel.go

@@ -0,0 +1,71 @@
+// Package cancel implements a plugin adds a canceling context to each request.
+package cancel
+
+import (
+	"context"
+	"fmt"
+	"time"
+
+	"github.com/coredns/coredns/core/dnsserver"
+	"github.com/coredns/coredns/plugin"
+
+	"github.com/mholt/caddy"
+	"github.com/miekg/dns"
+)
+
+func init() {
+	caddy.RegisterPlugin("cancel", caddy.Plugin{
+		ServerType: "dns",
+		Action:     setup,
+	})
+}
+
+func setup(c *caddy.Controller) error {
+	ca := Cancel{timeout: 5001 * time.Millisecond}
+
+	for c.Next() {
+		args := c.RemainingArgs()
+		switch len(args) {
+		case 0:
+			break
+		case 1:
+			dur, err := time.ParseDuration(args[0])
+			if err != nil {
+				return plugin.Error("cancel", fmt.Errorf("invalid duration: %q", args[0]))
+			}
+			if dur <= 0 {
+				return plugin.Error("cancel", fmt.Errorf("invalid negative duration: %q", args[0]))
+			}
+			ca.timeout = dur
+		default:
+			return plugin.Error("cancel", c.ArgErr())
+		}
+	}
+
+	dnsserver.GetConfig(c).AddPlugin(func(next plugin.Handler) plugin.Handler {
+		ca.Next = next
+		return ca
+	})
+
+	return nil
+}
+
+// Cancel is a plugin that adds a canceling context to each request's context.
+type Cancel struct {
+	timeout time.Duration
+	Next    plugin.Handler
+}
+
+// ServeDNS implements the plugin.Handler interface.
+func (c Cancel) ServeDNS(ctx context.Context, w dns.ResponseWriter, r *dns.Msg) (int, error) {
+	ctx, cancel := context.WithTimeout(ctx, c.timeout)
+
+	code, err := plugin.NextOrFailure(c.Name(), c.Next, ctx, w, r)
+
+	cancel()
+
+	return code, err
+}
+
+// Name implements the Handler interface.
+func (c Cancel) Name() string { return "cancel" }

plugin/cancel/cancel_test.go

@@ -0,0 +1,52 @@
+package cancel
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	"github.com/coredns/coredns/plugin"
+	"github.com/coredns/coredns/plugin/pkg/dnstest"
+	"github.com/coredns/coredns/plugin/test"
+
+	"github.com/miekg/dns"
+)
+
+type sleepPlugin struct{}
+
+func (s sleepPlugin) Name() string { return "sleep" }
+
+func (s sleepPlugin) ServeDNS(ctx context.Context, w dns.ResponseWriter, r *dns.Msg) (int, error) {
+	i := 0
+	m := new(dns.Msg)
+	m.SetReply(r)
+	for {
+		if plugin.Done(ctx) {
+			m.Rcode = dns.RcodeBadTime // use BadTime to return something time related
+			w.WriteMsg(m)
+			return 0, nil
+		}
+		time.Sleep(20 * time.Millisecond)
+		i++
+		if i > 2 {
+			m.Rcode = dns.RcodeServerFailure
+			w.WriteMsg(m)
+			return 0, nil
+		}
+	}
+	return 0, nil
+}
+
+func TestCancel(t *testing.T) {
+	ca := Cancel{Next: sleepPlugin{}, timeout: 20 * time.Millisecond}
+	ctx := context.Background()
+
+	w := dnstest.NewRecorder(&test.ResponseWriter{})
+	m := new(dns.Msg)
+	m.SetQuestion("aaa.example.com.", dns.TypeTXT)
+
+	ca.ServeDNS(ctx, w, m)
+	if w.Rcode != dns.RcodeBadTime {
+		t.Error("Expected ServeDNS to be canceled by context")
+	}
+}

plugin/cancel/setup_test.go

@@ -0,0 +1,29 @@
+package cancel
+
+import (
+	"testing"
+
+	"github.com/mholt/caddy"
+)
+
+func TestSetup(t *testing.T) {
+	c := caddy.NewTestController("dns", `cancel`)
+	if err := setup(c); err != nil {
+		t.Errorf("Test 1, expected no errors, but got: %q", err)
+	}
+
+	c = caddy.NewTestController("dns", `cancel 5s`)
+	if err := setup(c); err != nil {
+		t.Errorf("Test 2, expected no errors, but got: %q", err)
+	}
+
+	c = caddy.NewTestController("dns", `cancel 5`)
+	if err := setup(c); err == nil {
+		t.Errorf("Test 3, expected errors, but got none")
+	}
+
+	c = caddy.NewTestController("dns", `cancel -1s`)
+	if err := setup(c); err == nil {
+		t.Errorf("Test 4, expected errors, but got none")
+	}
+}