Fix connection leak (#680)

Signed-off-by: David Chung <[email protected]>

Author: David Chung

pkg/controller/ingress/swarm/init.go

@@ -37,8 +37,17 @@ type Spec struct {
 	Docker `json:",inline" yaml:",inline"`
 }
 
-// RoutesFromSwarmServices determines the routes based on the services running in the Docker swarm
-func RoutesFromSwarmServices(properties *types.Any,
+type handler struct {
+	dockerClient docker.APIClientCloser
+}
+
+// Close implements io.Closer
+func (h *handler) Close() error {
+	return h.dockerClient.Close()
+}
+
+// Routes implements ingress/types/RouteHandler
+func (h *handler) Routes(properties *types.Any,
 	options ingress.Options) (map[ingress.Vhost][]loadbalancer.Route, error) {
 
 	spec := Spec{}
@@ -63,10 +72,17 @@ func RoutesFromSwarmServices(properties *types.Any,
 	}
 
 	log.Info("Connected to Docker", "client", dockerClient)
+	h.dockerClient = dockerClient
+
 	routes, err := NewServiceRoutes(dockerClient).SetOptions(options).Build()
 	if err != nil {
 		return nil, err
 	}
 
 	return routes.List()
 }
+
+// RoutesFromSwarmServices determines the routes based on the services running in the Docker swarm
+func RoutesFromSwarmServices() (ingress.RouteHandler, error) {
+	return &handler{}, nil
+}

pkg/controller/ingress/types/funcs.go

@@ -1,6 +1,7 @@
 package types
 
 import (
+	"io"
 	"sync"
 
 	logutil "github.com/docker/infrakit/pkg/log"
@@ -12,12 +13,19 @@ import (
 var (
 	log               = logutil.New("module", "controller/ingress/types")
 	debugV            = logutil.V(300)
-	routeHandlers     = map[string]func(*types.Any, Options) (map[Vhost][]loadbalancer.Route, error){}
+	routeHandlers     = map[string]func() (RouteHandler, error){}
 	routeHandlersLock = sync.Mutex{}
 )
 
+// RouteHandler is the interface that different modules must support
+type RouteHandler interface {
+	io.Closer
+	// Routes returns a map of vhost and loadbalancer routes given the input blob
+	Routes(*types.Any, Options) (map[Vhost][]loadbalancer.Route, error)
+}
+
 // RegisterRouteHandler registers a package specific handler for determining the L4 routes (e.g. static or swarm)
-func RegisterRouteHandler(key string, f func(*types.Any, Options) (map[Vhost][]loadbalancer.Route, error)) {
+func RegisterRouteHandler(key string, f func() (RouteHandler, error)) {
 
 	routeHandlersLock.Lock()
 	defer routeHandlersLock.Unlock()
@@ -89,15 +97,22 @@ func (p Properties) Routes(options Options) (result map[Vhost][]loadbalancer.Rou
 		result[spec.Vhost] = spec.Routes
 
 		for key, config := range spec.RouteSources {
-			handler, has := routeHandlers[key]
+			handlerFunc, has := routeHandlers[key]
 
 			log.Debug("route handler", "key", key, "exists", has, "V", debugV)
 			if !has {
 				continue
 			}
 
 			log.Debug("calling route handler", "config", config, "options", options, "V", debugV)
-			vhostRoutes, err := handler(config, options)
+
+			handler, err := handlerFunc()
+			if err != nil {
+				return nil, err
+			}
+			defer handler.Close()
+
+			vhostRoutes, err := handler.Routes(config, options)
 
 			log.Debug("found routes", "routesByVhost", vhostRoutes, "err", err)
 			if err != nil {

pkg/controller/ingress/types/funcs_test.go

@@ -181,6 +181,15 @@ func TestGroupsInstanceIDs(t *testing.T) {
 	}, mm)
 }
 
+type routeHandlerFunc func(customConfig *types.Any, options Options) (map[Vhost][]loadbalancer.Route, error)
+
+func (r routeHandlerFunc) Close() error {
+	return nil
+}
+func (r routeHandlerFunc) Routes(customConfig *types.Any, options Options) (map[Vhost][]loadbalancer.Route, error) {
+	return r(customConfig, options)
+}
+
 func TestRegisterRouteHandler(t *testing.T) {
 
 	RegisterRouteHandler("nil", nil)
@@ -192,7 +201,9 @@ func TestRegisterRouteHandler(t *testing.T) {
 		calledChan <- customConfig
 		return <-routesChan, nil
 	}
-	RegisterRouteHandler("test", f)
+	RegisterRouteHandler("test", func() (RouteHandler, error) {
+		return routeHandlerFunc(f), nil
+	})
 	require.Equal(t, 1, len(routeHandlers))
 
 	vhost := Vhost("test")