github.com/diadata-org/diadata@v1.4.593/pkg/utils/probes/probe.go (about)

     1  package probes
     2  
     3  import (
     4  	"errors"
     5  	"net/http"
     6  
     7  	"github.com/diadata-org/diadata/pkg/http/restApi"
     8  	"github.com/diadata-org/diadata/pkg/utils"
     9  	"github.com/gin-gonic/gin"
    10  	log "github.com/sirupsen/logrus"
    11  )
    12  
    13  type probe func() bool
    14  
    15  var livenessProbe probe
    16  var readinessProbe probe
    17  
    18  func Start(liveness probe, readiness probe) {
    19  
    20  	log.Info("Ready and Live probes loading")
    21  	livenessProbe = liveness
    22  	readinessProbe = readiness
    23  
    24  	engine := gin.Default()
    25  	engine.Use(gin.Logger())
    26  	engine.Use(gin.Recovery())
    27  
    28  	engine.GET("/ready", execReadiness)
    29  	engine.GET("/live", execLiveness)
    30  	// This environment variable is either set in docker-compose or empty
    31  	go func() {
    32  		err := engine.Run(utils.Getenv("LISTEN_PORT_PROBES", ":2345"))
    33  		if err != nil {
    34  			log.Error(err)
    35  		}
    36  	}()
    37  	log.Info("Ready and Live probes starting")
    38  }
    39  
    40  func execReadiness(context *gin.Context) {
    41  	executeProbe(context, readinessProbe, http.StatusTooEarly)
    42  }
    43  
    44  func execLiveness(context *gin.Context) {
    45  	executeProbe(context, livenessProbe, http.StatusInternalServerError)
    46  }
    47  
    48  func executeProbe(context *gin.Context, fn probe, errorCode int) bool {
    49  	if fn() {
    50  		context.JSON(http.StatusOK, gin.H{"message": "success"})
    51  		return true
    52  	}
    53  	restApi.SendError(context, errorCode, errors.New("not live"))
    54  	return false
    55  }