github.com/lastbackend/toolkit@v0.0.0-20241020043710-cafa37b95aad/pkg/tools/probes/server/server.go (about)

     1  package server
     2  
     3  import (
     4  	"context"
     5  	"encoding/json"
     6  	"fmt"
     7  	"github.com/lastbackend/toolkit/pkg/runtime"
     8  	"github.com/lastbackend/toolkit/pkg/server"
     9  	"github.com/lastbackend/toolkit/pkg/tools/probes"
    10  	"net/http"
    11  	"sync"
    12  )
    13  
    14  const prefix = "probes"
    15  
    16  const (
    17  	defaultProbesHttpServerName string = "probes"
    18  	defaultContentType                 = "application/json; charset=utf-8"
    19  )
    20  
    21  type Options struct {
    22  	Enabled bool   `env:"SERVER_ENABLED" envDefault:"true" comment:"Enable or disable probes server"`
    23  	Host    string `env:"SERVER_LISTEN" envDefault:"0.0.0.0" comment:"Set probes listen host"`
    24  	Port    int    `env:"SERVER_PORT" envDefault:"8080" comment:"Set probes listen port"`
    25  
    26  	LivenessPath  string `env:"LIVENESS_PATH" envDefault:"/_healthz/liveness" comment:"Set liveness probe path"`
    27  	ReadinessPath string `env:"READINESS_PATH" envDefault:"/_healthz/readiness" comment:"Set readiness probe path"`
    28  }
    29  
    30  type probe struct {
    31  	mtx     sync.RWMutex
    32  	runtime runtime.Runtime
    33  
    34  	opts            Options
    35  	readinessProbes map[string]probes.HandleFunc
    36  	livenessProbes  map[string]probes.HandleFunc
    37  }
    38  
    39  func NewProbesServer(runtime runtime.Runtime) (probes.Probes, error) {
    40  	srv := new(probe)
    41  
    42  	srv.runtime = runtime
    43  	srv.opts = Options{}
    44  
    45  	srv.readinessProbes = make(map[string]probes.HandleFunc, 0)
    46  	srv.livenessProbes = make(map[string]probes.HandleFunc, 0)
    47  
    48  	return srv, runtime.Config().Parse(&srv.opts, prefix)
    49  }
    50  
    51  func (p *probe) RegisterCheck(name string, kind probes.ProbeKind, fn probes.HandleFunc) error {
    52  
    53  	p.mtx.Lock()
    54  	defer p.mtx.Unlock()
    55  
    56  	switch kind {
    57  	case probes.LivenessProbe:
    58  		if _, ok := p.livenessProbes[name]; ok {
    59  			return fmt.Errorf("trying to override liveness probe check: %s", name)
    60  		}
    61  		p.livenessProbes[name] = fn
    62  	case probes.ReadinessProbe:
    63  		if _, ok := p.readinessProbes[name]; ok {
    64  			return fmt.Errorf("trying to override liveness probe check: %s", name)
    65  		}
    66  		p.readinessProbes[name] = fn
    67  	}
    68  
    69  	return nil
    70  }
    71  
    72  func (p *probe) livenessProbeHandler(w http.ResponseWriter, r *http.Request) {
    73  	p.probeHandler("liveness_probe", w, r, p.livenessProbes)
    74  }
    75  
    76  func (p *probe) readinessProbeHandler(w http.ResponseWriter, r *http.Request) {
    77  	p.probeHandler("readiness_probe", w, r, p.readinessProbes)
    78  }
    79  
    80  func (p *probe) probeHandler(probeType string, w http.ResponseWriter, _ *http.Request, probes map[string]probes.HandleFunc) {
    81  
    82  	var (
    83  		result = make(map[string]string, 0)
    84  		status = http.StatusOK
    85  	)
    86  
    87  	for name, probeFunc := range probes {
    88  		fn := probeFunc
    89  
    90  		if err := fn(); err != nil {
    91  			status = http.StatusInternalServerError
    92  			result[name] = err.Error()
    93  
    94  			p.runtime.Log().Errorf("[%s][%s] Probe failed: %v", probeType, name, err)
    95  
    96  		} else {
    97  			result[name] = "OK"
    98  		}
    99  	}
   100  
   101  	w.WriteHeader(status)
   102  	w.Header().Set("Content-Type", defaultContentType)
   103  
   104  	encoder := json.NewEncoder(w)
   105  	encoder.SetIndent("", "    ")
   106  
   107  	if err := encoder.Encode(result); err != nil {
   108  		return
   109  	}
   110  }
   111  
   112  func (p *probe) Start(_ context.Context) error {
   113  
   114  	var (
   115  		s server.HTTPServer
   116  	)
   117  
   118  	if !p.opts.Enabled {
   119  		return nil
   120  	}
   121  
   122  	if p.opts.Port > 0 {
   123  		// check if provided port is used by grpc server
   124  		grpcServers := p.runtime.Server().GRPCList()
   125  		for _, srv := range grpcServers {
   126  			if srv.Info().Port == p.opts.Port {
   127  				return fmt.Errorf("can not bind probes handlers to grpc server. please change porbes port to different port")
   128  			}
   129  		}
   130  
   131  		// check if provided port is used on provided http server
   132  		httpServers := p.runtime.Server().HTTPList()
   133  		for _, srv := range httpServers {
   134  			if srv.Info().Port == p.opts.Port {
   135  				s = srv
   136  				break
   137  			}
   138  		}
   139  	}
   140  
   141  	if s == nil {
   142  		s = p.runtime.Server().HTTPNew(defaultProbesHttpServerName, &server.HTTPServerOptions{
   143  			Host:      p.opts.Host,
   144  			Port:      p.opts.Port,
   145  			TLSConfig: nil,
   146  		})
   147  	}
   148  
   149  	s.AddHandler(http.MethodGet, p.opts.LivenessPath, p.livenessProbeHandler)
   150  	s.AddHandler(http.MethodGet, p.opts.ReadinessPath, p.readinessProbeHandler)
   151  
   152  	return nil
   153  }