proberesponder/proberesponder.go at main · naughtygopher/proberesponder · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
package proberesponder

import (
	"fmt"
	"sync"
	"time"
)

type Statuskey string

func (sk Statuskey) String() string {
	return string(sk)
}

const (
	StatusStartup Statuskey = "startup"
	StatusReady   Statuskey = "ready"
	StatusLive    Statuskey = "live"
)

type healthstatus string

func (hs healthstatus) String() string {
	return string(hs)
}

func (hs healthstatus) Equal(s healthstatus) bool {
	return s == hs
}

func IsHealthOK[T ~string](s T) bool {
	return HealthOK.Equal(healthstatus(s))
}

const (
	HealthOK    healthstatus = "OK"
	HealthNotOK healthstatus = "NOT OK"
)

type StatusChangeListener func(status Statuskey, value bool)

// ProbeStatuses are maintained primarily for K8s probe responses. Though it can be used
// for any prober.
type ProbeResponder struct {
	notReady       bool
	notLive        bool
	notStarted     bool
	locker         *sync.Mutex
	msgPayload     map[string]string
	changeListener StatusChangeListener
}

func (pr *ProbeResponder) AppendHealthResponse(key, value string) {
	if pr == nil {
		return
	}
	pr.locker.Lock()
	defer pr.locker.Unlock()

	pr.appendHealthRespWithoutLock(key, value)
}

func (pr *ProbeResponder) appendHealthRespWithoutLock(key, value string) {
	pr.msgPayload[key] = value
}

func (pr *ProbeResponder) HealthResponse() map[string]string {
	if pr == nil {
		return nil
	}
	pr.locker.Lock()
	defer pr.locker.Unlock()

	copied := map[string]string{}
	for k, v := range pr.msgPayload {
		copied[k] = v
	}

	return copied
}

func (pr *ProbeResponder) onChange(status Statuskey, value bool) {
	hs := HealthOK
	if value {
		hs = HealthNotOK
	}

	pr.appendHealthRespWithoutLock(
		"probe->"+status.String(),
		fmt.Sprintf("%s: %s", hs, time.Now().Format(time.RFC3339)),
	)

	if pr.changeListener == nil {
		return
	}

	pr.changeListener(status, value)
}

func (pr *ProbeResponder) SetNotReady(b bool) {
	if pr == nil {
		return
	}

	pr.locker.Lock()
	defer pr.locker.Unlock()

	pr.notReady = b
	pr.onChange(StatusReady, b)
}

func (pr *ProbeResponder) SetNotLive(b bool) {
	if pr == nil {
		return
	}

	pr.locker.Lock()
	defer pr.locker.Unlock()

	pr.notLive = b
	pr.onChange(StatusLive, b)
}

func (pr *ProbeResponder) SetNotStarted(b bool) {
	if pr == nil {
		return
	}

	pr.locker.Lock()
	defer pr.locker.Unlock()

	pr.notStarted = b
	pr.onChange(StatusStartup, b)
}

// SetListener is used to set a callback function which will be invoked every time
// any of the statuses change (e.g. liveness)
func (pr *ProbeResponder) SetListener(l StatusChangeListener) {
	if pr == nil {
		return
	}

	pr.locker.Lock()
	defer pr.locker.Unlock()

	pr.changeListener = l
}

func (pr *ProbeResponder) NotReady() bool {
	return pr != nil && pr.notReady
}

func (pr *ProbeResponder) NotLive() bool {
	return pr != nil && pr.notLive
}

func (pr *ProbeResponder) NotStarted() bool {
	return pr != nil && pr.notStarted
}

func New() *ProbeResponder {
	pRes := &ProbeResponder{
		locker:     &sync.Mutex{},
		msgPayload: map[string]string{},
	}

	pRes.SetNotLive(true)
	pRes.SetNotReady(true)
	pRes.SetNotStarted(true)

	return pRes
}