collector: reimplement sockstat collector with procfs (#1552)

* collector: reimplement sockstat collector with procfs
* collector: handle sockstat IPv4 disabled, debug logging

Signed-off-by: Matt Layher <mdlayher@gmail.com>

1 files changed, 121 insertions, 63 deletions

diff --git a/collector/sockstat_linux.go b/collector/sockstat_linux.go
index 92847be..d6c4c72 100644
--- a/collector/sockstat_linux.go
+++ b/collector/sockstat_linux.go
@@ -16,14 +16,12 @@
 package collector
 
 import (
-        "bufio"
         "fmt"
-        "io"
         "os"
-        "strconv"
-        "strings"
 
         "github.com/prometheus/client_golang/prometheus"
+        "github.com/prometheus/common/log"
+        "github.com/prometheus/procfs"
 )
 
 const (
@@ -45,78 +43,138 @@ func NewSockStatCollector() (Collector, error) {
 }
 
 func (c *sockStatCollector) Update(ch chan<- prometheus.Metric) error {
-        sockStats, err := getSockStats(procFilePath("net/sockstat"))
+        fs, err := procfs.NewFS(*procPath)
         if err != nil {
-                return fmt.Errorf("couldn't get sockstats: %s", err)
+                return fmt.Errorf("failed to open procfs: %v", err)
         }
-        for protocol, protocolStats := range sockStats {
-                for name, value := range protocolStats {
-                        v, err := strconv.ParseFloat(value, 64)
-                        if err != nil {
-                                return fmt.Errorf("invalid value %s in sockstats: %s", value, err)
-                        }
-                        ch <- prometheus.MustNewConstMetric(
-                                prometheus.NewDesc(
-                                        prometheus.BuildFQName(namespace, sockStatSubsystem, protocol+"_"+name),
-                                        fmt.Sprintf("Number of %s sockets in state %s.", protocol, name),
-                                        nil, nil,
-                                ),
-                                prometheus.GaugeValue, v,
-                        )
-                }
+
+        // If IPv4 and/or IPv6 are disabled on this kernel, handle it gracefully.
+        stat4, err := fs.NetSockstat()
+        switch {
+        case err == nil:
+        case os.IsNotExist(err):
+                log.Debug("IPv4 sockstat statistics not found, skipping")
+        default:
+                return fmt.Errorf("failed to get IPv4 sockstat data: %v", err)
         }
-        return err
-}
 
-func getSockStats(fileName string) (map[string]map[string]string, error) {
-        file, err := os.Open(fileName)
-        if err != nil {
-                return nil, err
+        stat6, err := fs.NetSockstat6()
+        switch {
+        case err == nil:
+        case os.IsNotExist(err):
+                log.Debug("IPv6 sockstat statistics not found, skipping")
+        default:
+                return fmt.Errorf("failed to get IPv6 sockstat data: %v", err)
+        }
+
+        stats := []struct {
+                isIPv6 bool
+                stat   *procfs.NetSockstat
+        }{
+                {
+                        stat: stat4,
+                },
+                {
+                        isIPv6: true,
+                        stat:   stat6,
+                },
         }
-        defer file.Close()
 
-        return parseSockStats(file, fileName)
+        for _, s := range stats {
+                c.update(ch, s.isIPv6, s.stat)
+        }
+
+        return nil
 }
 
-func parseSockStats(r io.Reader, fileName string) (map[string]map[string]string, error) {
-        var (
-                sockStat = map[string]map[string]string{}
-                scanner  = bufio.NewScanner(r)
-        )
-
-        for scanner.Scan() {
-                line := strings.Split(scanner.Text(), " ")
-                // Remove trailing ':'.
-                protocol := line[0][:len(line[0])-1]
-                sockStat[protocol] = map[string]string{}
-
-                for i := 1; i < len(line) && i+1 < len(line); i++ {
-                        sockStat[protocol][line[i]] = line[i+1]
-                        i++
-                }
+func (c *sockStatCollector) update(ch chan<- prometheus.Metric, isIPv6 bool, s *procfs.NetSockstat) {
+        if s == nil {
+                // IPv6 disabled or similar; nothing to do.
+                return
         }
-        if err := scanner.Err(); err != nil {
-                return nil, err
+
+        // If sockstat contains the number of used sockets, export it.
+        if !isIPv6 && s.Used != nil {
+                // TODO: this must be updated if sockstat6 ever exports this data.
+                ch <- prometheus.MustNewConstMetric(
+                        prometheus.NewDesc(
+                                prometheus.BuildFQName(namespace, sockStatSubsystem, "sockets_used"),
+                                "Number of IPv4 sockets in use.",
+                                nil,
+                                nil,
+                        ),
+                        prometheus.GaugeValue,
+                        float64(*s.Used),
+                )
         }
 
-        // The mem metrics is the count of pages used. Multiply the mem metrics by
-        // the page size from the kernel to get the number of bytes used.
-        //
-        // Update the TCP mem from page count to bytes.
-        pageCount, err := strconv.Atoi(sockStat["TCP"]["mem"])
-        if err != nil {
-                return nil, fmt.Errorf("invalid value %s in sockstats: %s", sockStat["TCP"]["mem"], err)
+        // A name and optional value for a sockstat metric.
+        type ssPair struct {
+                name string
+                v    *int
         }
-        sockStat["TCP"]["mem_bytes"] = strconv.Itoa(pageCount * pageSize)
 
-        // Update the UDP mem from page count to bytes.
-        if udpMem := sockStat["UDP"]["mem"]; udpMem != "" {
-                pageCount, err = strconv.Atoi(udpMem)
-                if err != nil {
-                        return nil, fmt.Errorf("invalid value %s in sockstats: %s", sockStat["UDP"]["mem"], err)
+        // Previously these metric names were generated directly from the file output.
+        // In order to keep the same level of compatibility, we must map the fields
+        // to their correct names.
+        for _, p := range s.Protocols {
+                pairs := []ssPair{
+                        {
+                                name: "inuse",
+                                v:    &p.InUse,
+                        },
+                        {
+                                name: "orphan",
+                                v:    p.Orphan,
+                        },
+                        {
+                                name: "tw",
+                                v:    p.TW,
+                        },
+                        {
+                                name: "alloc",
+                                v:    p.Alloc,
+                        },
+                        {
+                                name: "mem",
+                                v:    p.Mem,
+                        },
+                        {
+                                name: "memory",
+                                v:    p.Memory,
+                        },
                 }
-                sockStat["UDP"]["mem_bytes"] = strconv.Itoa(pageCount * pageSize)
-        }
 
-        return sockStat, nil
+                // Also export mem_bytes values for sockets which have a mem value
+                // stored in pages.
+                if p.Mem != nil {
+                        v := *p.Mem * pageSize
+                        pairs = append(pairs, ssPair{
+                                name: "mem_bytes",
+                                v:    &v,
+                        })
+                }
+
+                for _, pair := range pairs {
+                        if pair.v == nil {
+                                // This value is not set for this protocol; nothing to do.
+                                continue
+                        }
+
+                        ch <- prometheus.MustNewConstMetric(
+                                prometheus.NewDesc(
+                                        prometheus.BuildFQName(
+                                                namespace,
+                                                sockStatSubsystem,
+                                                fmt.Sprintf("%s_%s", p.Protocol, pair.name),
+                                        ),
+                                        fmt.Sprintf("Number of %s sockets in state %s.", p.Protocol, pair.name),
+                                        nil,
+                                        nil,
+                                ),
+                                prometheus.GaugeValue,
+                                float64(*pair.v),
+                        )
+                }
+        }
 }