tests

Co-authored-by: Dmitry Fedotov <dmitry@uint32.ru>
Co-committed-by: Dmitry Fedotov <dmitry@uint32.ru>

watchdog.go

@@ -15,35 +15,36 @@ var (
 	ErrNotRunning    = errors.New("watchdog is not running")
 )
 
+// Watchdog keeps checks to run either periodically
+// or on demand.
 type Watchdog struct {
 	checks  []Check
 	running *running
 	mu      sync.Mutex
 }
 
-// NewWatchDog accepts settings for a number of groups to monitor.
-// keys of a map are group names. values are slices of their respective checks.
-func NewWatchDog(checks ...Check) *Watchdog {
-	w := &Watchdog{
-		checks: checks,
-	}
-
-	return w
-}
-
 type running struct {
 	out  chan CheckResult
 	stop chan struct{}
 }
 
-func (w *Watchdog) ListChecks() ([]Check, error) {
+// New accepts checks to run.
+func New(checks ...Check) *Watchdog {
+	w := &Watchdog{
+		checks: checks,
+	}
+
+	return w
+}
+
+func (w *Watchdog) ListChecks() []Check {
 	w.mu.Lock()
 	defer w.mu.Unlock()
 
 	out := make([]Check, len(w.checks))
 	copy(out, w.checks)
 
-	return out, nil
+	return out
 }
 
 // AddChecks adds checks to the group. This DOES NOT
@@ -51,18 +52,16 @@ func (w *Watchdog) ListChecks() ([]Check, error) {
 // then Start to restart monitoring when a new check is added.
 // Check may have duplicate Name fields but note that RemoveChecks removes checks
 // by their Name fields.
-func (w *Watchdog) AddChecks(checks ...Check) error {
+func (w *Watchdog) AddChecks(checks ...Check) {
 	w.mu.Lock()
 	defer w.mu.Unlock()
 
 	w.checks = append(w.checks, checks...)
-
-	return nil
 }
 
 // RemoveChecks removes the named checks.
 // This does not affect the already running monitoring for the group.
-func (w *Watchdog) RemoveChecks(names ...string) error {
+func (w *Watchdog) RemoveChecks(names ...string) {
 	w.mu.Lock()
 	defer w.mu.Unlock()
 
@@ -76,8 +75,6 @@ func (w *Watchdog) RemoveChecks(names ...string) error {
 	}
 
 	w.checks = remaining
-
-	return nil
 }
 
 // Start starts monitoring.
@@ -87,7 +84,9 @@ func (w *Watchdog) RemoveChecks(names ...string) error {
 // the group to the channel.
 // Subsequently only changes of status are pushed regardless of return error values of
 // CheckFunc provided to Watchdog.
-func (w *Watchdog) Start(ctx context.Context) (<-chan CheckResult, error) {
+// Concurrency limits the number of checks that can run concurrently. 0 means no
+// limit (all checks may run concurrently).
+func (w *Watchdog) Start(ctx context.Context, concurrency int) (<-chan CheckResult, error) {
 	w.mu.Lock()
 	defer w.mu.Unlock()
 
@@ -100,15 +99,19 @@ func (w *Watchdog) Start(ctx context.Context) (<-chan CheckResult, error) {
 		return nil, ErrNotConfigured
 	}
 
-	r := runMonitoringForGroup(ctx, w.checks)
-	w.running = r
+	cp := w.copyChecks()
+	if concurrency == 0 {
+		concurrency = len(cp)
+	}
 
-	return r.out, nil
+	w.runMonitoringForGroup(ctx, cp, concurrency)
+
+	return w.running.out, nil
 }
 
 // Stop stops execution of checks.
 // Subsequent calls of Stop for the same groupID
-// return ErrNoSuchGroup.
+// return ErrNotRunning.
 func (w *Watchdog) Stop() error {
 	w.mu.Lock()
 	defer w.mu.Unlock()
@@ -134,10 +137,7 @@ func (w *Watchdog) RunImmediately(ctx context.Context, concurrency int) ([]Check
 		return nil, ErrNotConfigured
 	}
 
-	// making a copy here because mutex should not be
-	// held while checks are running
-	cp := make([]Check, len(w.checks))
-	copy(cp, w.checks)
+	cp := w.copyChecks()
 	w.mu.Unlock() // release
 
 	if concurrency == 0 {
@@ -158,20 +158,34 @@ func (w *Watchdog) RunImmediately(ctx context.Context, concurrency int) ([]Check
 	return statuses, nil
 }
 
-func runMonitoringForGroup(ctx context.Context, ch []Check) *running {
-	events := make(chan CheckResult)
+func (w *Watchdog) copyChecks() []Check {
+	cp := make([]Check, len(w.checks))
+	copy(cp, w.checks)
+
+	return cp
+}
+
+func (w *Watchdog) runMonitoringForGroup(ctx context.Context, checks []Check, concurrency int) {
+	events := make(chan CheckResult, len(checks))
 	stop := make(chan struct{})
 
-	grp := errgroup.Group{}
+	w.running = &running{out: events, stop: stop}
 
-	for _, c := range ch {
+	grp := errgroup.Group{}
+	sema := make(chan struct{}, concurrency)
+
+	for _, c := range checks {
 		grp.Go(func() error {
 			state := CheckResult{}
 			ticker := time.Tick(c.Interval)
 
 			for {
+				sema <- struct{}{}
+
 				status, err := c.Check(ctx)
 
+				<-sema
+
 				s := CheckResult{
 					Name:   c.Name,
 					Status: status,
@@ -201,9 +215,13 @@ func runMonitoringForGroup(ctx context.Context, ch []Check) *running {
 	go func() {
 		grp.Wait()
 		close(events)
-	}()
 
-	return &running{out: events, stop: stop}
+		w.mu.Lock()
+		defer w.mu.Unlock()
+		if w.running != nil {
+			w.running = nil
+		}
+	}()
 }
 
 func runChecksConcurrently(ctx context.Context, ch []Check, concurrency int) []CheckResult {
@@ -221,6 +239,8 @@ func runChecksConcurrently(ctx context.Context, ch []Check, concurrency int) []C
 			defer func() {
 			}()
 
+			// relying on fact that CheckFunc obeys context
+			// cancellation
 			status, err := e.Check(ctx)
 
 			r := CheckResult{