warren/cmd/warrend/main.go

package main

import (
	"fmt"
	"github.com/dballard/warren/cmd/warrend/conf"
	"github.com/dballard/warren/lib/warren"
	"log"
	"net"
	"os"
	"os/signal"
	"strings"
	"syscall"
	"bufio"
)

var cmdSock net.Listener

type SockCommand struct {
	warren.Command
	Run func(cmd *SockCommand, writer *bufio.Writer, path string, args ...string)
}

func (cmd *SockCommand) GetName() string {
	return cmd.Name
}

func (cmd *SockCommand) GetSummary() string {
	return cmd.Summary
}

func (cmd *SockCommand) Exec(writer *bufio.Writer, path string, args []string) {
	cmd.Flag.Usage = func() {
		// helpFunc(c, c.Name)
	}
	cmd.Flag.Parse(args)
	cmd.Run(cmd, writer, path, cmd.Flag.Args()...)
}

func sockCmdHandle(c net.Conn) {
	reader := bufio.NewReader(c)
	writer := bufio.NewWriter(c)
	defer c.Close()
	fmt.Println("cmdHandle")
	data, err := reader.ReadString('\n')
	if err != nil {
		fmt.Println(err)
		writer.WriteString("read error:"+err.Error())
		return
	}
	line := strings.TrimSpace(data)
	fmt.Println("cmdHandle got command:", line)

	// BAD: TODO: handle quotes "
	// path|arg1|arg2|...
	args := strings.Split(line, "|")
	if len(args) < 2 {
		writer.WriteString("Format error")
		return
	}

	// parse args
	var icmd = warren.GetCommand(args[1])
	// get and run cmd
	if icmd == nil {
		writer.WriteString("Unknown command "+args[1])
		// usage ?

		return
	}
	icmd.(*SockCommand).Exec(writer, args[0], args[2:])
	writer.WriteByte(0)
	writer.Flush()
}

// determine if warrend is already running
func checkRunning() bool {
	if pfile, err := os.Open(warren.WarrendPidFile); err == nil {
		// pid file exists
		var pid int
		n, err := fmt.Fscanf(pfile, "%d", &pid)
		if n == 0 || err != nil {
			// read error
			return false
		}
		// is there a runnign process with the pid?
		err = syscall.Kill(pid, 0)
		return err == nil

	} else {
		// pid file does not exist
		return false
	}
}

func initRunning() {
	// rm old sock
	os.Remove(warren.CmdSockAddr)

	// incase gone
	os.Mkdir(warren.WarrendRunDir, 0775)

	// drop down pid
	pidfile, err := os.Create(warren.WarrendPidFile)
	if err != nil {
		log.Fatal("Could not create pidfile: ", warren.WarrendPidFile, " : ", err)
	}
	pid := os.Getpid()
	fmt.Fprintf(pidfile, "%d", pid)
	pidfile.Close()

	cmdSock, err = net.Listen(warren.CmdSockType, warren.CmdSockAddr)
	if err != nil {
		fmt.Println(cmdSock)
		Shutdown()
		log.Fatal("Listen on sock error:", err)
	}

}

func Shutdown() {
	os.Remove(warren.WarrendPidFile)
	if cmdSock != nil {
		cmdSock.Close()
	}
}

func main() {
	if checkRunning() {
		log.Fatal("warrend instance already running")
	}

	go conf.Run()
	if cinit := <-conf.Init; !cinit {
		return
	}

	initRunning()

	warren.RegisterCommand(listCmd)
	warren.RegisterCommand(regCmd)
	warren.RegisterCommand(unregCmd)
	warren.RegisterCommand(poolCmd)

	// Handle signals
	c := make(chan os.Signal, 1)
	signal.Notify(c,
		syscall.SIGINT,
		syscall.SIGTERM,
		syscall.SIGQUIT,
		syscall.SIGHUP,
	)
	go func() {
		<-c
		Shutdown()
	}()

	fmt.Println("warrend starting and listening...")

	// Listen
	for {
		fd, err := cmdSock.Accept()
		fmt.Println("accept")
		if err != nil {
			Shutdown()
			log.Fatal("accpet error:", err)
		}
		go sockCmdHandle(fd)
	}
	Shutdown()
}