github.com/iotexproject/iotex-core@v1.14.1-rc1/pkg/probe/probe.go (about)

     1  // Copyright (c) 2019 IoTeX Foundation
     2  // This source code is provided 'as is' and no warranties are given as to title or non-infringement, merchantability
     3  // or fitness for purpose and, to the extent permitted by law, all liability for your use of the code is disclaimed.
     4  // This source code is governed by Apache License 2.0 that can be found in the LICENSE file.
     5  
     6  package probe
     7  
     8  import (
     9  	"context"
    10  	"fmt"
    11  	"net/http"
    12  
    13  	"github.com/prometheus/client_golang/prometheus/promhttp"
    14  	"go.uber.org/zap"
    15  
    16  	"github.com/iotexproject/iotex-core/pkg/lifecycle"
    17  	"github.com/iotexproject/iotex-core/pkg/log"
    18  	"github.com/iotexproject/iotex-core/pkg/util/httputil"
    19  )
    20  
    21  // Server is a http server for service probe.
    22  type Server struct {
    23  	lifecycle.Readiness
    24  	server           http.Server
    25  	readinessHandler http.Handler
    26  }
    27  
    28  // Option is ued to set probe server's options.
    29  type Option interface {
    30  	SetOption(*Server)
    31  }
    32  
    33  // New creates a new probe server.
    34  func New(port int, opts ...Option) *Server {
    35  	s := &Server{
    36  		readinessHandler: http.HandlerFunc(successHandleFunc),
    37  	}
    38  
    39  	for _, opt := range opts {
    40  		opt.SetOption(s)
    41  	}
    42  
    43  	mux := http.NewServeMux()
    44  	mux.HandleFunc("/liveness", successHandleFunc)
    45  	readiness := func(w http.ResponseWriter, r *http.Request) {
    46  		if !s.IsReady() {
    47  			failureHandleFunc(w, r)
    48  			return
    49  		}
    50  		s.readinessHandler.ServeHTTP(w, r)
    51  	}
    52  
    53  	mux.HandleFunc("/readiness", readiness)
    54  	mux.HandleFunc("/health", readiness)
    55  	mux.Handle("/metrics", promhttp.Handler())
    56  
    57  	s.server = httputil.NewServer(fmt.Sprintf(":%d", port), mux)
    58  	return s
    59  }
    60  
    61  // Start starts the probe server and starts returning success status on liveness endpoint.
    62  func (s *Server) Start(_ context.Context) error {
    63  	go func() {
    64  		ln, err := httputil.LimitListener(s.server.Addr)
    65  		if err != nil {
    66  			log.L().Error("Failed to listen on probe port", zap.Error(err))
    67  			return
    68  		}
    69  		if err := s.server.Serve(ln); err != nil {
    70  			log.L().Info("Probe server stopped.", zap.Error(err))
    71  		}
    72  	}()
    73  	return nil
    74  }
    75  
    76  // Stop shutdown the probe server.
    77  func (s *Server) Stop(ctx context.Context) error { return s.server.Shutdown(ctx) }
    78  
    79  func successHandleFunc(w http.ResponseWriter, _ *http.Request) {
    80  	w.WriteHeader(http.StatusOK)
    81  	if _, err := w.Write([]byte("OK")); err != nil {
    82  		log.L().Warn("Failed to send http response.", zap.Error(err))
    83  	}
    84  }
    85  
    86  func failureHandleFunc(w http.ResponseWriter, _ *http.Request) {
    87  	w.WriteHeader(http.StatusServiceUnavailable)
    88  	if _, err := w.Write([]byte("FAIL")); err != nil {
    89  		log.L().Warn("Failed to send http response.", zap.Error(err))
    90  	}
    91  }