server.go

Documentation: edge-infra.dev/pkg/lib/server

     1  // Package server provides a generic http server and metrics server with overridable
     2  // implementations for Kubernetes probes
     3  package server
     4  
     5  import (
     6  	"context"
     7  	"fmt"
     8  	"net"
     9  	"net/http"
    10  	"sync"
    11  	"time"
    12  
    13  	"github.com/gin-gonic/gin"
    14  	"github.com/prometheus/client_golang/prometheus/promhttp"
    15  )
    16  
    17  func init() {
    18  	gin.SetMode(gin.ReleaseMode)
    19  	gin.DisableConsoleColor()
    20  }
    21  
    22  // Handler is a function that can handle a request from a gin server.
    23  type Handler func(c *gin.Context)
    24  
    25  // Config contains unexported configuration parameters for a server.
    26  // The fields are set using Option functions.
    27  type Config struct {
    28  	port     int
    29  	listener net.Listener
    30  
    31  	handlers map[string]Handler
    32  
    33  	secure         bool
    34  	secureCertFile string
    35  	secureKeyFile  string
    36  }
    37  
    38  // Validate checks that the config does not contain conflicting options.
    39  func (c Config) Validate() error {
    40  	if c.port != 0 && c.listener != nil {
    41  		return fmt.Errorf("OptionPort and OptionListener cannot both be set")
    42  	}
    43  
    44  	if c.secure && c.listener != nil {
    45  		return fmt.Errorf("OptionSecure can not be set with OptionListener")
    46  	}
    47  
    48  	return nil
    49  }
    50  
    51  // Addr creates the http.Server's Addr.
    52  func (c Config) Addr() string {
    53  	if c.port == 0 {
    54  		if c.secure {
    55  			return ":443"
    56  		}
    57  		return ":80"
    58  	}
    59  	return fmt.Sprintf(":%d", c.port)
    60  }
    61  
    62  // Server represents the HTTP server.
    63  type Server struct {
    64  	Router *gin.Engine
    65  
    66  	Config Config
    67  
    68  	// shutdownFunc is set to an http.Server.Shutdown function.
    69  	shutdownFunc func(context.Context) error
    70  }
    71  
    72  // NewServer creates an http server that runs with the desired options.
    73  //
    74  // To include a prometheus "/metrics" endpoint, provide the OptionMetrics option.
    75  // It uses the "github.com/prometheus/client_golang/prometheus/promhttp" package to serve your scrapable prometheus metrics.
    76  //
    77  // If the OptionPort and OptionListener options are omitted, then the default addr ":80" or ":443" will be used,
    78  // depending on the secure option.
    79  func NewServer(options ...Option) (*Server, error) {
    80  	var s Server
    81  	s.Router = gin.New()
    82  
    83  	for _, option := range options {
    84  		if err := option(&s.Config); err != nil {
    85  			return nil, err
    86  		}
    87  	}
    88  	if err := s.Config.Validate(); err != nil {
    89  		return nil, err
    90  	}
    91  
    92  	for path, h := range s.Config.handlers {
    93  		s.Router.GET(path, gin.HandlerFunc(h))
    94  	}
    95  
    96  	return &s, nil
    97  }
    98  
    99  // NewHealthServer creates a server that serves 200 OK on the endpoint "/healthz"
   100  func NewHealthServer(opts ...Option) (*Server, error) {
   101  	var h = func(c *gin.Context) {
   102  		c.JSON(http.StatusOK, "UP")
   103  	}
   104  
   105  	// copy for no side affects on passed in slice.
   106  	var alloptscopy = []Option{
   107  		OptionHandler("/healthz", h),
   108  	}
   109  	alloptscopy = append(alloptscopy, opts...)
   110  
   111  	return NewServer(alloptscopy...)
   112  }
   113  
   114  // NewReadyServer serves http.StatusServiceUnavailable on "/readyz" until the returned `ready` function is called, then it serves 200 OK.
   115  func NewReadyServer(opts ...Option) (s *Server, ready func(), err error) {
   116  	var ch = make(chan bool)
   117  	ready = func() {
   118  		close(ch)
   119  	}
   120  	var h = func(c *gin.Context) {
   121  		select {
   122  		case <-ch:
   123  			// when the channel is closed, then this case will always run.
   124  			c.JSON(http.StatusOK, "READY")
   125  		default:
   126  			c.JSON(http.StatusServiceUnavailable, "NOT READY")
   127  		}
   128  	}
   129  	// copy for no side affects on passed in slice.
   130  	var alloptscopy = []Option{
   131  		OptionHandler("/readyz", h),
   132  	}
   133  	alloptscopy = append(alloptscopy, opts...)
   134  
   135  	s, err = NewServer(alloptscopy...)
   136  	if err != nil {
   137  		close(ch)
   138  		return nil, nil, err
   139  	}
   140  	return s, ready, err
   141  }
   142  
   143  var PrometheusMetricsHandler = gin.WrapH(promhttp.Handler())
   144  
   145  // NewMetricsServer serves prometheus metrics on the "/metrics" endpoint.
   146  //
   147  // To add your own metrics, use the prometheus client_golang libraries as normal, and prometheus will wire them up for you.
   148  func NewMetricsServer(opts ...Option) (*Server, error) {
   149  	var oh = OptionHandler("/metrics", PrometheusMetricsHandler)
   150  	s, ready, err := NewReadyServer(append([]Option{oh}, opts...)...)
   151  	if err != nil {
   152  		return nil, err
   153  	}
   154  	ready()
   155  	return s, err
   156  }
   157  
   158  // NewLivenessServer serves http.StatusServiceUnavailable on "/livez" until the `setLivenessStatus` function is called.
   159  // After calling the `setLivenessStatus` function, it serves whatever status code you pass in.
   160  // The `setLivenessStatus` function can be called more than once.
   161  func NewLivenessServer(opts ...Option) (s *Server, setLivenessStatus func(code int), err error) {
   162  	// liveValue holds the livez status in a threadsafe manner.
   163  	var liveValue struct {
   164  		sync.Mutex
   165  		code int
   166  	}
   167  	// this func safely sets the liveValue fields after locking.
   168  	var set = func(code int) {
   169  		liveValue.Lock()
   170  		liveValue.code = code
   171  		liveValue.Unlock()
   172  	}
   173  	// this handler safely returns the liveValue values.
   174  	var h = func(c *gin.Context) {
   175  		liveValue.Lock()
   176  		var code = liveValue.code
   177  		liveValue.Unlock()
   178  
   179  		// Default to status http.StatusServiceUnavailable
   180  		if code == 0 {
   181  			code = http.StatusServiceUnavailable
   182  		}
   183  
   184  		if code >= 400 {
   185  			c.JSON(code, "DOWN")
   186  		} else {
   187  			c.JSON(code, "LIVE")
   188  		}
   189  	}
   190  
   191  	// copy for no side affects on passed in slice.
   192  	var alloptscopy = []Option{
   193  		OptionHandler("/livez", h),
   194  	}
   195  	alloptscopy = append(alloptscopy, opts...)
   196  
   197  	s, err = NewServer(alloptscopy...)
   198  	if err != nil {
   199  		return nil, nil, err
   200  	}
   201  	return s, set, nil
   202  }
   203  
   204  // ReadHeaderTimeout needs to be set to preven a "slowloris attack" according to the linter gosec.
   205  //
   206  // This is a public var so that consumers of this library can modify it before calling Run.
   207  var ReadHeaderTimeout = time.Minute
   208  
   209  // Run runs the Gin Router on the configured port or listener.
   210  func (s *Server) Run() error {
   211  	// Serve on a net/http.Server at the desired location.
   212  	var hs = &http.Server{
   213  		Handler:           s.Router,
   214  		Addr:              s.Config.Addr(),
   215  		ReadHeaderTimeout: ReadHeaderTimeout,
   216  	}
   217  	s.shutdownFunc = hs.Shutdown // Allows the gin server be shut down gracefully.
   218  
   219  	if s.Config.listener != nil {
   220  		// Serve on a custom listener
   221  		return hs.Serve(s.Config.listener)
   222  	} else if s.Config.secure {
   223  		// Serve TLS on a tcp port.
   224  		return hs.ListenAndServeTLS(s.Config.secureCertFile, s.Config.secureKeyFile)
   225  	}
   226  	// Serve insecure HTTP 1.1 on a tcp port.
   227  	return hs.ListenAndServe()
   228  }
   229  
   230  // Shutdown stops the running server. See golang.org/pkg/net/http#Server.Shutdown
   231  func (s *Server) Shutdown(ctx context.Context) error {
   232  	if s.shutdownFunc == nil {
   233  		return nil
   234  	}
   235  	return s.shutdownFunc(ctx)
   236  }
   237
View as plain text