Add support for running as windows service

exporter.go

@@ -48,15 +48,29 @@ var (
 	metricsPath       = flag.String("web.metrics-path", "/metrics", "Path under which the metrics are available")
 )
 
+type startupTask func() error
+
+var quitSig = make(chan os.Signal, 1)
+var startupTasks = make(map[string]startupTask, 0)
+
 func main() {
 	flag.Parse()
 
+	// Startup tasks can be registered to perform (typically os-specific) initialization
+	for taskName, task := range startupTasks {
+		log.Debugf("Executing startup task '%s'", taskName)
+		err := task()
+		if err != nil {
+			log.Fatal(err)
+		}
+	}
+
 	if *listInputTypes {
 		fmt.Println(input.GetAvailableInputs())
 		os.Exit(0)
 	}
 	if *inputType == "" {
-		fmt.Printf("-input.type is required. The following input types are available:\n%v", input.GetAvailableInputs())
+		log.Errorf("-input.type is required. The following input types are available:\n%v", input.GetAvailableInputs())
 		os.Exit(1)
 	}
 
@@ -77,7 +91,6 @@ func main() {
 	prometheus.MustRegister(totalLines)
 
 	// Setup signal handling
-	quitSig := make(chan os.Signal, 1)
 	signal.Notify(quitSig, os.Interrupt)
 
 	// Configure input

service_windows.go

@@ -0,0 +1,201 @@
+package main
+
+import (
+	"flag"
+	"fmt"
+	"os"
+
+	"github.com/prometheus/common/log"
+	"golang.org/x/sys/windows/svc"
+	"golang.org/x/sys/windows/svc/eventlog"
+	"golang.org/x/sys/windows/svc/mgr"
+)
+
+var (
+	installService   = flag.Bool("service.install", false, "Install as Windows service and exit")
+	uninstallService = flag.Bool("service.uninstall", false, "Uninstall Windows service and exit")
+	serviceName      = flag.String("service.name", "stream_exporter", "Name of Windows service to run, install or uninstall")
+)
+
+func init() {
+	startupTasks["processServiceInstallationFlags"] = processServiceInstallationFlags
+	startupTasks["runAsService"] = runAsServiceIfNoninteractive
+}
+
+func processServiceInstallationFlags() error {
+	if *installService {
+		if *serviceName == "" {
+			return fmt.Errorf("Cannot install service, name not set")
+		}
+		err := install(*serviceName)
+		if err != nil {
+			return err
+		} else {
+			log.Info("Successfully installed service")
+			os.Exit(0)
+		}
+	} else if *uninstallService {
+		if *serviceName == "" {
+			return fmt.Errorf("Cannot uninstall service, name not set")
+		}
+		err := uninstall(*serviceName)
+		if err != nil {
+			return err
+		} else {
+			log.Info("Successfully uninstalled service")
+			os.Exit(0)
+		}
+	}
+
+	return nil
+}
+
+func runAsServiceIfNoninteractive() error {
+	if *installService || *uninstallService {
+		return nil
+	}
+
+	isInteractive, err := svc.IsAnInteractiveSession()
+	if err != nil {
+		return err
+	}
+
+	if !isInteractive {
+		if *serviceName == "" {
+			return fmt.Errorf("Cannot start service, name not set")
+		}
+		log.Debugf("Starting service %s", *serviceName)
+		go svc.Run(*serviceName, &streamExporterService{stopCh: quitSig})
+	}
+
+	return nil
+}
+
+type streamExporterService struct {
+	stopCh chan<- os.Signal
+}
+
+func (s *streamExporterService) Execute(args []string, r <-chan svc.ChangeRequest, changes chan<- svc.Status) (ssec bool, errno uint32) {
+	const cmdsAccepted = svc.AcceptStop | svc.AcceptShutdown
+	changes <- svc.Status{State: svc.StartPending}
+	changes <- svc.Status{State: svc.Running, Accepts: cmdsAccepted}
+loop:
+	for {
+		select {
+		case c := <-r:
+			switch c.Cmd {
+			case svc.Interrogate:
+				changes <- c.CurrentStatus
+			case svc.Stop, svc.Shutdown:
+				s.stopCh <- os.Interrupt
+				break loop
+			default:
+				log.Error(fmt.Sprintf("Unexpected control request #%d", c))
+			}
+		}
+	}
+	changes <- svc.Status{State: svc.StopPending}
+	return
+}
+
+// ServiceManager integration adapted from https://github.com/martinlindhe/winservice
+
+func install(name string) error {
+	exepath, err := os.Executable()
+	if err != nil {
+		return fmt.Errorf("Could not determine own path, %s", err)
+	}
+
+	m, err := mgr.Connect()
+	if err != nil {
+		return err
+	}
+	defer m.Disconnect()
+
+	s, err := m.OpenService(name)
+	if err == nil {
+		s.Close()
+		return fmt.Errorf("Service %s already exists", name)
+	}
+	config := mgr.Config{
+		DisplayName: name,
+		StartType:   mgr.StartAutomatic,
+		Description: "Prometheus stream_exporter",
+	}
+	args := getServiceArgs(os.Args)
+	s, err = m.CreateService(name, exepath, config, args...)
+	if err != nil {
+		return err
+	}
+	defer s.Close()
+
+	err = eventlog.InstallAsEventCreate(name, eventlog.Error|eventlog.Warning|eventlog.Info)
+	if err != nil {
+		s.Delete()
+		return fmt.Errorf("Install() failed: %s", err)
+	}
+	return nil
+}
+
+// Remove first arg (program name), and -service.install flag
+// Ensure -service.name is included
+func getServiceArgs(inputArgs []string) []string {
+	args := make([]string, len(inputArgs)-1)
+	copy(args, inputArgs[1:])
+	foundServiceName := false
+
+	for idx, arg := range args {
+		switch arg {
+		case "-service.install":
+			args = append(args[:idx], args[idx+1:]...)
+		case "-service.name":
+			foundServiceName = true
+		}
+	}
+
+	if !foundServiceName {
+		args = append(args, "-service.name")
+		args = append(args, "stream_exporter")
+	}
+
+	args = append(args, "-log.format")
+	args = append(args, "logger:eventlog?name=stream_exporter")
+
+	return args
+}
+
+func uninstall(name string) error {
+	m, err := mgr.Connect()
+	if err != nil {
+		return err
+	}
+	defer m.Disconnect()
+
+	s, err := m.OpenService(name)
+	if err != nil {
+		return fmt.Errorf("Service %q is not installed", name)
+	}
+	defer s.Close()
+
+	status, err := s.Query()
+	if err != nil {
+		return fmt.Errorf("Could not query for service %q state: %v", name, err)
+	}
+
+	if status.State != svc.Stopped {
+		_, err := s.Control(svc.Stop)
+		if err != nil {
+			return fmt.Errorf("Could not send stop signal to service %q", err)
+		}
+	}
+
+	err = s.Delete()
+	if err != nil {
+		return err
+	}
+	err = eventlog.Remove(name)
+	if err != nil {
+		return fmt.Errorf("Remove() failed: %s", err)
+	}
+	return nil
+}