Moved expvarhandler from fasthttputil to fasthttp root

2 files changed, 129 insertions, 0 deletions

diff --git a/expvarhandler/expvar.go b/expvarhandler/expvar.go
new file mode 100644
index 0000000..cc812b2
--- /dev/null
+++ b/expvarhandler/expvar.go
@@ -0,0 +1,62 @@
+// Package expvarhandler provides fasthttp-compatible request handler
+// serving expvars.
+package expvarhandler
+
+import (
+	"expvar"
+	"fmt"
+	"regexp"
+
+	"github.com/valyala/fasthttp"
+)
+
+var (
+	expvarHandlerCalls = expvar.NewInt("expvarHandlerCalls")
+	expvarRegexpErrors = expvar.NewInt("expvarRegexpErrors")
+)
+
+// ExpvarHandler dumps json representation of expvars to http response.
+//
+// Expvars may be filtered by regexp provided via 'r' query argument.
+//
+// See https://golang.org/pkg/expvar/ for details.
+func ExpvarHandler(ctx *fasthttp.RequestCtx) {
+	expvarHandlerCalls.Add(1)
+
+	ctx.Response.Reset()
+
+	r, err := getExpvarRegexp(ctx)
+	if err != nil {
+		expvarRegexpErrors.Add(1)
+		fmt.Fprintf(ctx, "Error when obtaining expvar regexp: %s", err)
+		ctx.SetStatusCode(fasthttp.StatusBadRequest)
+		return
+	}
+
+	fmt.Fprintf(ctx, "{\n")
+	first := true
+	expvar.Do(func(kv expvar.KeyValue) {
+		if !first {
+			fmt.Fprintf(ctx, ",\n")
+		}
+		if r.MatchString(kv.Key) {
+			first = false
+			fmt.Fprintf(ctx, "\t%q: %s", kv.Key, kv.Value)
+		}
+	})
+	fmt.Fprintf(ctx, "\n}\n")
+
+	ctx.SetContentType("application/json; charset=utf-8")
+}
+
+func getExpvarRegexp(ctx *fasthttp.RequestCtx) (*regexp.Regexp, error) {
+	r := string(ctx.QueryArgs().Peek("r"))
+	if len(r) == 0 {
+		r = "."
+	}
+	rr, err := regexp.Compile(r)
+	if err != nil {
+		return nil, fmt.Errorf("cannot parse r=%q: %s", r, err)
+	}
+	return rr, nil
+}
diff --git a/expvarhandler/expvar_test.go b/expvarhandler/expvar_test.go
new file mode 100644
index 0000000..df27dd7
--- /dev/null
+++ b/expvarhandler/expvar_test.go
@@ -0,0 +1,67 @@
+package expvarhandler
+
+import (
+	"encoding/json"
+	"expvar"
+	"strings"
+	"testing"
+
+	"github.com/valyala/fasthttp"
+)
+
+func TestExpvarHandlerBasic(t *testing.T) {
+	expvar.Publish("customVar", expvar.Func(func() interface{} {
+		return "foobar"
+	}))
+
+	var ctx fasthttp.RequestCtx
+
+	expvarHandlerCalls.Set(0)
+
+	ExpvarHandler(&ctx)
+
+	body := ctx.Response.Body()
+
+	var m map[string]interface{}
+	if err := json.Unmarshal(body, &m); err != nil {
+		t.Fatalf("unexpected error: %s", err)
+	}
+
+	if _, ok := m["cmdline"]; !ok {
+		t.Fatalf("cannot locate cmdline expvar")
+	}
+	if _, ok := m["memstats"]; !ok {
+		t.Fatalf("cannot locate memstats expvar")
+	}
+
+	v := m["customVar"]
+	sv, ok := v.(string)
+	if !ok {
+		t.Fatalf("unexpected custom var type %T. Expecting string", v)
+	}
+	if sv != "foobar" {
+		t.Fatalf("unexpected custom var value: %q. Expecting %q", v, "foobar")
+	}
+
+	v = m["expvarHandlerCalls"]
+	fv, ok := v.(float64)
+	if !ok {
+		t.Fatalf("unexpected expvarHandlerCalls type %T. Expecting float64", v)
+	}
+	if int(fv) != 1 {
+		t.Fatalf("unexpected value for expvarHandlerCalls: %v. Expecting %v", fv, 1)
+	}
+}
+
+func TestExpvarHandlerRegexp(t *testing.T) {
+	var ctx fasthttp.RequestCtx
+	ctx.QueryArgs().Set("r", "cmd")
+	ExpvarHandler(&ctx)
+	body := string(ctx.Response.Body())
+	if !strings.Contains(body, `"cmdline"`) {
+		t.Fatalf("missing 'cmdline' expvar")
+	}
+	if strings.Contains(body, `"memstats"`) {
+		t.Fatalf("unexpected memstats expvar found")
+	}
+}