Merge pull request #2030 from fcrisciani/change-server

Call the Shutdown on the server

libnetwork/controller.go

@@ -1306,19 +1306,19 @@ func (c *controller) Stop() {
 // StartDiagnose start the network diagnose mode
 func (c *controller) StartDiagnose(port int) {
 	c.Lock()
-	defer c.Unlock()
 	if !c.DiagnoseServer.IsDebugEnable() {
 		c.DiagnoseServer.EnableDebug("127.0.0.1", port)
 	}
+	c.Unlock()
 }
 
 // StopDiagnose start the network diagnose mode
 func (c *controller) StopDiagnose() {
 	c.Lock()
-	defer c.Unlock()
 	if c.DiagnoseServer.IsDebugEnable() {
 		c.DiagnoseServer.DisableDebug()
 	}
+	c.Unlock()
 }
 
 // IsDiagnoseEnabled returns true if the diagnose is enabled

libnetwork/diagnose/server.go

@@ -1,12 +1,12 @@
 package diagnose
 
 import (
+	"context"
 	"encoding/json"
 	"fmt"
-	"log"
-	"net"
 	"net/http"
 	"sync"
+	"sync/atomic"
 
 	stackdump "github.com/docker/docker/pkg/signal"
 	"github.com/docker/libnetwork/common"
@@ -36,7 +36,8 @@ var diagPaths2Func = map[string]HTTPHandlerFunc{
 // Server when the debug is enabled exposes a
 // This data structure is protected by the Agent mutex so does not require and additional mutex here
 type Server struct {
-	sk                net.Listener
+	enable            int32
+	srv               *http.Server
 	port              int
 	mux               *http.ServeMux
 	registeredHanders map[string]bool
@@ -51,64 +52,74 @@ func New() *Server {
 }
 
 // Init initialize the mux for the http handling and register the base hooks
-func (n *Server) Init() {
-	n.mux = http.NewServeMux()
+func (s *Server) Init() {
+	s.mux = http.NewServeMux()
 
 	// Register local handlers
-	n.RegisterHandler(n, diagPaths2Func)
+	s.RegisterHandler(s, diagPaths2Func)
 }
 
 // RegisterHandler allows to register new handlers to the mux and to a specific path
-func (n *Server) RegisterHandler(ctx interface{}, hdlrs map[string]HTTPHandlerFunc) {
-	n.Lock()
-	defer n.Unlock()
+func (s *Server) RegisterHandler(ctx interface{}, hdlrs map[string]HTTPHandlerFunc) {
+	s.Lock()
+	defer s.Unlock()
 	for path, fun := range hdlrs {
-		if _, ok := n.registeredHanders[path]; ok {
+		if _, ok := s.registeredHanders[path]; ok {
 			continue
 		}
-		n.mux.Handle(path, httpHandlerCustom{ctx, fun})
-		n.registeredHanders[path] = true
+		s.mux.Handle(path, httpHandlerCustom{ctx, fun})
+		s.registeredHanders[path] = true
 	}
 }
 
+// ServeHTTP this is the method called bu the ListenAndServe, and is needed to allow us to
+// use our custom mux
+func (s *Server) ServeHTTP(w http.ResponseWriter, r *http.Request) {
+	s.mux.ServeHTTP(w, r)
+}
+
 // EnableDebug opens a TCP socket to debug the passed network DB
-func (n *Server) EnableDebug(ip string, port int) {
-	n.Lock()
-	defer n.Unlock()
+func (s *Server) EnableDebug(ip string, port int) {
+	s.Lock()
+	defer s.Unlock()
 
-	n.port = port
+	s.port = port
 
-	if n.sk != nil {
+	if s.enable == 1 {
 		logrus.Info("The server is already up and running")
 		return
 	}
 
-	logrus.Infof("Starting the server listening on %d for commands", port)
-	// Create the socket
-	var err error
-	n.sk, err = net.Listen("tcp", fmt.Sprintf("%s:%d", ip, port))
-	if err != nil {
-		log.Fatal(err)
-	}
+	logrus.Infof("Starting the diagnose server listening on %d for commands", port)
+	srv := &http.Server{Addr: fmt.Sprintf("127.0.0.1:%d", port), Handler: s}
+	s.srv = srv
+	s.enable = 1
+	go func(n *Server) {
+		// Ingore ErrServerClosed that is returned on the Shutdown call
+		if err := srv.ListenAndServe(); err != nil && err != http.ErrServerClosed {
+			logrus.Errorf("ListenAndServe error: %s", err)
+			atomic.SwapInt32(&n.enable, 0)
+		}
+	}(s)
 
-	go func() {
-		http.Serve(n.sk, n.mux)
-	}()
 }
 
 // DisableDebug stop the dubug and closes the tcp socket
-func (n *Server) DisableDebug() {
-	n.Lock()
-	defer n.Unlock()
-	n.sk.Close()
-	n.sk = nil
+func (s *Server) DisableDebug() {
+	s.Lock()
+	defer s.Unlock()
+
+	s.srv.Shutdown(context.Background())
+	s.srv = nil
+	s.enable = 0
+	logrus.Info("Disabling the diagnose server")
 }
 
 // IsDebugEnable returns true when the debug is enabled
-func (n *Server) IsDebugEnable() bool {
-	n.Lock()
-	defer n.Unlock()
-	return n.sk != nil
+func (s *Server) IsDebugEnable() bool {
+	s.Lock()
+	defer s.Unlock()
+	return s.enable == 1
 }
 
 func notImplemented(ctx interface{}, w http.ResponseWriter, r *http.Request) {
@@ -197,6 +208,7 @@ func ParseHTTPFormOptions(r *http.Request) (bool, *JSONOutput) {
 func HTTPReply(w http.ResponseWriter, r *HTTPResult, j *JSONOutput) (int, error) {
 	var response []byte
 	if j.enable {
+		w.Header().Set("Content-Type", "application/json")
 		var err error
 		if j.prettyPrint {
 			response, err = json.MarshalIndent(r, "", "  ")