edgedb-challenge/main.go

89 lines
1.7 KiB
Go

package main
import (
"context"
"flag"
"fmt"
"os"
"github.com/nberlee/go-netstat/netstat"
"github.com/prometheus/client_golang/prometheus"
"github.com/prometheus/common/expfmt"
)
var (
maybe_port = flag.Uint("port", 8500, "Port to monitor")
)
func main() {
flag.Parse()
// Check port number before attempting to cast it to a uint16
if *maybe_port < 1 || *maybe_port > 65535 {
fmt.Fprintln(os.Stderr, "Invalid port number")
os.Exit(1)
}
port := uint16(*maybe_port)
reg := prometheus.NewRegistry()
open_consul_connections := prometheus.NewGauge(prometheus.GaugeOpts{
Subsystem: "consul",
Name: "open_tcp_connections",
Help: "Number of open consul tcp connections",
ConstLabels: prometheus.Labels{
"port": fmt.Sprintf("%d", port),
},
})
reg.MustRegister(open_consul_connections)
open_consul_connections.Set(float64(ConsulConnectionCount(port)))
metricFamilies, err := reg.Gather()
if err != nil {
panic(err)
}
for _, mf := range metricFamilies {
expfmt.MetricFamilyToText(os.Stdout, mf)
}
}
// Filters
func PortFilter(port uint16) func(s *netstat.SockTabEntry) bool {
return func(s *netstat.SockTabEntry) bool {
return s.LocalEndpoint.Port == port
}
}
func StateFilter(s *netstat.SockTabEntry) bool {
return s.State == netstat.Established
}
func ConsulConnectionCount(port uint16) int {
ctx := context.Background()
features := netstat.EnableFeatures{
TCP: true,
TCP6: true,
UDP: true,
UDP6: true,
PID: true,
}
targetFilterFn := PortFilter(port)
resp, err := netstat.Netstat(ctx, features, func(ste *netstat.SockTabEntry) bool {
return targetFilterFn(ste) && StateFilter(ste)
})
if err != nil {
fmt.Fprintf(os.Stderr, "Error: %v\n", err)
return 0
}
return len(resp)
}