godns/main.go

package main

import (
	"flag"
	"fmt"
	"github.com/spf13/viper"
	"meow.tf/joker/godns/log"
	"meow.tf/joker/godns/settings"
	"os"
	"os/signal"
	"runtime"
	"runtime/pprof"
	"time"
)

const (
	Version = "0.3.0"
)

var (
	cfgFile string
)

func init() {
	flag.StringVar(&cfgFile, "config", "/etc/godns/godns.conf", "")
}

func main() {
	initLogger()

	viper.SetConfigFile(cfgFile)

	viper.AutomaticEnv()

	if err := viper.ReadInConfig(); err == nil {
		fmt.Println("Using config file:", viper.ConfigFileUsed())
	}

	serverSettings := settings.Server()

	server := &Server{
		host:     serverSettings.Host,
		port:     serverSettings.Port,
		rTimeout: 5 * time.Second,
		wTimeout: 5 * time.Second,
	}

	server.Run()

	log.Info("godns %s (%s) start", Version, runtime.Version())

	if settings.Debug() {
		go profileCPU()
		go profileMEM()
	}

	sig := make(chan os.Signal)
	signal.Notify(sig, os.Interrupt)


	<- sig
	log.Info("signal received, stopping")
}

func profileCPU() {
	f, err := os.Create("godns.cprof")
	if err != nil {
		log.Error("%s", err)
		return
	}

	pprof.StartCPUProfile(f)
	time.AfterFunc(6*time.Minute, func() {
		pprof.StopCPUProfile()
		f.Close()

	})
}

func profileMEM() {
	f, err := os.Create("godns.mprof")

	if err != nil {
		log.Error("%s", err)
		return
	}

	time.AfterFunc(5*time.Minute, func() {
		pprof.WriteHeapProfile(f)
		f.Close()
	})

}

func initLogger() {
	logSettings := settings.Log()

	if viper.GetBool("log.stdout") {
		log.SetLogger("console", nil)
	}

	if file := viper.GetString("log.file"); file != "" {
		log.SetLogger("file", map[string]interface{}{"file": file})
	}

	log.SetLevel(logSettings.LogLevel())
}

func init() {
	runtime.GOMAXPROCS(runtime.NumCPU())
}