github.com/Tyktechnologies/tyk@v2.9.5+incompatible/gateway/host_checker.go

package gateway

import (
	"crypto/tls"
	"math/rand"
	"net"
	"net/http"
	"net/url"
	"runtime"
	"strings"
	"sync"
	"time"

	"github.com/jeffail/tunny"
	proxyproto "github.com/pires/go-proxyproto"
	cache "github.com/pmylund/go-cache"

	"github.com/TykTechnologies/tyk/apidef"
	"github.com/TykTechnologies/tyk/config"
)

const (
	defaultTimeout             = 10
	defaultSampletTriggerLimit = 3
)

var (
	HostCheckerClient = &http.Client{
		Timeout: 500 * time.Millisecond,
	}
	defaultWorkerPoolSize = runtime.NumCPU()
	hostCheckTicker       = make(chan struct{})
)

type HostData struct {
	CheckURL            string
	Protocol            string
	Timeout             time.Duration
	EnableProxyProtocol bool
	Commands            []apidef.CheckCommand
	Method              string
	Headers             map[string]string
	Body                string
	MetaData            map[string]string
}

type HostHealthReport struct {
	HostData
	ResponseCode int
	Latency      float64
	IsTCPError   bool
}

type HostUptimeChecker struct {
	failureCallback    func(HostHealthReport)
	upCallback         func(HostHealthReport)
	pingCallback       func(HostHealthReport)
	workerPoolSize     int
	sampleTriggerLimit int
	checkTimeout       int
	HostList           map[string]HostData
	unHealthyList      map[string]bool
	pool               *tunny.WorkPool

	errorChan       chan HostHealthReport
	okChan          chan HostHealthReport
	stopPollingChan chan bool
	sampleCache     *cache.Cache
	stopLoop        bool
	muStopLoop      sync.RWMutex

	resetListMu sync.Mutex
	doResetList bool
	newList     map[string]HostData
}

func (h *HostUptimeChecker) getStopLoop() bool {
	h.muStopLoop.RLock()
	defer h.muStopLoop.RUnlock()
	return h.stopLoop
}

func (h *HostUptimeChecker) setStopLoop(newValue bool) {
	h.muStopLoop.Lock()
	h.stopLoop = newValue
	h.muStopLoop.Unlock()
}

func (h *HostUptimeChecker) getStaggeredTime() time.Duration {
	if h.checkTimeout <= 5 {
		return time.Duration(h.checkTimeout) * time.Second
	}

	rand.Seed(time.Now().Unix())
	min := h.checkTimeout - 3
	max := h.checkTimeout + 3

	dur := rand.Intn(max-min) + min

	return time.Duration(dur) * time.Second
}

func (h *HostUptimeChecker) HostCheckLoop() {
	for !h.getStopLoop() {
		if isRunningTests() {
			<-hostCheckTicker
		}
		h.resetListMu.Lock()
		if h.doResetList && h.newList != nil {
			h.HostList = h.newList
			h.newList = nil
			h.doResetList = false
			log.Debug("[HOST CHECKER] Host list reset")
		}
		h.resetListMu.Unlock()
		for _, host := range h.HostList {
			_, err := h.pool.SendWork(host)
			if err != nil && err != tunny.ErrPoolNotRunning {
				log.Warnf("[HOST CHECKER] could not send work, error: %v", err)
			}
		}

		if !isRunningTests() {
			time.Sleep(h.getStaggeredTime())
		}
	}
	log.Info("[HOST CHECKER] Checker stopped")
}

func (h *HostUptimeChecker) HostReporter() {
	for {
		select {
		case okHost := <-h.okChan:
			// Clear host from unhealthylist if it exists
			if h.unHealthyList[okHost.CheckURL] {
				newVal := 1
				if count, found := h.sampleCache.Get(okHost.CheckURL); found {
					newVal = count.(int) - 1
				}

				if newVal <= 0 {
					// Reset the count
					h.sampleCache.Delete(okHost.CheckURL)
					log.Warning("[HOST CHECKER] [HOST UP]: ", okHost.CheckURL)
					h.upCallback(okHost)
					delete(h.unHealthyList, okHost.CheckURL)
				} else {
					log.Warning("[HOST CHECKER] [HOST UP BUT NOT REACHED LIMIT]: ", okHost.CheckURL)
					h.sampleCache.Set(okHost.CheckURL, newVal, cache.DefaultExpiration)
				}
			}
			go h.pingCallback(okHost)

		case failedHost := <-h.errorChan:
			newVal := 1
			if count, found := h.sampleCache.Get(failedHost.CheckURL); found {
				newVal = count.(int) + 1
			}

			if newVal >= h.sampleTriggerLimit {
				log.Warning("[HOST CHECKER] [HOST DOWN]: ", failedHost.CheckURL)
				// track it
				h.unHealthyList[failedHost.CheckURL] = true
				// Call the custom callback hook
				go h.failureCallback(failedHost)
			} else {
				log.Warning("[HOST CHECKER] [HOST DOWN BUT NOT REACHED LIMIT]: ", failedHost.CheckURL)
				h.sampleCache.Set(failedHost.CheckURL, newVal, cache.DefaultExpiration)
			}
			go h.pingCallback(failedHost)

		case <-h.stopPollingChan:
			log.Debug("[HOST CHECKER] Received kill signal")
			return
		}
	}
}

func (h *HostUptimeChecker) CheckHost(toCheck HostData) {
	log.Debug("[HOST CHECKER] Checking: ", toCheck.CheckURL)

	t1 := time.Now()
	report := HostHealthReport{
		HostData: toCheck,
	}
	switch toCheck.Protocol {
	case "tcp", "tls":
		host := toCheck.CheckURL
		base := toCheck.Protocol + "://"
		if !strings.HasPrefix(host, base) {
			host = base + host
		}
		u, err := url.Parse(host)
		if err != nil {
			log.Error("Could not parse host: ", err)
			return
		}
		var ls net.Conn
		var d net.Dialer
		d.Timeout = toCheck.Timeout
		if toCheck.Protocol == "tls" {
			ls, err = tls.DialWithDialer(&d, "tls", u.Host, nil)
		} else {
			ls, err = d.Dial("tcp", u.Host)
		}
		if err != nil {
			log.Error("Could not connect to host: ", err)
			report.IsTCPError = true
			break
		}
		if toCheck.EnableProxyProtocol {
			log.Debug("using proxy protocol")
			ls = proxyproto.NewConn(ls, 0)
		}
		defer ls.Close()
		for _, cmd := range toCheck.Commands {
			switch cmd.Name {
			case "send":
				log.Debugf("%s: sending %s", host, cmd.Message)
				_, err = ls.Write([]byte(cmd.Message))
				if err != nil {
					log.Errorf("Failed to send %s :%v", cmd.Message, err)
					report.IsTCPError = true
					break
				}
			case "expect":
				buf := make([]byte, len(cmd.Message))
				_, err = ls.Read(buf)
				if err != nil {
					log.Errorf("Failed to read %s :%v", cmd.Message, err)
					report.IsTCPError = true
					break
				}
				g := string(buf)
				if g != cmd.Message {
					log.Errorf("Failed expectation  expected %s got %s", cmd.Message, g)
					report.IsTCPError = true
					break
				}
				log.Debugf("%s: received %s", host, cmd.Message)
			}
		}
		report.ResponseCode = http.StatusOK
	default:
		useMethod := toCheck.Method
		if toCheck.Method == "" {
			useMethod = http.MethodGet
		}
		req, err := http.NewRequest(useMethod, toCheck.CheckURL, strings.NewReader(toCheck.Body))
		if err != nil {
			log.Error("Could not create request: ", err)
			return
		}
		for headerName, headerValue := range toCheck.Headers {
			req.Header.Set(headerName, headerValue)
		}
		req.Header.Set("Connection", "close")
		HostCheckerClient.Transport = &http.Transport{
			TLSClientConfig: &tls.Config{
				InsecureSkipVerify: config.Global().ProxySSLInsecureSkipVerify,
			},
		}
		if toCheck.Timeout != 0 {
			HostCheckerClient.Timeout = toCheck.Timeout
		}
		response, err := HostCheckerClient.Do(req)
		if err != nil {
			report.IsTCPError = true
			break
		}
		response.Body.Close()
		report.ResponseCode = response.StatusCode
	}

	millisec := DurationToMillisecond(time.Since(t1))
	report.Latency = millisec
	if report.IsTCPError {
		h.errorChan <- report
		return
	}

	if report.ResponseCode != http.StatusOK {
		h.errorChan <- report
		return
	}

	// host is healthy, report it
	h.okChan <- report
}

func (h *HostUptimeChecker) Init(workers, triggerLimit, timeout int, hostList map[string]HostData, failureCallback, upCallback, pingCallback func(HostHealthReport)) {
	h.sampleCache = cache.New(30*time.Second, 30*time.Second)
	h.stopPollingChan = make(chan bool)
	h.errorChan = make(chan HostHealthReport)
	h.okChan = make(chan HostHealthReport)
	h.HostList = hostList
	h.unHealthyList = make(map[string]bool)
	h.failureCallback = failureCallback
	h.upCallback = upCallback
	h.pingCallback = pingCallback

	h.workerPoolSize = workers
	if workers == 0 {
		h.workerPoolSize = defaultWorkerPoolSize
	}

	h.sampleTriggerLimit = triggerLimit
	if triggerLimit == 0 {
		h.sampleTriggerLimit = defaultSampletTriggerLimit
	}

	h.checkTimeout = timeout
	if timeout == 0 {
		h.checkTimeout = defaultTimeout
	}

	log.Debug("[HOST CHECKER] Config:TriggerLimit: ", h.sampleTriggerLimit)
	log.Debug("[HOST CHECKER] Config:Timeout: ~", h.checkTimeout)
	log.Debug("[HOST CHECKER] Config:WorkerPool: ", h.workerPoolSize)

	var err error
	h.pool, err = tunny.CreatePool(h.workerPoolSize, func(hostData interface{}) interface{} {
		input, _ := hostData.(HostData)
		h.CheckHost(input)
		return nil
	}).Open()

	log.Debug("[HOST CHECKER] Init complete")

	if err != nil {
		log.Errorf("[HOST CHECKER POOL] Error: %v\n", err)
	}
}

func (h *HostUptimeChecker) Start() {
	// Start the loop that checks for bum hosts
	h.setStopLoop(false)
	log.Debug("[HOST CHECKER] Starting...")
	go h.HostCheckLoop()
	log.Debug("[HOST CHECKER] Check loop started...")
	go h.HostReporter()
	log.Debug("[HOST CHECKER] Host reporter started...")
}

func (h *HostUptimeChecker) Stop() {
	h.setStopLoop(true)

	h.stopPollingChan <- true
	log.Info("[HOST CHECKER] Stopping poller")
	h.pool.Close()
}

func (h *HostUptimeChecker) ResetList(hostList map[string]HostData) {
	h.resetListMu.Lock()
	h.doResetList = true
	h.newList = hostList
	h.resetListMu.Unlock()
	log.Debug("[HOST CHECKER] Checker reset queued!")
}