github.com/m-lab/locate@v0.17.6/handler/heartbeat.go (about)

     1  package handler
     2  
     3  import (
     4  	"context"
     5  	"encoding/json"
     6  	"net/http"
     7  	"time"
     8  
     9  	"github.com/gorilla/websocket"
    10  	v2 "github.com/m-lab/locate/api/v2"
    11  	"github.com/m-lab/locate/metrics"
    12  	"github.com/m-lab/locate/static"
    13  	log "github.com/sirupsen/logrus"
    14  )
    15  
    16  var readDeadline = static.WebsocketReadDeadline
    17  
    18  type conn interface {
    19  	ReadMessage() (int, []byte, error)
    20  	SetReadDeadline(time.Time) error
    21  	Close() error
    22  }
    23  
    24  // Heartbeat implements /v2/heartbeat requests.
    25  // It starts a new persistent connection and a new goroutine
    26  // to read incoming messages.
    27  func (c *Client) Heartbeat(rw http.ResponseWriter, req *http.Request) {
    28  	upgrader := websocket.Upgrader{
    29  		ReadBufferSize:  static.WebsocketBufferSize,
    30  		WriteBufferSize: static.WebsocketBufferSize,
    31  	}
    32  	ws, err := upgrader.Upgrade(rw, req, nil)
    33  	if err != nil {
    34  		log.Errorf("failed to establish a connection: %v", err)
    35  		metrics.RequestsTotal.WithLabelValues("heartbeat", "establish connection",
    36  			"error upgrading the HTTP server connection to the WebSocket protocol").Inc()
    37  		return
    38  	}
    39  	metrics.RequestsTotal.WithLabelValues("heartbeat", "establish connection", "OK").Inc()
    40  	go c.handleHeartbeats(ws)
    41  }
    42  
    43  // handleHeartbeats handles incoming messages from the connection.
    44  func (c *Client) handleHeartbeats(ws conn) error {
    45  	defer ws.Close()
    46  	setReadDeadline(ws)
    47  
    48  	var hostname string
    49  	var experiment string
    50  	for {
    51  		_, message, err := ws.ReadMessage()
    52  		if err != nil {
    53  			closeConnection(experiment, err)
    54  			return err
    55  		}
    56  		if message != nil {
    57  			setReadDeadline(ws)
    58  
    59  			var hbm v2.HeartbeatMessage
    60  			if err := json.Unmarshal(message, &hbm); err != nil {
    61  				log.Errorf("failed to unmarshal heartbeat message, err: %v", err)
    62  				continue
    63  			}
    64  
    65  			switch {
    66  			case hbm.Registration != nil:
    67  				if err := c.RegisterInstance(*hbm.Registration); err != nil {
    68  					closeConnection(experiment, err)
    69  					return err
    70  				}
    71  
    72  				if hostname == "" {
    73  					hostname = hbm.Registration.Hostname
    74  					experiment = hbm.Registration.Experiment
    75  					metrics.CurrentHeartbeatConnections.WithLabelValues(experiment).Inc()
    76  				}
    77  
    78  				// Update Prometheus signals every time a Registration message is received.
    79  				c.UpdatePrometheusForMachine(context.Background(), hbm.Registration.Hostname)
    80  			case hbm.Health != nil:
    81  				if err := c.UpdateHealth(hostname, *hbm.Health); err != nil {
    82  					closeConnection(experiment, err)
    83  					return err
    84  				}
    85  			}
    86  		}
    87  	}
    88  }
    89  
    90  // setReadDeadline sets/resets the read deadline for the connection.
    91  func setReadDeadline(ws conn) {
    92  	deadline := time.Now().Add(readDeadline)
    93  	ws.SetReadDeadline(deadline)
    94  }
    95  
    96  func closeConnection(experiment string, err error) {
    97  	if experiment != "" {
    98  		metrics.CurrentHeartbeatConnections.WithLabelValues(experiment).Dec()
    99  	}
   100  	log.Errorf("closing connection, err: %v", err)
   101  }