readiness.go 1.9 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980
  1. package metrics
  2. import (
  3. "encoding/json"
  4. "fmt"
  5. "net/http"
  6. "sync"
  7. conn "github.com/cloudflare/cloudflared/connection"
  8. "github.com/rs/zerolog"
  9. )
  10. // ReadyServer serves HTTP 200 if the tunnel can serve traffic. Intended for k8s readiness checks.
  11. type ReadyServer struct {
  12. sync.RWMutex
  13. isConnected map[int]bool
  14. log *zerolog.Logger
  15. }
  16. // NewReadyServer initializes a ReadyServer and starts listening for dis/connection events.
  17. func NewReadyServer(log *zerolog.Logger) *ReadyServer {
  18. return &ReadyServer{
  19. isConnected: make(map[int]bool, 0),
  20. log: log,
  21. }
  22. }
  23. func (rs *ReadyServer) OnTunnelEvent(c conn.Event) {
  24. switch c.EventType {
  25. case conn.Connected:
  26. rs.Lock()
  27. rs.isConnected[int(c.Index)] = true
  28. rs.Unlock()
  29. case conn.Disconnected, conn.Reconnecting, conn.RegisteringTunnel, conn.Unregistering:
  30. rs.Lock()
  31. rs.isConnected[int(c.Index)] = false
  32. rs.Unlock()
  33. case conn.SetURL:
  34. break
  35. default:
  36. rs.log.Error().Msgf("Unknown connection event case %v", c)
  37. }
  38. }
  39. type body struct {
  40. Status int `json:"status"`
  41. ReadyConnections int `json:"readyConnections"`
  42. }
  43. // ServeHTTP responds with HTTP 200 if the tunnel is connected to the edge.
  44. func (rs *ReadyServer) ServeHTTP(w http.ResponseWriter, r *http.Request) {
  45. statusCode, readyConnections := rs.makeResponse()
  46. w.WriteHeader(statusCode)
  47. body := body{
  48. Status: statusCode,
  49. ReadyConnections: readyConnections,
  50. }
  51. msg, err := json.Marshal(body)
  52. if err != nil {
  53. _, _ = fmt.Fprintf(w, `{"error": "%s"}`, err)
  54. }
  55. _, _ = w.Write(msg)
  56. }
  57. // This is the bulk of the logic for ServeHTTP, broken into its own pure function
  58. // to make unit testing easy.
  59. func (rs *ReadyServer) makeResponse() (statusCode, readyConnections int) {
  60. statusCode = http.StatusServiceUnavailable
  61. rs.RLock()
  62. defer rs.RUnlock()
  63. for _, connected := range rs.isConnected {
  64. if connected {
  65. statusCode = http.StatusOK
  66. readyConnections++
  67. }
  68. }
  69. return statusCode, readyConnections
  70. }