write a PID file

main.go

@@ -10,14 +10,15 @@ import (
 	"net"
 	"net/http"
 	"os"
+	"os/signal"
 	"strings"
+	"syscall"
 	"time"
 
 	"github.com/BurntSushi/toml"
 	"github.com/google/gopacket"
 	"github.com/google/gopacket/layers"
 	"github.com/google/gopacket/pcap"
-	"github.com/kr/pretty"
 	"github.com/nats-io/nats"
 	"github.com/nats-io/nats/encoders/protobuf"
 
@@ -39,6 +40,7 @@ var (
 	useXForwardedAsSource = flag.Bool("use-x-forwarded", false, "Use the IP address in X-Forwarded-For as source")
 	pidFile               = flag.String("pidfile", "/var/run/munchclient.pid", "The location of the PID file")
 	configFile            = flag.String("config", "", "The location of the TOML config file")
+	beQuiet               = flag.Bool("quiet", true, "Be quiet")
 	doVersion             = flag.Bool("version", false, "Show version information")
 
 	natsEC  *nats.EncodedConn
@@ -67,6 +69,7 @@ type Config struct {
 	RequestsFile          string
 	UseXForwardedAsSource bool
 	PidFile               string
+	Quiet                 bool
 }
 
 func init() {
@@ -89,6 +92,7 @@ func loadConfig() {
 	config.RequestsFile = *requestsFile
 	config.UseXForwardedAsSource = *useXForwardedAsSource
 	config.PidFile = *pidFile
+	config.Quiet = *beQuiet
 
 	if *configFile == "" {
 		return
@@ -103,8 +107,6 @@ func loadConfig() {
 	if _, err = toml.DecodeFile(*configFile, &config); err != nil {
 		log.Printf("%s: %s\n", *configFile, err)
 	}
-
-	pretty.Println(config)
 }
 
 func main() {
@@ -115,19 +117,39 @@ func main() {
 
 	loadConfig()
 
-	if err := pidfile.Write(config.PidFile); err != nil {
-		log.Fatalf("munchclient: %s\n", err)
-		os.Exit(1)
-	}
-	defer pidfile.Remove(config.PidFile)
+	// Clean up before exiting
+	//
+	sigs := make(chan os.Signal, 1)
+	signal.Notify(sigs, syscall.SIGINT, syscall.SIGTERM)
+	go func() {
+		<-sigs
+		pidfile.Remove(config.PidFile)
+		os.Exit(0)
+	}()
 
-	go func(c *uint64) {
-		for {
-			fmt.Printf("%d requests per second\n", *c)
-			*c = 0
-			time.Sleep(time.Second)
+	// Write the PID file
+	//
+	err := pidfile.Write(config.PidFile)
+	// if the PID file is stale remove it and create a new one with the current PID
+	if err == pidfile.ErrFileStale {
+		err2 := os.Remove(config.PidFile)
+		if err2 != nil {
+			log.Fatalf("%s: %s\n", config.PidFile, err)
 		}
-	}(&count)
+	} else if err != nil { // permission denied or program is still running
+		log.Fatalf("%s: %s\n", config.PidFile, err)
+	}
+
+	// Output how many requests per second were sent
+	if !config.Quiet {
+		go func(c *uint64) {
+			for {
+				fmt.Printf("%d requests per second\n", *c)
+				*c = 0
+				time.Sleep(time.Second)
+			}
+		}(&count)
+	}
 
 	// NATS
 	//

munchclient.toml

@@ -7,5 +7,6 @@ NatsURL = "nats://127.0.0.1:4222"
 NatsQueue = "requests"
 UseXForwardedAsSource = true
 PidFile = "/tmp/munchclient.pid"
+Quiet = true
 # SleepFor = 100µs
 # RequestsFile = requests.csv