Fix graceful reload (#141)

Fix CoreDNS graceful reloading. This uses the same stuff as Caddy
(obviously), but extends it for UDP listeners as well. Also add to the
README that we *will* call Shutdown for middleware.

Fixes #4

core/core.go

@@ -21,6 +21,7 @@ import (
 	"fmt"
 	"io/ioutil"
 	"log"
+	"net"
 	"os"
 	"path"
 	"strings"
@@ -197,36 +198,52 @@ func startServers(groupings bindingGroup) error {
 			s.TLSConfig.GetCertificate = https.GetCertificate
 		}
 
-		var ln server.ListenerFile
+		var (
-		/*
+			ln net.Listener
-			if IsRestart() {
+			pc net.PacketConn
-				// Look up this server's listener in the map of inherited file descriptors;
+		)
-				// if we don't have one, we must make a new one (later).
-				if fdIndex, ok := loadedGob.ListenerFds[s.Addr]; ok {
-					file := os.NewFile(fdIndex, "")
 
-					fln, err := net.FileListener(file)
+		if IsRestart() {
-					if err != nil {
+			// Look up this server's listener in the map of inherited file descriptors; if we don't have one, we must make a new one (later).
-						return err
+			if fdIndex, ok := loadedGob.ListenerFds["tcp"+s.Addr]; ok {
-					}
+				file := os.NewFile(fdIndex, "")
 
-					ln, ok = fln.(server.ListenerFile)
+				fln, err := net.FileListener(file)
-					if !ok {
+				if err != nil {
-						return errors.New("listener for " + s.Addr + " was not a ListenerFile")
+					return err
-					}
-
-					file.Close()
-					delete(loadedGob.ListenerFds, s.Addr)
 				}
+
+				ln, ok = fln.(*net.TCPListener)
+				if !ok {
+					return errors.New("listener for " + s.Addr + " was not a *net.TCPListener")
+				}
+
+				file.Close()
+				delete(loadedGob.ListenerFds, "tcp"+s.Addr)
 			}
-		*/
+			if fdIndex, ok := loadedGob.ListenerFds["udp"+s.Addr]; ok {
+				file := os.NewFile(fdIndex, "")
+
+				fpc, err := net.FilePacketConn(file)
+				if err != nil {
+					return err
+				}
+
+				pc, ok = fpc.(*net.UDPConn)
+				if !ok {
+					return errors.New("packetConn for " + s.Addr + " was not a *net.PacketConn")
+				}
+
+				file.Close()
+				delete(loadedGob.ListenerFds, "udp"+s.Addr)
+			}
+		}
 
 		wg.Add(1)
-		go func(s *server.Server, ln server.ListenerFile) {
+		go func(s *server.Server, ln net.Listener, pc net.PacketConn) {
 			defer wg.Done()
 
-			// run startup functions that should only execute when
+			// run startup functions that should only execute when the original parent process is starting.
-			// the original parent process is starting.
 			if !IsRestart() && !startedBefore {
 				err := s.RunFirstStartupFuncs()
 				if err != nil {
@@ -236,14 +253,12 @@ func startServers(groupings bindingGroup) error {
 			}
 
 			// start the server
-			// TODO(miek): for now will always be nil, so we will run ListenAndServe()
+			if ln != nil && pc != nil {
-			// TODO(miek): this is also why graceful restarts don't work.
+				errChan <- s.Serve(ln, pc)
-			if ln != nil {
-				//errChan <- s.Serve(ln)
 			} else {
 				errChan <- s.ListenAndServe()
 			}
-		}(s, ln)
+		}(s, ln, pc)
 
 		startupWg.Add(1)
 		go func(s *server.Server) {

core/restart.go

@@ -81,9 +81,14 @@ func Restart(newCorefile Input) error {
 
 	// Add file descriptors of all the sockets
 	serversMu.Lock()
-	for i, s := range servers {
+	j := 0
+	for _, s := range servers {
 		extraFiles = append(extraFiles, s.ListenerFd())
-		crfileGob.ListenerFds[s.Addr] = uintptr(4 + i) // 4 fds come before any of the listeners
+		extraFiles = append(extraFiles, s.PacketConnFd())
+		// So this will be 0 1 2 3 TCP UDP TCP UDP ... etc.
+		crfileGob.ListenerFds["tcp"+s.Addr] = uintptr(4 + j)     // 4 fds come before any of the listeners
+		crfileGob.ListenerFds["udp"+s.Addr] = uintptr(4 + j + 1) // add udp after that
+		j += 2
 	}
 	serversMu.Unlock()
 

server/config.go

@@ -42,7 +42,8 @@ type Config struct {
 	FirstStartup []func() error
 
 	// Functions (or methods) to execute when the server quits;
-	// these are executed in response to SIGINT and are blocking
+	// these are executed in response to SIGINT and are blocking. These
+	// function are *also* called when we are restarting.
 	Shutdown []func() error
 
 	// The path to the configuration file from which this was loaded

server/graceful.go

@@ -1,76 +0,0 @@
-package server
-
-import (
-	"net"
-	"sync"
-	"syscall"
-)
-
-// newGracefulListener returns a gracefulListener that wraps l and
-// uses wg (stored in the host server) to count connections.
-func newGracefulListener(l ListenerFile, wg *sync.WaitGroup) *gracefulListener {
-	gl := &gracefulListener{ListenerFile: l, stop: make(chan error), httpWg: wg}
-	go func() {
-		<-gl.stop
-		gl.Lock()
-		gl.stopped = true
-		gl.Unlock()
-		gl.stop <- gl.ListenerFile.Close()
-	}()
-	return gl
-}
-
-// gracefuListener is a net.Listener which can
-// count the number of connections on it. Its
-// methods mainly wrap net.Listener to be graceful.
-type gracefulListener struct {
-	ListenerFile
-	stop       chan error
-	stopped    bool
-	sync.Mutex                 // protects the stopped flag
-	httpWg     *sync.WaitGroup // pointer to the host's wg used for counting connections
-}
-
-// Accept accepts a connection.
-func (gl *gracefulListener) Accept() (c net.Conn, err error) {
-	c, err = gl.ListenerFile.Accept()
-	if err != nil {
-		return
-	}
-	c = gracefulConn{Conn: c, httpWg: gl.httpWg}
-	gl.httpWg.Add(1)
-	return
-}
-
-// Close immediately closes the listener.
-func (gl *gracefulListener) Close() error {
-	gl.Lock()
-	if gl.stopped {
-		gl.Unlock()
-		return syscall.EINVAL
-	}
-	gl.Unlock()
-	gl.stop <- nil
-	return <-gl.stop
-}
-
-// gracefulConn represents a connection on a
-// gracefulListener so that we can keep track
-// of the number of connections, thus facilitating
-// a graceful shutdown.
-type gracefulConn struct {
-	net.Conn
-	httpWg *sync.WaitGroup // pointer to the host server's connection waitgroup
-}
-
-// Close closes c's underlying connection while updating the wg count.
-func (c gracefulConn) Close() error {
-	err := c.Conn.Close()
-	if err != nil {
-		return err
-	}
-	// close can fail on http2 connections (as of Oct. 2015, before http2 in std lib)
-	// so don't decrement count unless close succeeds
-	c.httpWg.Done()
-	return nil
-}

server/server.go

@@ -32,15 +32,15 @@ type Server struct {
 	Addr   string // Address we listen on
 	mux    *dns.ServeMux
 	server [2]*dns.Server // by convention 0 is tcp and 1 is udp
-	tcp    net.Listener
+
-	udp    net.PacketConn
+	tcp        net.Listener
+	udp        net.PacketConn
+	listenerMu sync.Mutex // protects listener and packetconn
 
 	tls         bool // whether this server is serving all HTTPS hosts or not
 	TLSConfig   *tls.Config
 	OnDemandTLS bool             // whether this server supports on-demand TLS (load certs at handshake-time)
 	zones       map[string]zone  // zones keyed by their address
-	listener    ListenerFile     // the listener which is bound to the socket
-	listenerMu  sync.Mutex       // protects listener
 	dnsWg       sync.WaitGroup   // used to wait on outstanding connections
 	startChan   chan struct{}    // used to block until server is finished starting
 	connTimeout time.Duration    // the maximum duration of a graceful shutdown
@@ -48,12 +48,6 @@ type Server struct {
 	SNICallback func(clientHello *tls.ClientHelloInfo) (*tls.Certificate, error)
 }
 
-// ListenerFile represents a listener.
-type ListenerFile interface {
-	net.Listener
-	File() (*os.File, error)
-}
-
 // OptionalCallback is a function that may or may not handle a request.
 // It returns whether or not it handled the request. If it handled the
 // request, it is presumed that no further request handling should occur.
@@ -145,25 +139,31 @@ func (s *Server) LocalAddr() (net.Addr, net.Addr) {
 	return tcp, udp
 }
 
-// Serve starts the server with an existing listener. It blocks until the
+// Serve starts the server with an existing listener. It blocks until the server stops.
-// server stops.
+func (s *Server) Serve(ln net.Listener, pc net.PacketConn) error {
-/*
-func (s *Server) Serve(ln ListenerFile) error {
-	// TODO(miek): Go DNS has no server stuff that allows you to give it a listener
-	// and use that.
 	err := s.setup()
 	if err != nil {
-		defer close(s.startChan) // MUST defer so error is properly reported, same with all cases in this file
+		close(s.startChan) // MUST defer so error is properly reported, same with all cases in this file
 		return err
 	}
-	return s.serve(ln)
+	s.listenerMu.Lock()
+	s.server[0] = &dns.Server{Listener: ln, Net: "tcp", Handler: s.mux}
+	s.tcp = ln
+	s.server[1] = &dns.Server{PacketConn: pc, Net: "udp", Handler: s.mux}
+	s.udp = pc
+	s.listenerMu.Unlock()
+
+	go func() {
+		s.server[0].ActivateAndServe()
+	}()
+	close(s.startChan)
+	return s.server[1].ActivateAndServe()
 }
-*/
 
 // ListenAndServe starts the server with a new listener. It blocks until the server stops.
 func (s *Server) ListenAndServe() error {
 	err := s.setup()
-	// defer close(s.startChan) // Don't understand why defer wouldn't actually work in this method.
+	// defer close(s.startChan) // Don't understand why defer wouldn't actually work in this method (prolly cause the last ActivateAndServe does not actually return?
 	if err != nil {
 		close(s.startChan)
 		return err
@@ -266,8 +266,11 @@ func (s *Server) Stop() (err error) {
 
 	// Close the listener now; this stops the server without delay
 	s.listenerMu.Lock()
-	if s.listener != nil {
+	if s.tcp != nil {
-		err = s.listener.Close()
+		err = s.tcp.Close()
+	}
+	if s.udp != nil {
+		err = s.udp.Close()
 	}
 
 	for _, s1 := range s.server {
@@ -291,8 +294,21 @@ func (s *Server) WaitUntilStarted() {
 func (s *Server) ListenerFd() *os.File {
 	s.listenerMu.Lock()
 	defer s.listenerMu.Unlock()
-	if s.listener != nil {
+	if s.tcp != nil {
-		file, _ := s.listener.File()
+		file, _ := s.tcp.(*net.TCPListener).File()
+		return file
+	}
+	return nil
+}
+
+// PacketConnFd gets a dup'ed file of the packetconn. If there
+// is no underlying file, the return value will be nil. It
+// is the caller's responsibility to close the file.
+func (s *Server) PacketConnFd() *os.File {
+	s.listenerMu.Lock()
+	defer s.listenerMu.Unlock()
+	if s.udp != nil {
+		file, _ := s.udp.(*net.UDPConn).File()
 		return file
 	}
 	return nil