Blacktip/health.go

package blacktip

import (
	"context"
	"crypto/tls"
	"fmt"
	"io"
	"net"
	"net/http"
	"strings"
	"sync"
	"time"

	"git.veillax.com/Veillax/reef/model"
)

// healthLoop checks every target of ms until ctx is cancelled. Routing prefers
// healthy targets, so a dead backend is skipped without anyone editing config.
func (b *Blacktip) healthLoop(ctx context.Context, ms *mapState) {
	every := time.Duration(ms.m.HealthSeconds) * time.Second
	if ms.m.HealthSeconds <= 0 {
		every = 30 * time.Second
	}
	run := func() {
		results := make([]model.TargetHealth, len(ms.m.Targets))
		var wg sync.WaitGroup
		for i := range ms.m.Targets {
			wg.Add(1)
			go func(i int) {
				defer wg.Done()
				results[i] = ms.check(ctx, i)
			}(i)
		}
		wg.Wait()
		if ctx.Err() != nil {
			return
		}
		ms.mu.Lock()
		changed := false
		for i, r := range results {
			if ms.health[i].Healthy != r.Healthy {
				changed = true
			}
		}
		ms.health = results
		ms.mu.Unlock()
		if changed {
			go b.report()
		}
	}
	// healthy targets are re-checked every `every`; unhealthy ones every few
	// seconds so recovery (or a site that was still connecting) shows up fast
	run()
	for {
		wait := every
		ms.mu.Lock()
		for _, h := range ms.health {
			if !h.Healthy && wait > 5*time.Second {
				wait = 5 * time.Second
			}
		}
		ms.mu.Unlock()
		select {
		case <-ctx.Done():
			return
		case <-time.After(wait):
			run()
		}
	}
}

func (ms *mapState) check(ctx context.Context, i int) model.TargetHealth {
	t := ms.m.Targets[i]
	if model.Web(ms.m.Proto) && ms.m.HealthPath != "" {
		return ms.checkHTTP(ctx, t)
	}
	c, tunneled, err := ms.b.dialTarget(t, "tcp")
	if err != nil {
		return model.TargetHealth{Detail: err.Error()}
	}
	defer c.Close()
	if tunneled {
		// A tunnel stream opens even when the far side can't connect, but the
		// agent then closes it straight away. Silence means the target is up.
		c.SetReadDeadline(time.Now().Add(800 * time.Millisecond))
		one := make([]byte, 1)
		if _, err := c.Read(one); err != nil {
			if ne, ok := err.(net.Error); ok && ne.Timeout() {
				return model.TargetHealth{Healthy: true}
			}
			if err == io.EOF || strings.Contains(err.Error(), "closed") || strings.Contains(err.Error(), "reset") {
				return model.TargetHealth{Detail: "the site couldn't connect to the target"}
			}
			return model.TargetHealth{Detail: err.Error()}
		}
	}
	return model.TargetHealth{Healthy: true}
}

func (ms *mapState) checkHTTP(ctx context.Context, t model.Target) model.TargetHealth {
	scheme := "http"
	if ms.m.BackendHTTPS {
		scheme = "https"
	}
	tr := &http.Transport{
		DialContext: func(ctx context.Context, _, _ string) (net.Conn, error) {
			c, _, err := ms.b.dialTarget(t, "tcp")
			return c, err
		},
		TLSClientConfig:   &tls.Config{InsecureSkipVerify: true},
		DisableKeepAlives: true,
	}
	defer tr.CloseIdleConnections()
	ctx, cancel := context.WithTimeout(ctx, 6*time.Second)
	defer cancel()
	req, _ := http.NewRequestWithContext(ctx, "GET", scheme+"://target"+ms.m.HealthPath, nil)
	if ms.m.RewriteHost != "" {
		req.Host = ms.m.RewriteHost
	} else if h := strings.TrimPrefix(ms.m.Hostname, "*."); h != "" {
		req.Host = h
	}
	req.Header.Set("User-Agent", "blacktip-healthcheck")
	resp, err := (&http.Client{Transport: tr, CheckRedirect: func(*http.Request, []*http.Request) error { return http.ErrUseLastResponse }}).Do(req)
	if err != nil {
		return model.TargetHealth{Detail: err.Error()}
	}
	resp.Body.Close()
	if resp.StatusCode >= 200 && resp.StatusCode < 400 {
		return model.TargetHealth{Healthy: true}
	}
	return model.TargetHealth{Detail: fmt.Sprintf("health check got HTTP %d", resp.StatusCode)}
}