Fix main startup (#232)

Set version and name of the program. And then call coremain.Run().

The coremain split makes CoreDNS embeddable.

Also see #189 for an old PR.

README.md

@@ -51,10 +51,8 @@ Caddyfile when I forked it).
 
 ## Compilation
 
-CoreDNS (as a servertype plugin for Caddy) has a dependency on Caddy - this is *almost* like
-the normal Go dependencies, but with a small twist: caddy (the source) need to know that CoreDNS
-exists and for this we need to add 1 line `_ "github.com/miekg/coredns/core"` to file in caddy.
-
+CoreDNS (as a servertype plugin for Caddy) has a dependency on Caddy, but this is not different than
+any other Go dependency.
 You have the source of CoreDNS, this should preferably be downloaded under your `$GOPATH`. Get all
 dependencies:
 

coredns.go

 package main
 
-import (
-	"flag"
-
-	"github.com/mholt/caddy"
-	"github.com/mholt/caddy/caddy/caddymain"
-)
-
-//go:generate go run plugin_generate.go
+import "github.com/miekg/coredns/coremain"
 
 func main() {
-	setFlag()
-	setName()
-
-	caddymain.Run()
-}
-
-// setFlag sets flags to predefined values for CoreDNS.
-func setFlag() {
-	flag.Set("type", "dns")
-}
-
-// setName sets application name and versioning information for CoreDNS.
-func setName() {
-	caddy.DefaultConfigFile = "Corefile"
-	caddy.AppName = "CoreDNS"
-	caddy.AppVersion = version
+	coremain.Run()
 }
-
-const version = "001"

coremain/run.go

+package coremain
+
+import (
+	"errors"
+	"flag"
+	"fmt"
+	"io/ioutil"
+	"log"
+	"os"
+	"runtime"
+	"strconv"
+	"strings"
+
+	"github.com/mholt/caddy"
+	"gopkg.in/natefinch/lumberjack.v2"
+
+	// Plug in CoreDNS
+	_ "github.com/miekg/coredns/core"
+)
+
+func init() {
+	caddy.TrapSignals()
+	setVersion()
+
+	flag.StringVar(&conf, "conf", "", "Caddyfile to load (default \""+caddy.DefaultConfigFile+"\")")
+	flag.StringVar(&cpu, "cpu", "100%", "CPU cap")
+	flag.BoolVar(&plugins, "plugins", false, "List installed plugins")
+	flag.StringVar(&logfile, "log", "", "Process log file")
+	flag.StringVar(&caddy.PidFile, "pidfile", "", "Path to write pid file")
+	flag.BoolVar(&caddy.Quiet, "quiet", false, "Quiet mode (no initialization output)")
+	flag.BoolVar(&version, "version", false, "Show version")
+
+	caddy.RegisterCaddyfileLoader("flag", caddy.LoaderFunc(confLoader))
+	caddy.SetDefaultCaddyfileLoader("default", caddy.LoaderFunc(defaultLoader))
+}
+
+// Run is CoreDNS's main() function.
+func Run() {
+	flag.Parse()
+
+	caddy.AppName = coreName
+	caddy.AppVersion = coreVersion
+
+	// Set up process log before anything bad happens
+	switch logfile {
+	case "stdout":
+		log.SetOutput(os.Stdout)
+	case "stderr":
+		log.SetOutput(os.Stderr)
+	case "":
+		log.SetOutput(ioutil.Discard)
+	default:
+		log.SetOutput(&lumberjack.Logger{
+			Filename:   logfile,
+			MaxSize:    100,
+			MaxAge:     14,
+			MaxBackups: 10,
+		})
+	}
+
+	if version {
+		fmt.Printf("%s-%s\n", caddy.AppName, caddy.AppVersion)
+		if devBuild && gitShortStat != "" {
+			fmt.Printf("%s\n%s\n", gitShortStat, gitFilesModified)
+		}
+		os.Exit(0)
+	}
+	if plugins {
+		fmt.Println(caddy.DescribePlugins())
+		os.Exit(0)
+	}
+
+	// Set CPU cap
+	err := setCPU(cpu)
+	if err != nil {
+		mustLogFatal(err)
+	}
+
+	// Get Caddyfile input
+	caddyfile, err := caddy.LoadCaddyfile(serverType)
+	if err != nil {
+		mustLogFatal(err)
+	}
+
+	// Start your engines
+	instance, err := caddy.Start(caddyfile)
+	if err != nil {
+		mustLogFatal(err)
+	}
+
+	// Twiddle your thumbs
+	instance.Wait()
+}
+
+// mustLogFatal wraps log.Fatal() in a way that ensures the
+// output is always printed to stderr so the user can see it
+// if the user is still there, even if the process log was not
+// enabled. If this process is an upgrade, however, and the user
+// might not be there anymore, this just logs to the process
+// log and exits.
+func mustLogFatal(args ...interface{}) {
+	if !caddy.IsUpgrade() {
+		log.SetOutput(os.Stderr)
+	}
+	log.Fatal(args...)
+}
+
+// confLoader loads the Caddyfile using the -conf flag.
+func confLoader(serverType string) (caddy.Input, error) {
+	if conf == "" {
+		return nil, nil
+	}
+
+	if conf == "stdin" {
+		return caddy.CaddyfileFromPipe(os.Stdin)
+	}
+
+	contents, err := ioutil.ReadFile(conf)
+	if err != nil {
+		return nil, err
+	}
+	return caddy.CaddyfileInput{
+		Contents:       contents,
+		Filepath:       conf,
+		ServerTypeName: serverType,
+	}, nil
+}
+
+// defaultLoader loads the Caddyfile from the current working directory.
+func defaultLoader(serverType string) (caddy.Input, error) {
+	contents, err := ioutil.ReadFile(caddy.DefaultConfigFile)
+	if err != nil {
+		if os.IsNotExist(err) {
+			return nil, nil
+		}
+		return nil, err
+	}
+	return caddy.CaddyfileInput{
+		Contents:       contents,
+		Filepath:       caddy.DefaultConfigFile,
+		ServerTypeName: serverType,
+	}, nil
+}
+
+// setVersion figures out the version information
+// based on variables set by -ldflags.
+func setVersion() {
+	// A development build is one that's not at a tag or has uncommitted changes
+	devBuild = gitTag == "" || gitShortStat != ""
+
+	// Only set the appVersion if -ldflags was used
+	if gitNearestTag != "" || gitTag != "" {
+		if devBuild && gitNearestTag != "" {
+			appVersion = fmt.Sprintf("%s (+%s %s)",
+				strings.TrimPrefix(gitNearestTag, "v"), gitCommit, buildDate)
+		} else if gitTag != "" {
+			appVersion = strings.TrimPrefix(gitTag, "v")
+		}
+	}
+}
+
+// setCPU parses string cpu and sets GOMAXPROCS
+// according to its value. It accepts either
+// a number (e.g. 3) or a percent (e.g. 50%).
+func setCPU(cpu string) error {
+	var numCPU int
+
+	availCPU := runtime.NumCPU()
+
+	if strings.HasSuffix(cpu, "%") {
+		// Percent
+		var percent float32
+		pctStr := cpu[:len(cpu)-1]
+		pctInt, err := strconv.Atoi(pctStr)
+		if err != nil || pctInt < 1 || pctInt > 100 {
+			return errors.New("invalid CPU value: percentage must be between 1-100")
+		}
+		percent = float32(pctInt) / 100
+		numCPU = int(float32(availCPU) * percent)
+	} else {
+		// Number
+		num, err := strconv.Atoi(cpu)
+		if err != nil || num < 1 {
+			return errors.New("invalid CPU value: provide a number or percent greater than 0")
+		}
+		numCPU = num
+	}
+
+	if numCPU > availCPU {
+		numCPU = availCPU
+	}
+
+	runtime.GOMAXPROCS(numCPU)
+	return nil
+}
+
+// Flags that control program flow or startup
+var (
+	conf    string
+	cpu     string
+	logfile string
+	version bool
+	plugins bool
+)
+
+// Build information obtained with the help of -ldflags
+var (
+	appVersion = "(untracked dev build)" // inferred at startup
+	devBuild   = true                    // inferred at startup
+
+	buildDate        string // date -u
+	gitTag           string // git describe --exact-match HEAD 2> /dev/null
+	gitNearestTag    string // git describe --abbrev=0 --tags HEAD
+	gitCommit        string // git rev-parse HEAD
+	gitShortStat     string // git diff-index --shortstat
+	gitFilesModified string // git diff-index --name-only HEAD
+)
+
+const (
+	coreName    = "CoreDNS"
+	coreVersion = "001"
+	serverType  = "dns"
+)

coremain/run_test.go

+package coremain
+
+import (
+	"runtime"
+	"testing"
+)
+
+func TestSetCPU(t *testing.T) {
+	currentCPU := runtime.GOMAXPROCS(-1)
+	maxCPU := runtime.NumCPU()
+	halfCPU := int(0.5 * float32(maxCPU))
+	if halfCPU < 1 {
+		halfCPU = 1
+	}
+	for i, test := range []struct {
+		input     string
+		output    int
+		shouldErr bool
+	}{
+		{"1", 1, false},
+		{"-1", currentCPU, true},
+		{"0", currentCPU, true},
+		{"100%", maxCPU, false},
+		{"50%", halfCPU, false},
+		{"110%", currentCPU, true},
+		{"-10%", currentCPU, true},
+		{"invalid input", currentCPU, true},
+		{"invalid input%", currentCPU, true},
+		{"9999", maxCPU, false}, // over available CPU
+	} {
+		err := setCPU(test.input)
+		if test.shouldErr && err == nil {
+			t.Errorf("Test %d: Expected error, but there wasn't any", i)
+		}
+		if !test.shouldErr && err != nil {
+			t.Errorf("Test %d: Expected no error, but there was one: %v", i, err)
+		}
+		if actual, expected := runtime.GOMAXPROCS(-1), test.output; actual != expected {
+			t.Errorf("Test %d: GOMAXPROCS was %d but expected %d", i, actual, expected)
+		}
+		// teardown
+		runtime.GOMAXPROCS(currentCPU)
+	}
+}

middleware/chaos/setup.go

@@ -47,4 +47,4 @@ func chaosParse(c *caddy.Controller) (string, map[string]bool, error) {
 	return version, authors, nil
 }
 
-const defaultVersion = "CoreDNS"
+var defaultVersion = caddy.AppName + "-" + caddy.AppVersion

plugin_generate.go

-//+build ignore
-
-package main
-
-import (
-	"bytes"
-	"go/ast"
-	"go/parser"
-	"go/printer"
-	"go/token"
-	"io/ioutil"
-	"log"
-
-	"golang.org/x/tools/go/ast/astutil"
-)
-
-func GenerateFile(fset *token.FileSet, file *ast.File) ([]byte, error) {
-	var output []byte
-	buffer := bytes.NewBuffer(output)
-	if err := printer.Fprint(buffer, fset, file); err != nil {
-		return nil, err
-	}
-
-	return buffer.Bytes(), nil
-}
-
-func main() {
-	fset := token.NewFileSet()
-	f, err := parser.ParseFile(fset, caddyrun, nil, parser.ParseComments)
-	if err != nil {
-		log.Fatalf("failed to parse %s: %s", caddyrun, err)
-	}
-	astutil.AddNamedImport(fset, f, "_", coredns)
-	astutil.DeleteNamedImport(fset, f, "_", caddy)
-
-	out, err := GenerateFile(fset, f)
-	if err := ioutil.WriteFile(caddyrun, out, 0644); err != nil {
-		log.Fatalf("failed to write go file: %s", err)
-	}
-}
-
-const (
-	coredns = "github.com/miekg/coredns/core"
-	caddy   = "github.com/mholt/caddy/caddyhttp"
-
-	// If everything is OK and we are sitting in CoreDNS' dir, this is where run.go should be.
-	caddyrun = "../../mholt/caddy/caddy/caddymain/run.go"
-)