cyberark/secretless-broker

package util

import (
    "context"
    "fmt"
    "log"
    "net/http"
    "strconv"
    "sync"
    "time"

    "github.com/heptiolabs/healthcheck"
)

const defaultHealthCheckPort = 5335

type healthWatcher struct {
    handler        healthcheck.Handler
    isAppReady     bool
    isAppLive      bool
    serverInstance *http.Server
}

var _healthWatcher *healthWatcher
var _healthWatcherSyncOnce = sync.Once{}

func (healthWatcher *healthWatcher) registerLivelinessCheck() {
    checkFunc := func() error {
        if healthWatcher.isAppLive {
            return nil
        }

        // If we are not ready, we return a failed liveliness check
        return fmt.Errorf("secretless is not listening")
    }

    healthWatcher.handler.AddLivenessCheck("listening", checkFunc)
}

func (healthWatcher *healthWatcher) registerReadinessCheck() {
    checkFunc := func() error {
        if healthWatcher.isAppReady {
            return nil
        }

        // If we are not ready, we return a failed readiness check
        return fmt.Errorf("secretless is not ready")
    }

    healthWatcher.handler.AddReadinessCheck("ready", checkFunc)
}

func (healthWatcher *healthWatcher) enable(port int) {
    log.Printf("Initializing health check on :%d...", port)

    healthWatcher.handler = healthcheck.NewHandler()

    healthWatcher.registerLivelinessCheck()
    healthWatcher.registerReadinessCheck()

    healthWatcher.serverInstance = &http.Server{
        Addr:           ":" + strconv.Itoa(port),
        Handler:        healthWatcher.handler,
        ReadTimeout:    5 * time.Second,
        WriteTimeout:   5 * time.Second,
        MaxHeaderBytes: 1 << 20,
    }

    go healthWatcher.serverInstance.ListenAndServe()

    log.Printf("Initialization of health check done. " +
        "You can access the endpoint at `/live` and `/ready`.")
}

func enableHealthCheck() {
    _healthWatcherSyncOnce.Do(func() {
        _healthWatcher = &healthWatcher{}
        _healthWatcher.enable(defaultHealthCheckPort)
    })
}

func disableHealthCheck() {
    if _healthWatcher == nil || _healthWatcher.serverInstance == nil {
        return
    }

    // Clean up everything as best we can to ensure prompt GC
    if err := _healthWatcher.serverInstance.Shutdown(context.Background()); err != nil {
        panic(err)
    }

    _healthWatcher.serverInstance = nil
    _healthWatcher.handler = nil
    _healthWatcher = nil
    _healthWatcherSyncOnce = sync.Once{}
}

// SetAppInitializedFlag enables health check and sets the ready flag.
func SetAppInitializedFlag() {
    enableHealthCheck()
    _healthWatcher.isAppReady = true
}

// SetAppIsLive enables health check and marks the app as live.
func SetAppIsLive(isLive bool) {
    enableHealthCheck()
    _healthWatcher.isAppLive = isLive
}