Prism/internal/health/checker.go

package health

import (
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"strings"
	"sync/atomic"
	"time"

	"github.com/prism/proxy/internal/cache"
	"github.com/prism/proxy/internal/store"
)

type ReloadFunc func(ctx context.Context) error

type Checker struct {
	store    *store.Store
	cache    *cache.Cache
	apiBase  string
	secret   string
	client   *http.Client
	testing  atomic.Bool
}

func NewChecker(s *store.Store, c *cache.Cache, apiBase, secret string) *Checker {
	return &Checker{
		store:   s,
		cache:   c,
		apiBase: stringsTrimRight(apiBase, "/"),
		secret:  secret,
		client: &http.Client{
			Timeout: 10 * time.Second,
		},
	}
}

func stringsTrimRight(s, cut string) string {
	for len(s) > 0 && stringsHasSuffix(s, cut) {
		s = s[:len(s)-len(cut)]
	}
	return s
}

func stringsHasSuffix(s, suffix string) bool {
	return len(s) >= len(suffix) && s[len(s)-len(suffix):] == suffix
}

func (h *Checker) TestNode(ctx context.Context, nodeID int64, reload ReloadFunc) (int, error) {
	node, err := h.store.GetProxyNode(ctx, nodeID)
	if err != nil {
		return -1, err
	}
	latency, err := h.testRuntimeName(node.RuntimeName)
	if isNotFound(err) && reload != nil {
		if reloadErr := reload(ctx); reloadErr == nil {
			latency, err = h.testRuntimeName(node.RuntimeName)
		}
	}
	if isNotFound(err) {
		err = fmt.Errorf("proxy %q not loaded in engine; reload engine in settings", node.RuntimeName)
	}
	alive := err == nil && latency >= 0
	lat := latency
	if !alive {
		lat = -1
	}
	_ = h.store.UpdateProxyNodeHealth(ctx, nodeID, lat, alive)
	h.cache.SetProxyLatency(node.RuntimeName, lat)
	return lat, err
}

func (h *Checker) TestAll(ctx context.Context, _ ReloadFunc) error {
	if !h.testing.CompareAndSwap(false, true) {
		return nil
	}
	go func() {
		defer h.testing.Store(false)
		nodes, err := h.store.ListProxyNodes(context.Background(), "", nil)
		if err != nil {
			return
		}
		for _, n := range nodes {
			if !n.Enabled {
				continue
			}
			_, _ = h.TestNode(context.Background(), n.ID, nil)
		}
	}()
	return nil
}

func (h *Checker) ProxyLoaded(name string) bool {
	data, err := h.getJSON("/proxies")
	if err != nil {
		return false
	}
	proxies, ok := data["proxies"].(map[string]any)
	if !ok {
		return false
	}
	_, ok = proxies[name]
	return ok
}

func (h *Checker) testRuntimeName(name string) (int, error) {
	testURL := fmt.Sprintf("%s/proxies/%s/delay?timeout=5000&url=%s", h.apiBase, url.PathEscape(name), url.QueryEscape("http://www.gstatic.com/generate_204"))
	req, err := http.NewRequest(http.MethodGet, testURL, nil)
	if err != nil {
		return -1, err
	}
	if h.secret != "" {
		req.Header.Set("Authorization", "Bearer "+h.secret)
	}
	resp, err := h.client.Do(req)
	if err != nil {
		return -1, err
	}
	defer resp.Body.Close()
	body, _ := io.ReadAll(resp.Body)
	if resp.StatusCode >= 400 {
		return -1, fmt.Errorf("delay test failed: %s", string(body))
	}
	var result struct {
		Delay int `json:"delay"`
	}
	if err := json.Unmarshal(body, &result); err != nil {
		return -1, err
	}
	return result.Delay, nil
}

func (h *Checker) getJSON(path string) (map[string]any, error) {
	req, err := http.NewRequest(http.MethodGet, h.apiBase+path, nil)
	if err != nil {
		return nil, err
	}
	if h.secret != "" {
		req.Header.Set("Authorization", "Bearer "+h.secret)
	}
	resp, err := h.client.Do(req)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	body, err := io.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}
	if resp.StatusCode >= 400 {
		return nil, fmt.Errorf("api error %d: %s", resp.StatusCode, string(body))
	}
	var out map[string]any
	if err := json.Unmarshal(body, &out); err != nil {
		return nil, err
	}
	return out, nil
}

func isNotFound(err error) bool {
	return err != nil && strings.Contains(err.Error(), "Resource not found")
}

func (h *Checker) UpdateAPIBase(apiBase, secret string) {
	h.apiBase = stringsTrimRight(apiBase, "/")
	h.secret = secret
}