[Heartbeat] Monitor Retries

heartbeat/monitors/jobs/job.go

@@ -73,3 +73,19 @@ func Wrap(job Job, wrapper JobWrapper) Job {
  return WrapAll(cont, wrapper), err
  }
 }
+
+type StatefulWrapper[T any] interface {
+ Wrap(j Job) Job
+}
+
+type StatefulWrapperFactory[T any] func(rootJob Job) StatefulWrapper[T]
+
+func WrapStateful[T StatefulWrapper[T]](makeSW StatefulWrapperFactory[T]) JobWrapper {
+ return func(j Job) Job {
+ return func(event *beat.Event) ([]Job, error) {
+ sw := makeSW(j)
+ conts, err := sw.Wrap(j)(event)
+ return WrapAll(conts, sw.Wrap), err
+ }
+ }
+}

heartbeat/monitors/jobs/job_test.go

@@ -20,6 +20,7 @@ package jobs
 import (
  "testing"
 
+ "github.com/stretchr/testify/assert"
  "github.com/stretchr/testify/require"
 
  "github.com/elastic/beats/v7/heartbeat/eventext"
@@ -29,6 +30,67 @@ import (
  "github.com/elastic/go-lookslike/testslike"
 )
 
+type TestSWrapper struct {
+ rootCounter int
+ contCounter int
+}
+
+func (tsw *TestSWrapper) Wrap(j Job) Job {
+ return func(event *beat.Event) ([]Job, error) {
+ tsw.contCounter++
+ eventext.MergeEventFields(event, mapstr.M{"cont_counter": tsw.contCounter, "root_counter": tsw.rootCounter})
+ return j(event)
+ }
+}
+
+func TestStatefulWrapper(t *testing.T) {
+ innerj1Fields := mapstr.M{
+ "innerjob": "innervalue",
+ "innercount": 1,
+ }
+ innerj2Fields := mapstr.M{
+ "innerjob": "innervalue",
+ "innercount": 2,
+ }
+
+ var j Job = func(event *beat.Event) ([]Job, error) {
+ eventext.MergeEventFields(event, innerj1Fields)
+
+ return []Job{
+ func(event *beat.Event) ([]Job, error) {
+ eventext.MergeEventFields(event, innerj2Fields)
+ return nil, nil
+ },
+ }, nil
+ }
+
+ wrapper := WrapStateful[*TestSWrapper](func(rootJob Job) StatefulWrapper[*TestSWrapper] { return &TestSWrapper{rootCounter: 1} })
+
+ // Run this ten times to ensure state is not carried across retries
+ for i := 0; i < 10; i++ {
+ events, err := ExecJobAndConts(t, wrapper(j))
+ assert.NoError(t, err)
+ assert.Len(t, events, 2)
+ testslike.Test(
+ t,
+ lookslike.Compose(
+ lookslike.MustCompile(innerj1Fields),
+ lookslike.MustCompile(mapstr.M{"cont_counter": 1, "root_counter": 1}),
+ ),
+ events[0].Fields,
+ )
+ testslike.Test(
+ t,
+ lookslike.Compose(
+ lookslike.MustCompile(innerj2Fields),
+ lookslike.MustCompile(mapstr.M{"cont_counter": 2, "root_counter": 1}),
+ ),
+ events[1].Fields,
+ )
+ }
+
+}
+
 func TestWrapAll(t *testing.T) {
  type args struct {
  jobs []Job

heartbeat/monitors/monitor.go

@@ -70,6 +70,8 @@ type Monitor struct {
  // stats is the countersRecorder used to record lifecycle events
  // for global metrics + telemetry
  stats plugin.RegistryRecorder
+
+ monitorStateTracker *monitorstate.Tracker
 }
 
 // String prints a description of the monitor in a threadsafe way. It is important that this use threadsafe
@@ -125,15 +127,16 @@ func newMonitorUnsafe(
  }
 
  m := &Monitor{
- stdFields: standardFields,
- pluginName: pluginFactory.Name,
- addTask: addTask,
- configuredJobs: []*configuredJob{},
- pubClient: pubClient,
- internalsMtx: sync.Mutex{},
- config: config,
- stats: pluginFactory.Stats,
- state: MON_INIT,
+ stdFields: standardFields,
+ pluginName: pluginFactory.Name,
+ addTask: addTask,
+ configuredJobs: []*configuredJob{},
+ pubClient: pubClient,
+ internalsMtx: sync.Mutex{},
+ config: config,
+ stats: pluginFactory.Stats,
+ state: MON_INIT,
+ monitorStateTracker: monitorstate.NewTracker(stateLoader, false),
  }
 
  if m.stdFields.ID == "" {
@@ -156,7 +159,7 @@ func newMonitorUnsafe(
 
  var wrappedJobs []jobs.Job
  if err == nil {
- wrappedJobs = wrappers.WrapCommon(p.Jobs, m.stdFields, stateLoader)
+ wrappedJobs = wrappers.WrapCommon(p.Jobs, m.stdFields, stateLoader, 2)
  } else {
  // If we've hit an error at this point, still run on schedule, but always return an error.
  // This way the error is clearly communicated through to kibana.
@@ -179,7 +182,7 @@ func newMonitorUnsafe(
  // We need to use the lightweight wrapping for error jobs
  // since browser wrapping won't write summaries, but the fake job here is
  // effectively a lightweight job
- wrappedJobs = wrappers.WrapLightweight(p.Jobs, m.stdFields, monitorstate.NewTracker(stateLoader, false))
+ wrappedJobs = wrappers.WrapLightweight(p.Jobs, m.stdFields, monitorstate.NewTracker(stateLoader, false), 2)
  }
 
  m.endpoints = p.Endpoints

heartbeat/monitors/plugin/plugin.go

@@ -60,7 +60,7 @@ func (p Plugin) Close() error {
 
 // RunWrapped runs the plug-in with the provided wrappers returning a channel of resultant events.
 func (p Plugin) RunWrapped(fields stdfields.StdMonitorFields) chan *beat.Event {
- wj := wrappers.WrapCommon(p.Jobs, fields, nil)
+ wj := wrappers.WrapCommon(p.Jobs, fields, nil, 2)
  results := make(chan *beat.Event)
 
  var runJob func(j jobs.Job)

heartbeat/monitors/wrappers/monitorstate/esloader_test.go

@@ -51,7 +51,7 @@ func TestStatesESLoader(t *testing.T) {
 
  monID := etc.createTestMonitorStateInES(t, testStatus)
  // Since we've continued this state it should register the initial state
- ms := etc.tracker.getCurrentState(monID)
+ ms := etc.tracker.GetCurrentState(monID)
  require.True(t, ms.StartedAt.After(testStart.Add(-time.Nanosecond)), "timestamp for new state is off")
  requireMSStatusCount(t, ms, testStatus, 1)
 

heartbeat/monitors/wrappers/monitorstate/monitorstate.go

@@ -37,6 +37,8 @@ const (
  StatusUp StateStatus = "up"
  StatusDown StateStatus = "down"
  StatusFlapping StateStatus = "flap"
+ // Nil, essentially
+ StatusEmpty StateStatus = ""
 )
 
 func newMonitorState(sf stdfields.StdMonitorFields, status StateStatus, ctr int, flappingEnabled bool) *State {

heartbeat/monitors/wrappers/monitorstate/tracker.go

@@ -62,7 +62,7 @@ func (t *Tracker) RecordStatus(sf stdfields.StdMonitorFields, newStatus StateSta
  t.mtx.Lock()
  defer t.mtx.Unlock()
 
- state := t.getCurrentState(sf)
+ state := t.GetCurrentState(sf)
  if state == nil {
  state = newMonitorState(sf, newStatus, 0, t.flappingEnabled)
  logp.L().Infof("initializing new state for monitor %s: %s", sf.ID, state.String())
@@ -74,7 +74,15 @@ func (t *Tracker) RecordStatus(sf stdfields.StdMonitorFields, newStatus StateSta
  return state.copy()
 }
 
-func (t *Tracker) getCurrentState(sf stdfields.StdMonitorFields) (state *State) {
+func (t *Tracker) GetCurrentStatus(sf stdfields.StdMonitorFields) StateStatus {
+ s := t.GetCurrentState(sf)
+ if s == nil {
+ return StatusEmpty
+ }
+ return s.Status
+}
+
+func (t *Tracker) GetCurrentState(sf stdfields.StdMonitorFields) (state *State) {
  if state, ok := t.states[sf.ID]; ok {
  return state
  }