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>

core/dnsserver/zdirectives.go

@@ -11,6 +11,7 @@ package dnsserver
 // care what plugin above them are doing.
 var Directives = []string{
 	"metadata",
+	"cancel",
 	"tls",
 	"reload",
 	"nsid",

core/plugin/zplugin.go

@@ -8,6 +8,7 @@ import (
 	_ "github.com/coredns/coredns/plugin/autopath"
 	_ "github.com/coredns/coredns/plugin/bind"
 	_ "github.com/coredns/coredns/plugin/cache"
+	_ "github.com/coredns/coredns/plugin/cancel"
 	_ "github.com/coredns/coredns/plugin/chaos"
 	_ "github.com/coredns/coredns/plugin/debug"
 	_ "github.com/coredns/coredns/plugin/deprecated"

plugin.cfg

@@ -20,6 +20,7 @@
 # log:log
 
 metadata:metadata
+cancel:cancel
 tls:tls
 reload:reload
 nsid:nsid

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")
+	}
+}

plugin/done.go

@@ -0,0 +1,14 @@
+package plugin
+
+import "context"
+
+// Done is a non-blocking function that returns true if the context has been canceled.
+func Done(ctx context.Context) bool {
+	select {
+	case <-ctx.Done():
+		return true
+	default:
+		return false
+	}
+	return false
+}