Skip to content

Commit 56a11ac

Browse files
pashagolubpashagolub
authored
[+] add /readiness and /liveness endpoints, closes #611 (#615)
* [+] add `/readiness` and `/liveness` endpoints, closes #611 * [+] add `--web-disable=ui` option to disable only UI and keep REST API
1 parent f3f11f1 commit 56a11ac

File tree

7 files changed

+181
-56
lines changed

7 files changed

+181
-56
lines changed

‎cmd/pgwatch/main.go‎

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@ import (
1313
"github.com/cybertec-postgresql/pgwatch/v3/internal/cmdopts"
1414
"github.com/cybertec-postgresql/pgwatch/v3/internal/log"
1515
"github.com/cybertec-postgresql/pgwatch/v3/internal/reaper"
16+
"github.com/cybertec-postgresql/pgwatch/v3/internal/webserver"
1617
"github.com/cybertec-postgresql/pgwatch/v3/internal/webui"
1718
)
1819

@@ -90,13 +91,15 @@ func main() {
9091
return
9192
}
9293

93-
if err = opts.InitWebUI(webui.WebUIFs, logger); err != nil {
94+
reaper := reaper.NewReaper(opts, opts.SourcesReaderWriter, opts.MetricsReaderWriter)
95+
96+
if _, err = webserver.Init(mainCtx, opts.WebUI, webui.WebUIFs, opts.MetricsReaderWriter,
97+
opts.SourcesReaderWriter, reaper); err != nil {
9498
exitCode.Store(cmdopts.ExitCodeWebUIError)
95-
logger.Error(err)
99+
logger.Error("failed to initialize web UI: ", err)
96100
return
97101
}
98102

99-
reaper := reaper.NewReaper(opts, opts.SourcesReaderWriter, opts.MetricsReaderWriter)
100103
if err = reaper.Reap(mainCtx); err != nil {
101104
logger.Error(err)
102105
exitCode.Store(cmdopts.ExitCodeFatalError)

‎internal/cmdopts/cmdoptions.go‎

Lines changed: 0 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,6 @@ import (
55
"errors"
66
"fmt"
77
"io"
8-
"io/fs"
98
"os"
109
"time"
1110

@@ -179,18 +178,6 @@ func (c *Options) NeedsSchemaUpgrade() (upgrade bool, err error) {
179178
return
180179
}
181180

182-
// InitWebUI initializes the web UI server
183-
func (c *Options) InitWebUI(fs fs.FS, logger log.LoggerIface) error {
184-
if c.WebUI.WebDisable {
185-
logger.Info("web user interface is disabled")
186-
return nil
187-
}
188-
if webserver.Init(c.WebUI, fs, c.MetricsReaderWriter, c.SourcesReaderWriter, logger) == nil {
189-
return errors.New("failed to initialize web UI")
190-
}
191-
return nil
192-
}
193-
194181
func validateConfig(c *Options) error {
195182
if len(c.Sources.Sources)+len(c.Metrics.Metrics) == 0 {
196183
return errors.New("both --sources and --metrics are empty")

‎internal/reaper/reaper.go‎

Lines changed: 24 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,8 @@ import (
88
"sync"
99
"time"
1010

11+
"sync/atomic"
12+
1113
"github.com/cybertec-postgresql/pgwatch/v3/internal/cmdopts"
1214
"github.com/cybertec-postgresql/pgwatch/v3/internal/log"
1315
"github.com/cybertec-postgresql/pgwatch/v3/internal/metrics"
@@ -24,6 +26,7 @@ var metricDefinitionMap *metrics.Metrics = &metrics.Metrics{}
2426
var metricDefMapLock = sync.RWMutex{}
2527

2628
type Reaper struct {
29+
ready atomic.Bool
2730
opts *cmdopts.Options
2831
sourcesReaderWriter sources.ReaderWriter
2932
metricsReaderWriter metrics.ReaderWriter
@@ -35,15 +38,24 @@ func NewReaper(opts *cmdopts.Options, sourcesReaderWriter sources.ReaderWriter,
3538
opts: opts,
3639
sourcesReaderWriter: sourcesReaderWriter,
3740
metricsReaderWriter: metricsReaderWriter,
41+
measurementCh: make(chan []metrics.MeasurementEnvelope, 10000),
3842
}
3943
}
4044

45+
// Ready() returns true if the service is healthy and operating correctly
46+
func (r *Reaper) Ready() bool {
47+
return r.ready.Load()
48+
}
49+
50+
// Reap() starts the main monitoring loop. It is responsible for fetching metrics measurements
51+
// from the sources and storing them to the sinks. It also manages the lifecycle of
52+
// the metric gatherers. In case of a source or metric definition change, it will
53+
// start or stop the gatherers accordingly.
4154
func (r *Reaper) Reap(mainContext context.Context) (err error) {
4255
var measurementsWriter *sinks.MultiWriter
4356

4457
cancelFuncs := make(map[string]context.CancelFunc) // [db1+metric1]=chan
4558

46-
firstLoop := true
4759
mainLoopCount := 0
4860
logger := log.GetLogger(mainContext)
4961
metricsReaderWriter := r.metricsReaderWriter
@@ -59,23 +71,19 @@ func (r *Reaper) Reap(mainContext context.Context) (err error) {
5971
if measurementsWriter, err = sinks.NewMultiWriter(mainContext, &opts.Sinks, metricDefinitionMap); err != nil {
6072
logger.Fatal(err)
6173
}
62-
r.measurementCh = make(chan []metrics.MeasurementEnvelope, 10000)
6374
go measurementsWriter.WriteMeasurements(mainContext, r.measurementCh)
6475

76+
if monitoredDbs, err = monitoredDbs.SyncFromReader(sourcesReaderWriter); err != nil {
77+
logger.Fatal("could not fetch active hosts - check config!", err)
78+
}
79+
80+
// at this stage we have all the metric definitions, the sinks and the sources configured
81+
r.ready.Store(true)
82+
6583
for { //main loop
6684
hostsToShutDownDueToRoleChange := make(map[string]bool) // hosts went from master to standby and have "only if master" set
6785
gatherersShutDown := 0
6886

69-
if monitoredDbs, err = monitoredDbs.SyncFromReader(sourcesReaderWriter); err != nil {
70-
if firstLoop {
71-
logger.Fatal("could not fetch active hosts - check config!", err)
72-
} else {
73-
logger.Error("could not fetch active hosts, using last valid config data:", err)
74-
time.Sleep(time.Second * time.Duration(opts.Sources.Refresh))
75-
continue
76-
}
77-
}
78-
7987
if DoesEmergencyTriggerfileExist(opts.Metrics.EmergencyPauseTriggerfile) {
8088
logger.Warningf("Emergency pause triggerfile detected at %s, ignoring currently configured DBs", opts.Metrics.EmergencyPauseTriggerfile)
8189
monitoredDbs = make([]*sources.MonitoredDatabase, 0)
@@ -93,14 +101,12 @@ func (r *Reaper) Reap(mainContext context.Context) (err error) {
93101
WithField("metrics", len(metricDefinitionMap.MetricDefs)).
94102
WithField("presets", len(metricDefinitionMap.PresetDefs)).
95103
Log(func() logrus.Level {
96-
if firstLoop && len(monitoredDbs)*len(metricDefinitionMap.MetricDefs) == 0 {
104+
if len(monitoredDbs)*len(metricDefinitionMap.MetricDefs) == 0 {
97105
return logrus.WarnLevel
98106
}
99107
return logrus.InfoLevel
100108
}(), "sources and metrics refreshed")
101109

102-
firstLoop = false // only used for failing when 1st config reading fails
103-
104110
for _, monitoredDB := range monitoredDbs {
105111
logger.WithField("source", monitoredDB.Name).
106112
WithField("metric", monitoredDB.Metrics).
@@ -358,7 +364,9 @@ func (r *Reaper) Reap(mainContext context.Context) (err error) {
358364
logger.Debugf("main sleeping %ds...", opts.Sources.Refresh)
359365
select {
360366
case <-time.After(time.Second * time.Duration(opts.Sources.Refresh)):
361-
// pass
367+
if monitoredDbs, err = monitoredDbs.SyncFromReader(sourcesReaderWriter); err != nil {
368+
logger.Error("could not fetch active hosts, using last valid config data:", err)
369+
}
362370
case <-mainContext.Done():
363371
return
364372
}

‎internal/webserver/cmdoptions_test.go‎

Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,39 @@
1+
package webserver
2+
3+
import (
4+
"os"
5+
"testing"
6+
7+
flags "github.com/jessevdk/go-flags"
8+
"github.com/stretchr/testify/assert"
9+
)
10+
11+
func TestWebDisableOpt(t *testing.T) {
12+
a := assert.New(t)
13+
testCases := []struct {
14+
args []string
15+
expected string
16+
expectError bool
17+
}{
18+
{[]string{0: "config_test"}, "", false},
19+
{[]string{0: "config_test", "--web-disable"}, WebDisableAll, false},
20+
{[]string{0: "config_test", "--web-disable=all"}, WebDisableAll, false},
21+
{[]string{0: "config_test", "--web-disable=ui"}, WebDisableUI, false},
22+
{[]string{0: "config_test", "--web-disable=foo"}, "", true},
23+
}
24+
25+
for _, tc := range testCases {
26+
opts := new(CmdOpts)
27+
os.Args = tc.args
28+
_, err := flags.NewParser(opts, flags.HelpFlag).Parse()
29+
30+
if tc.expectError {
31+
a.Error(err)
32+
a.Empty(opts.WebDisable)
33+
} else {
34+
a.NoError(err)
35+
a.Equal(tc.expected, opts.WebDisable)
36+
}
37+
}
38+
39+
}

‎internal/webserver/cmdopts.go‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,13 @@
11
package webserver
22

3+
const (
4+
WebDisableAll string = "all"
5+
WebDisableUI string = "ui"
6+
)
7+
38
// CmdOpts specifies the internal web UI server command-line options
49
type CmdOpts struct {
5-
WebDisable bool `long:"web-disable" mapstructure:"web-disable" description:"Disable the web UI" env:"PW_WEBDISABLE"`
10+
WebDisable string `long:"web-disable" mapstructure:"web-disable" description:"Disable REST API and/or web UI" env:"PW_WEBDISABLE" optional:"true" optional-value:"all" choice:"all" choice:"ui"`
611
WebAddr string `long:"web-addr" mapstructure:"web-addr" description:"TCP address in the form 'host:port' to listen on" default:":8080" env:"PW_WEBADDR"`
712
WebUser string `long:"web-user" mapstructure:"web-user" description:"Admin login" env:"PW_WEBUSER"`
813
WebPassword string `long:"web-password" mapstructure:"web-password" description:"Admin password" env:"PW_WEBPASSWORD"`

‎internal/webserver/server_test.go‎

Lines changed: 52 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
package webserver_test
22

33
import (
4+
"context"
45
"encoding/json"
56
"fmt"
67
"io"
@@ -10,7 +11,6 @@ import (
1011
"strings"
1112
"testing"
1213

13-
"github.com/cybertec-postgresql/pgwatch/v3/internal/log"
1414
"github.com/cybertec-postgresql/pgwatch/v3/internal/webserver"
1515
"github.com/stretchr/testify/assert"
1616
)
@@ -20,20 +20,61 @@ type Credentials struct {
2020
Password string `json:"password"`
2121
}
2222

23-
func TestStatus(t *testing.T) {
24-
restsrv := webserver.Init(webserver.CmdOpts{WebAddr: "127.0.0.1:8080"}, os.DirFS("../webui/build"), nil, nil, log.FallbackLogger)
23+
type ReadyBool bool
24+
25+
func (ready *ReadyBool) Ready() bool {
26+
return bool(*ready)
27+
}
28+
29+
func TestWebDisableOpt(t *testing.T) {
30+
var ready ReadyBool
31+
restsrv, err := webserver.Init(context.Background(), webserver.CmdOpts{WebDisable: "all"}, os.DirFS("../webui/build"), nil, nil, &ready)
32+
assert.Nil(t, restsrv, "no webserver should be started")
33+
assert.NoError(t, err)
34+
35+
restsrv, err = webserver.Init(context.Background(), webserver.CmdOpts{WebAddr: "127.0.0.1:8079", WebDisable: "ui"}, os.DirFS("../webui/build"), nil, nil, &ready)
2536
assert.NotNil(t, restsrv)
26-
// r, err := http.Get("http://localhost:8080/")
27-
// assert.NoError(t, err)
28-
// assert.Equal(t, http.StatusOK, r.StatusCode)
29-
// b, err := io.ReadAll(r.Body)
30-
// assert.NoError(t, err)
31-
// assert.True(t, len(b) > 0)
37+
assert.NoError(t, err)
38+
r, err := http.Get("http://localhost:8079/")
39+
assert.NoError(t, err)
40+
assert.Equal(t, http.StatusNotFound, r.StatusCode, "no webui should be served")
41+
r, err = http.Get("http://localhost:8079/liveness")
42+
assert.NoError(t, err)
43+
assert.Equal(t, http.StatusOK, r.StatusCode, "rest api should be served though")
44+
45+
restsrv, err = webserver.Init(context.Background(), webserver.CmdOpts{WebAddr: "127.0.0.1:8079"}, os.DirFS("../webui/build"), nil, nil, &ready)
46+
assert.Nil(t, restsrv)
47+
assert.Error(t, err, "port should be in use")
48+
}
49+
50+
func TestHealth(t *testing.T) {
51+
var ready ReadyBool
52+
ctx, cancel := context.WithCancel(context.Background())
53+
restsrv, _ := webserver.Init(ctx, webserver.CmdOpts{WebAddr: "127.0.0.1:8080"}, os.DirFS("../webui/build"), nil, nil, &ready)
54+
assert.NotNil(t, restsrv)
55+
56+
r, err := http.Get("http://localhost:8080/liveness")
57+
assert.NoError(t, err)
58+
assert.Equal(t, http.StatusOK, r.StatusCode)
59+
60+
cancel()
61+
r, err = http.Get("http://localhost:8080/liveness")
62+
assert.NoError(t, err)
63+
assert.Equal(t, http.StatusServiceUnavailable, r.StatusCode)
64+
65+
r, err = http.Get("http://localhost:8080/readiness")
66+
assert.NoError(t, err)
67+
assert.Equal(t, http.StatusServiceUnavailable, r.StatusCode)
68+
69+
ready = true
70+
r, err = http.Get("http://localhost:8080/readiness")
71+
assert.NoError(t, err)
72+
assert.Equal(t, http.StatusOK, r.StatusCode)
3273
}
3374

3475
func TestServerNoAuth(t *testing.T) {
3576
host := "http://localhost:8081"
36-
restsrv := webserver.Init(webserver.CmdOpts{WebAddr: "localhost:8081"}, os.DirFS("../webui/build"), nil, nil, log.FallbackLogger)
77+
restsrv, _ := webserver.Init(context.Background(), webserver.CmdOpts{WebAddr: "localhost:8081"}, os.DirFS("../webui/build"), nil, nil, nil)
3778
assert.NotNil(t, restsrv)
3879
rr := httptest.NewRecorder()
3980
// test request metrics
@@ -64,7 +105,7 @@ func TestServerNoAuth(t *testing.T) {
64105

65106
func TestGetToken(t *testing.T) {
66107
host := "http://localhost:8082"
67-
restsrv := webserver.Init(webserver.CmdOpts{WebAddr: "localhost:8082"}, os.DirFS("../webui/build"), nil, nil, log.FallbackLogger)
108+
restsrv, _ := webserver.Init(context.Background(), webserver.CmdOpts{WebAddr: "localhost:8082"}, os.DirFS("../webui/build"), nil, nil, nil)
68109
rr := httptest.NewRecorder()
69110

70111
credentials := Credentials{

0 commit comments

Comments
 (0)