github.com/m-lab/locate@v0.17.6/cmd/heartbeat/health/port-probe.go (about)

     1  package health
     2  
     3  import (
     4  	"net"
     5  	"net/url"
     6  	"time"
     7  
     8  	"github.com/m-lab/locate/metrics"
     9  )
    10  
    11  const (
    12  	defaultPort       = "80"
    13  	defaultPortSecure = "443"
    14  )
    15  
    16  // PortProbe checks whether a set of ports are open.
    17  type PortProbe struct {
    18  	ports map[string]bool
    19  }
    20  
    21  // NewPortProbe creates a new PortProbe.
    22  func NewPortProbe(services map[string][]string) *PortProbe {
    23  	pp := PortProbe{
    24  		ports: getPorts(services),
    25  	}
    26  	return &pp
    27  }
    28  
    29  // checkPorts returns true if all the given ports are open and false
    30  // otherwise.
    31  func (ps *PortProbe) checkPorts() bool {
    32  	for p := range ps.ports {
    33  		conn, err := net.DialTimeout("tcp", "localhost:"+p, time.Second)
    34  		if err != nil {
    35  			metrics.PortChecksTotal.WithLabelValues(err.Error()).Inc()
    36  			return false
    37  		}
    38  
    39  		conn.Close()
    40  		metrics.PortChecksTotal.WithLabelValues("OK").Inc()
    41  	}
    42  	return true
    43  }
    44  
    45  // getPorts extracts the set of ports from a map of service names to
    46  // their URL templates.
    47  func getPorts(services map[string][]string) map[string]bool {
    48  	ports := make(map[string]bool)
    49  
    50  	for _, s := range services {
    51  		for _, u := range s {
    52  			url, err := url.Parse(u)
    53  
    54  			if err != nil {
    55  				continue
    56  			}
    57  
    58  			port := getPort(*url)
    59  			ports[port] = true
    60  		}
    61  	}
    62  
    63  	return ports
    64  }
    65  
    66  // getPort extracts the port from a single URL. If no port is specified,
    67  // it sets a default.
    68  func getPort(url url.URL) string {
    69  	port := url.Port()
    70  
    71  	// Set default ports.
    72  	if port == "" {
    73  		if url.Scheme == "https" || url.Scheme == "wss" {
    74  			return defaultPortSecure
    75  		}
    76  		return defaultPort
    77  	}
    78  
    79  	return port
    80  }