mirror of https://github.com/k3s-io/k3s
140 lines
2.5 KiB
Go
140 lines
2.5 KiB
Go
package probing
|
|
|
|
import (
|
|
"encoding/json"
|
|
"errors"
|
|
"fmt"
|
|
"net/http"
|
|
"sync"
|
|
"time"
|
|
)
|
|
|
|
var (
|
|
ErrNotFound = errors.New("probing: id not found")
|
|
ErrExist = errors.New("probing: id exists")
|
|
)
|
|
|
|
type Prober interface {
|
|
AddHTTP(id string, probingInterval time.Duration, endpoints []string) error
|
|
Remove(id string) error
|
|
RemoveAll()
|
|
Reset(id string) error
|
|
Status(id string) (Status, error)
|
|
}
|
|
|
|
type prober struct {
|
|
mu sync.Mutex
|
|
targets map[string]*status
|
|
tr http.RoundTripper
|
|
}
|
|
|
|
func NewProber(tr http.RoundTripper) Prober {
|
|
p := &prober{targets: make(map[string]*status)}
|
|
if tr == nil {
|
|
p.tr = http.DefaultTransport
|
|
} else {
|
|
p.tr = tr
|
|
}
|
|
return p
|
|
}
|
|
|
|
func (p *prober) AddHTTP(id string, probingInterval time.Duration, endpoints []string) error {
|
|
p.mu.Lock()
|
|
defer p.mu.Unlock()
|
|
if _, ok := p.targets[id]; ok {
|
|
return ErrExist
|
|
}
|
|
|
|
s := &status{stopC: make(chan struct{})}
|
|
p.targets[id] = s
|
|
|
|
ticker := time.NewTicker(probingInterval)
|
|
|
|
go func() {
|
|
pinned := 0
|
|
for {
|
|
select {
|
|
case <-ticker.C:
|
|
start := time.Now()
|
|
req, err := http.NewRequest("GET", endpoints[pinned], nil)
|
|
if err != nil {
|
|
panic(err)
|
|
}
|
|
resp, err := p.tr.RoundTrip(req)
|
|
if err == nil && resp.StatusCode != http.StatusOK {
|
|
err = fmt.Errorf("got unexpected HTTP status code %s from %s", resp.Status, endpoints[pinned])
|
|
resp.Body.Close()
|
|
}
|
|
if err != nil {
|
|
s.recordFailure(err)
|
|
pinned = (pinned + 1) % len(endpoints)
|
|
continue
|
|
}
|
|
|
|
var hh Health
|
|
d := json.NewDecoder(resp.Body)
|
|
err = d.Decode(&hh)
|
|
resp.Body.Close()
|
|
if err != nil || !hh.OK {
|
|
s.recordFailure(err)
|
|
pinned = (pinned + 1) % len(endpoints)
|
|
continue
|
|
}
|
|
|
|
s.record(time.Since(start), hh.Now)
|
|
case <-s.stopC:
|
|
ticker.Stop()
|
|
return
|
|
}
|
|
}
|
|
}()
|
|
|
|
return nil
|
|
}
|
|
|
|
func (p *prober) Remove(id string) error {
|
|
p.mu.Lock()
|
|
defer p.mu.Unlock()
|
|
|
|
s, ok := p.targets[id]
|
|
if !ok {
|
|
return ErrNotFound
|
|
}
|
|
close(s.stopC)
|
|
delete(p.targets, id)
|
|
return nil
|
|
}
|
|
|
|
func (p *prober) RemoveAll() {
|
|
p.mu.Lock()
|
|
defer p.mu.Unlock()
|
|
|
|
for _, s := range p.targets {
|
|
close(s.stopC)
|
|
}
|
|
p.targets = make(map[string]*status)
|
|
}
|
|
|
|
func (p *prober) Reset(id string) error {
|
|
p.mu.Lock()
|
|
defer p.mu.Unlock()
|
|
|
|
s, ok := p.targets[id]
|
|
if !ok {
|
|
return ErrNotFound
|
|
}
|
|
s.reset()
|
|
return nil
|
|
}
|
|
|
|
func (p *prober) Status(id string) (Status, error) {
|
|
p.mu.Lock()
|
|
defer p.mu.Unlock()
|
|
|
|
s, ok := p.targets[id]
|
|
if !ok {
|
|
return nil, ErrNotFound
|
|
}
|
|
return s, nil
|
|
}
|