add systemd socket activation support

server.go

@@ -12,26 +12,57 @@ import (
 	"path"
 	"regexp"
 	"sort"
+	"strconv"
 	"strings"
+	"sync"
+	"sync/atomic"
 	"syscall"
+	"time"
 )
 
 type Server struct {
 	DataPath         string
+	LastConnection   time.Time
+	SocketActivation bool
+	stopListening    int32
+	activeWorkers    sync.WaitGroup
 }
 
-func (s Server) Run(listener *net.TCPListener) {
+func New(dataPath string) *Server {
-	for {
+	return &Server{dataPath, time.Now(), false, 0, sync.WaitGroup{}}
-		conn, e := listener.AcceptTCP()
+}
-		if e != nil {
+
-			fmt.Fprintf(os.Stderr, "Error: %v\n", e)
+func (s *Server) Run(listener *net.TCPListener) {
+	atomic.StoreInt32(&s.stopListening, 0)
+	s.activeWorkers.Add(1)
+	defer s.activeWorkers.Done()
+	defer listener.Close()
+	for atomic.LoadInt32(&s.stopListening) != 1 {
+		if e := listener.SetDeadline(time.Now().Add(time.Second)); e != nil {
+			fmt.Fprintf(os.Stderr, "Error setting deadline: %v\n", e)
 			continue
 		}
+		conn, err := listener.AcceptTCP()
+		if err != nil {
+			if err, ok := err.(net.Error); ok && err.Timeout() {
+				continue
+			} else {
+				fmt.Fprintf(os.Stderr, "Error: %v\n", err)
+				os.Exit(1)
+			}
+		}
 
+		s.activeWorkers.Add(1)
+		s.LastConnection = time.Now()
 		go s.handleConn(conn)
 	}
 }
 
+func (s *Server) Shutdown() {
+	atomic.StoreInt32(&s.stopListening, 1)
+	s.activeWorkers.Wait()
+}
+
 func readCidrs(path string) ([]net.IPNet, error) {
 	files, err := ioutil.ReadDir(path)
 	if err != nil {
@@ -118,8 +149,12 @@ func parseQuery(conn *net.TCPConn) map[int]interface{} {
 	return queryArgs
 }
 
-func (s Server) handleConn(conn *net.TCPConn) {
+func (s *Server) handleConn(conn *net.TCPConn) {
-	defer conn.Close()
+	defer func() {
+		conn.Close()
+		s.activeWorkers.Done()
+	}()
+
 	queryArgs := parseQuery(conn)
 	if queryArgs == nil {
 		return
@@ -145,7 +180,7 @@ func (s Server) handleConn(conn *net.TCPConn) {
 	}
 }
 
-func (s Server) printNet(conn *net.TCPConn, name string, ip net.IP) bool {
+func (s *Server) printNet(conn *net.TCPConn, name string, ip net.IP) bool {
 	routePath := path.Join(s.DataPath, name)
 	cidrs, err := readCidrs(routePath)
 	if err != nil {
@@ -163,7 +198,7 @@ func (s Server) printNet(conn *net.TCPConn, name string, ip net.IP) bool {
 	return found
 }
 
-func (s Server) printObject(conn *net.TCPConn, objType string, obj string) {
+func (s *Server) printObject(conn *net.TCPConn, objType string, obj string) {
 	f, err := os.Open(path.Join(s.DataPath, objType, obj))
 	defer f.Close()
 	if err != nil && !os.IsNotExist(err) {
@@ -176,8 +211,7 @@ type options struct {
 	Port          uint
 	Address       string
 	Registry      string
-	User     string
+	SocketTimeout float64
-	Group    string
 }
 
 func parseFlags() options {
@@ -185,6 +219,8 @@ func parseFlags() options {
 	flag.UintVar(&o.Port, "port", 43, "port to listen")
 	flag.StringVar(&o.Address, "address", "*", "address to listen")
 	flag.StringVar(&o.Registry, "registry", ".", "path to dn42 registry")
+	msg := "timeout in seconds before suspending the service when using socket activation"
+	flag.Float64Var(&o.SocketTimeout, "timeout", 10, msg)
 	flag.Parse()
 	if o.Address == "*" {
 		o.Address = ""
@@ -192,39 +228,98 @@ func parseFlags() options {
 	return o
 }
 
-func main() {
+func Listeners() []*net.TCPListener {
-	opts := parseFlags()
+	defer os.Unsetenv("LISTEN_PID")
-	registryPath := path.Join(opts.Registry, "data")
+	defer os.Unsetenv("LISTEN_FDS")
 
-	if _, err := os.Stat(registryPath); err != nil {
+	pid, err := strconv.Atoi(os.Getenv("LISTEN_PID"))
-		fmt.Fprintf(os.Stderr,
+	if err != nil || pid != os.Getpid() {
-			"Cannot access '%s', should be in the registry repository: %s\n",
+		return nil
-			registryPath,
-			err)
-		os.Exit(1)
 	}
 
-	address := opts.Address + ":" + fmt.Sprint(opts.Port)
+	nfds, err := strconv.Atoi(os.Getenv("LISTEN_FDS"))
+	if err != nil || nfds == 0 {
+		return nil
+	}
+
+	listeners := make([]*net.TCPListener, 0)
+	for fd := 3; fd < 3+nfds; fd++ {
+		syscall.CloseOnExec(fd)
+		file := os.NewFile(uintptr(fd), "LISTEN_FD_"+strconv.Itoa(fd))
+		if listener, err := net.FileListener(file); err == nil {
+			if l, ok := listener.(*net.TCPListener); ok {
+				listeners = append(listeners, l)
+			}
+		}
+	}
+
+	return listeners
+}
+
+func checkDataPath(registry string) (string, error) {
+	dataPath := path.Join(registry, "data")
+
+	if _, err := os.Stat(dataPath); err != nil {
+		return "", fmt.Errorf("Cannot access '%s', should be in the registry repository: %s\n",
+			dataPath,
+			err)
+	}
+	return dataPath, nil
+}
+
+func createServer(opts options) (*Server, error) {
+	dataPath, err := checkDataPath(opts.Registry)
+	if err != nil {
+		return nil, err
+	}
+	server := New(dataPath)
+
+	if listeners := Listeners(); len(listeners) > 0 {
+		fmt.Printf("socket action detected\n")
+		server.SocketActivation = true
+		for _, listener := range listeners {
+			go server.Run(listener)
+		}
+	} else {
+		address := opts.Address + ":" + strconv.Itoa(int(opts.Port))
 		listener, err := net.Listen("tcp", address)
+		if err != nil {
+			return nil, err
+		}
+		go server.Run(listener.(*net.TCPListener))
+	}
+	return server, nil
+}
+
+func main() {
+	opts := parseFlags()
+	server, err := createServer(opts)
 	if err != nil {
 		fmt.Fprintf(os.Stderr, "Error: %v\n", err)
 		os.Exit(1)
 	}
 
-	server := Server{registryPath}
+	signals := make(chan os.Signal, 1)
-	go server.Run(listener.(*net.TCPListener))
+	signal.Notify(signals, os.Interrupt)
-
+	signal.Notify(signals, syscall.SIGTERM)
-	c := make(chan os.Signal, 1)
+	signal.Notify(signals, syscall.SIGINT)
-	signal.Notify(c, os.Interrupt)
-	signal.Notify(c, syscall.SIGTERM)
-	signal.Notify(c, syscall.SIGINT)
 
+	if server.SocketActivation {
+	Out:
 		for {
 			select {
-		case <-c:
+			case <-signals:
-			fmt.Printf("Shutting socket down\n")
+				break Out
-			listener.Close()
+			case <-time.After(time.Second * 3):
-			os.Exit(0)
+				if time.Since(server.LastConnection).Seconds() >= opts.SocketTimeout {
+					break Out
 				}
 			}
+		}
+	} else {
+		<-signals
+	}
+
+	fmt.Printf("Shutting socket(s) down (takes up to 1s)\n")
+	server.Shutdown()
 }