add monitoring of metrics proxy servers

This change introduces a Monitor type which is made aware of a metrics
proxy's closure via a channel. It contains the logic for port selection
for a new metrics proxy and port freeing for a metrics proxy that has
been shut down.

Unit test for metrics method added.

resolves #607

Signed-off-by: Gavin Inglis <giinglis@amazon.com>

Author: ginglis13

snapshotter/app/service.go

@@ -17,7 +17,6 @@ import (
 	"context"
 	"fmt"
 	"net"
-	"net/url"
 	"os"
 	"os/signal"
 	"strconv"
@@ -56,7 +55,29 @@ func Run(config config.Config) error {
 	group, ctx := errgroup.WithContext(ctx)
 
 	cache := cache.NewSnapshotterCache()
-	snapshotter, err := initSnapshotter(ctx, config, cache)
+
+	var (
+		monitor          *metrics.Monitor
+		serviceDiscovery *discovery.ServiceDiscovery
+	)
+	if config.Snapshotter.Metrics.Enable {
+		sdHost := config.Snapshotter.Metrics.Host
+		sdPort := config.Snapshotter.Metrics.ServiceDiscoveryPort
+		serviceDiscovery := discovery.NewServiceDiscovery(sdHost, sdPort, cache)
+		monitor, err := initMetricsProxyMonitor(config.Snapshotter.Metrics.PortRange)
+		if err != nil {
+			log.G(ctx).WithError(err).Fatal("failed creating metrics proxy monitor")
+			return err
+		}
+		group.Go(func() error {
+			return serviceDiscovery.Serve()
+		})
+		group.Go(func() error {
+			return monitor.Start()
+		})
+	}
+
+	snapshotter, err := initSnapshotter(ctx, config, cache, monitor)
 	if err != nil {
 		log.G(ctx).WithFields(
 			logrus.Fields{"resolver": config.Snapshotter.Proxy.Address.Resolver.Type},
@@ -80,15 +101,6 @@ func Run(config config.Config) error {
 		return err
 	}
 
-	var serviceDiscovery *discovery.ServiceDiscovery
-	if config.Snapshotter.Metrics.Enable {
-		sdPort := config.Snapshotter.Metrics.ServiceDiscoveryPort
-		serviceDiscovery = discovery.NewServiceDiscovery(config.Snapshotter.Metrics.Host, sdPort, cache)
-		group.Go(func() error {
-			return serviceDiscovery.Serve()
-		})
-	}
-
 	group.Go(func() error {
 		return grpcServer.Serve(listener)
 	})
@@ -101,10 +113,13 @@ func Run(config config.Config) error {
 			if err := snapshotter.Close(); err != nil {
 				log.G(ctx).WithError(err).Error("failed to close snapshotter")
 			}
-			if serviceDiscovery != nil {
+			if config.Snapshotter.Metrics.Enable {
 				if err := serviceDiscovery.Shutdown(ctx); err != nil {
 					log.G(ctx).WithError(err).Error("failed to shutdown service discovery server")
 				}
+				// Senders to this channel would panic if it is closed. However snapshotter.Close() will
+				// shutdown all metrics proxies and ensure there are no more senders over the channel.
+				monitor.Stop()
 			}
 		}()
 
@@ -141,23 +156,19 @@ func initResolver(config config.Config) (proxyaddress.Resolver, error) {
 const base10 = 10
 const bits32 = 32
 
-func initSnapshotter(ctx context.Context, config config.Config, cache cache.Cache) (snapshots.Snapshotter, error) {
+func initSnapshotter(ctx context.Context, config config.Config, cache cache.Cache, monitor *metrics.Monitor) (snapshots.Snapshotter, error) {
 	resolver, err := initResolver(config)
 	if err != nil {
 		return nil, err
 	}
 
-	newProxySnapshotterFunc := func(ctx context.Context, namespace string) (*proxy.RemoteSnapshotter, error) {
+	newRemoteSnapshotterFunc := func(ctx context.Context, namespace string) (*proxy.RemoteSnapshotter, error) {
 		r := resolver
 		response, err := r.Get(namespace)
 		if err != nil {
 			return nil, err
 		}
-		u, err := url.Parse(response.Address)
-		if err != nil {
-			return nil, err
-		}
-		host := u.Hostname()
+		host := response.Address
 		port, err := strconv.ParseUint(response.SnapshotterPort, base10, bits32)
 		if err != nil {
 			return nil, err
@@ -167,57 +178,48 @@ func initSnapshotter(ctx context.Context, config config.Config, cache cache.Cach
 		}
 
 		var metricsProxy *metrics.Proxy
-		// TODO (ginglis13): port management and lifecycle ties in to overall metrics proxy
-		// server lifecycle. tracked here: https://github.com/firecracker-microvm/firecracker-containerd/issues/607
-		portMap := make(map[int]bool)
 		if config.Snapshotter.Metrics.Enable {
-			metricsPort, err := strconv.ParseUint(response.MetricsPort, base10, bits32)
+			metricsProxy, err = initMetricsProxy(config, monitor, host, response.MetricsPort, response.Labels)
 			if err != nil {
 				return nil, err
 			}
+		}
 
-			metricsDialer := func(ctx context.Context, _, _ string) (net.Conn, error) {
-				return vsock.DialContext(ctx, host, uint32(metricsPort), vsock.WithLogger(log.G(ctx)))
-			}
+		return proxy.NewRemoteSnapshotter(ctx, host, snapshotterDialer, metricsProxy)
+	}
 
-			portRange := config.Snapshotter.Metrics.PortRange
-			metricsHost := config.Snapshotter.Metrics.Host
+	return demux.NewSnapshotter(cache, newRemoteSnapshotterFunc), nil
+}
 
-			// Assign a port for metrics proxy server.
-			ports := strings.Split(portRange, "-")
-			portRangeError := fmt.Errorf("invalid port range %s", portRange)
-			if len(ports) < 2 {
-				return nil, portRangeError
-			}
-			lower, err := strconv.Atoi(ports[0])
-			if err != nil {
-				return nil, portRangeError
-			}
-			upper, err := strconv.Atoi(ports[1])
-			if err != nil {
-				return nil, portRangeError
-			}
-			port := -1
-			for p := lower; p <= upper; p++ {
-				if _, ok := portMap[p]; !ok {
-					port = p
-					portMap[p] = true
-					break
-				}
-			}
-			if port < 0 {
-				return nil, fmt.Errorf("invalid port: %d", port)
-			}
+func initMetricsProxyMonitor(portRange string) (*metrics.Monitor, error) {
+	ports := strings.Split(portRange, "-")
+	portRangeError := fmt.Errorf("invalid port range %s", portRange)
+	if len(ports) < 2 {
+		return nil, portRangeError
+	}
+	lower, err := strconv.Atoi(ports[0])
+	if err != nil {
+		return nil, portRangeError
+	}
+	upper, err := strconv.Atoi(ports[1])
+	if err != nil {
+		return nil, portRangeError
+	}
 
-			metricsProxy, err = metrics.NewProxy(metricsHost, port, response.Labels, metricsDialer)
-			if err != nil {
-				return nil, err
-			}
+	return metrics.NewMonitor(lower, upper)
+}
 
-		}
+func initMetricsProxy(config config.Config, monitor *metrics.Monitor, host, port string, labels map[string]string) (*metrics.Proxy, error) {
+	metricsPort, err := strconv.ParseUint(port, base10, bits32)
+	if err != nil {
+		return nil, err
+	}
 
-		return proxy.NewProxySnapshotter(ctx, host, snapshotterDialer, metricsProxy)
+	metricsDialer := func(ctx context.Context, _, _ string) (net.Conn, error) {
+		return vsock.DialContext(ctx, host, uint32(metricsPort), vsock.WithLogger(log.G(ctx)))
 	}
 
-	return demux.NewSnapshotter(cache, newProxySnapshotterFunc), nil
+	metricsHost := config.Snapshotter.Metrics.Host
+
+	return metrics.NewProxy(metricsHost, monitor, labels, metricsDialer)
 }

snapshotter/demux/metrics/proxy.go

@@ -15,10 +15,12 @@ package metrics
 
 import (
 	"context"
+	"fmt"
 	"io"
 	"net"
 	"net/http"
 	"strconv"
+	"sync"
 
 	"github.com/containerd/containerd/log"
 )
@@ -30,10 +32,76 @@ const (
 	MaxMetricsResponseSize = 32768 // 32 KB
 )
 
+// Monitor is used for port and lifecycle management of metrics proxies.
+type Monitor struct {
+	// ch is used for communicating when a metrics proxy and its port are no longer in use
+	// so that the Monitor can mark the port as free.
+	ch chan int
+	// portsInUse is a set that tracks used and unused metrics proxy server ports.
+	portsInUse map[int]bool
+	// startPort is first port in the configurable range of ports to serve metrics proxies from.
+	startPort int
+	// endPort is last port in the configurable range of ports to serve metrics proxies from.
+	endPort int
+	// mu is the lock used for threadsafe reading and writing of the metrics proxy port map.
+	mu *sync.Mutex
+}
+
+// NewMonitor returns a new metrics proxy monitor.
+func NewMonitor(startPort, endPort int) (*Monitor, error) {
+	if startPort <= 0 || endPort <= 0 || startPort > endPort {
+		return nil, fmt.Errorf("invalid port range %d-%d", startPort, endPort)
+	}
+	return &Monitor{make(chan int), make(map[int]bool), startPort, endPort, &sync.Mutex{}}, nil
+}
+
+// Start receives messages from the proxy server channel indiciating a server has closed or terminated.
+// It marks the metrics proxy server's port as available by removing it from the portsInUse set.
+func (m *Monitor) Start() error {
+	for freePort := range m.ch {
+		m.mu.Lock()
+		delete(m.portsInUse, freePort)
+		m.mu.Unlock()
+	}
+
+	return nil
+}
+
+// Stop closes the metrics proxy monitoring channel.
+func (m *Monitor) Stop() {
+	close(m.ch)
+}
+
+// findOpenPort returns a port for use by a metrics proxy server.
+func (m *Monitor) findOpenPort() (int, error) {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	metricsProxyPort := -1
+	for p := m.startPort; p <= m.endPort; p++ {
+		if _, ok := m.portsInUse[p]; !ok {
+			// Mark the port as in use in the set
+			metricsProxyPort = p
+			m.portsInUse[p] = true
+			break
+		}
+	}
+
+	if metricsProxyPort < 0 {
+		return 0, fmt.Errorf("unable to find port in range %d-%d", m.startPort, m.endPort)
+	}
+
+	return metricsProxyPort, nil
+}
+
+// HTTPClient defines the interface for the client getting metrics.
+type HTTPClient interface {
+	Get(string) (*http.Response, error)
+}
+
 // Proxy represents a metrics proxy server for getting and serving remote snapshotter metrics.
 type Proxy struct {
 	// client is the HTTP client used to send metrics requests to a remote snapshotter.
-	client *http.Client
+	client HTTPClient
 	// server is the proxy server on host for serving remote snapshotter metrics.
 	server *http.Server
 	// host is the nonroutable address listening for metrics requests.
@@ -42,13 +110,19 @@ type Proxy struct {
 	Port int
 	// Labels are labels used to apply to metrics.
 	Labels map[string]string
+	// monitor is the metrics proxy monitor
+	monitor *Monitor
 }
 
 // NewProxy creates a new Proxy with initialized HTTP client and port.
 // dialer is used as the DialContext underlying the metrics HTTP client's RoundTripper.
 //
 // Reference: https://pkg.go.dev/net/http#Transport
-func NewProxy(host string, port int, labels map[string]string, dialer func(context.Context, string, string) (net.Conn, error)) (*Proxy, error) {
+func NewProxy(host string, monitor *Monitor, labels map[string]string, dialer func(context.Context, string, string) (net.Conn, error)) (*Proxy, error) {
+	port, err := monitor.findOpenPort()
+	if err != nil {
+		return nil, err
+	}
 	return &Proxy{
 		client: &http.Client{
 			Transport: &http.Transport{
@@ -74,9 +148,12 @@ func (mp *Proxy) Serve(ctx context.Context) error {
 
 	err := mp.server.ListenAndServe()
 	if err != http.ErrServerClosed {
+		log.G(ctx).Errorf("metrics proxy server: %v\n", err)
+		mp.monitor.ch <- mp.Port
 		return err
 	}
 
+	mp.monitor.ch <- mp.Port
 	return nil
 }
 
@@ -92,13 +169,18 @@ func (mp *Proxy) metrics(w http.ResponseWriter, req *http.Request) {
 	// Pull metrics and copy to HTTP response.
 	res, err := mp.client.Get(RequestPath)
 	if err != nil {
-		log.G(ctx).Errorf("error reading response body: %v\n", err)
+		log.G(ctx).Errorf("error reading response body: %v", err)
 		w.WriteHeader(500)
 		return
 	}
 
 	defer res.Body.Close()
 	response := io.LimitReader(res.Body, MaxMetricsResponseSize)
 
-	io.Copy(w, response)
+	_, err = io.Copy(w, response)
+	if err != nil {
+		log.G(ctx).Errorf("error writing response body: %v", err)
+		w.WriteHeader(500)
+		return
+	}
 }