[Heartbeat] Add publish pipeline timeout to run_once

heartbeat/beater/heartbeat.go

@@ -59,6 +59,7 @@ type Heartbeat struct {
 	autodiscover       *autodiscover.Autodiscover
 	replaceStateLoader func(sl monitorstate.StateLoader)
 	trace              tracer.Tracer
+	pipeline           beat.Pipeline
 }
 
 // New creates a new heartbeat.
@@ -118,16 +119,8 @@ func New(b *beat.Beat, rawConfig *conf.C) (beat.Beater, error) {
 	sched := scheduler.Create(limit, hbregistry.SchedulerRegistry, location, jobConfig, parsedConfig.RunOnce)
 
 	pipelineClientFactory := func(p beat.Pipeline) (pipeline.ISyncClient, error) {
-		if parsedConfig.RunOnce {
-			client, err := pipeline.NewSyncClient(logp.L(), p, beat.ClientConfig{})
-			if err != nil {
-				return nil, fmt.Errorf("could not create pipeline sync client for run_once: %w", err)
-			}
-			return client, nil
-		} else {
-			client, err := p.Connect()
-			return monitors.SyncPipelineClientAdaptor{C: client}, err
-		}
+		client, err := p.Connect()
+		return monitors.SyncPipelineClientAdaptor{C: client}, err
 	}
 
 	bt := &Heartbeat{
@@ -145,7 +138,8 @@ func New(b *beat.Beat, rawConfig *conf.C) (beat.Beater, error) {
 			PipelineClientFactory: pipelineClientFactory,
 			BeatRunFrom:           parsedConfig.RunFrom,
 		}),
-		trace: trace,
+		trace:    trace,
+		pipeline: b.Publisher,
 	}
 	runFromID := "<unknown location>"
 	if parsedConfig.RunFrom != nil {
@@ -160,6 +154,15 @@ func (bt *Heartbeat) Run(b *beat.Beat) error {
 	bt.trace.Start()
 	defer bt.trace.Close()
 
+	waitMonitors := newSignalWait()
+	waitPublished := newSignalWait()
+	defer waitPublished.Wait()
+
+	pipelineWrapper := &PipelineClientWrapper{log: logp.L().Named("run_once pipeline wrapper")}
+	if bt.config.RunOnce {
+		bt.pipeline = withPipelineWrapper(bt.pipeline, pipelineWrapper)
+	}
+
 	logp.L().Info("heartbeat is running! Hit CTRL-C to stop it.")
 	groups, _ := syscall.Getgroups()
 	logp.L().Info("Effective user/group ids: %d/%d, with groups: %v", syscall.Geteuid(), syscall.Getegid(), groups)
@@ -174,9 +177,11 @@ func (bt *Heartbeat) Run(b *beat.Beat) error {
 	defer stopStaticMonitors()
 
 	if bt.config.RunOnce {
-		bt.scheduler.WaitForRunOnce()
-		logp.L().Info("Ending run_once run")
-		return nil
+		runOnce := func() {
+			bt.scheduler.WaitForRunOnce()
+			logp.L().Info("Ending run_once run")
+		}
+		waitMonitors.Add(runOnce)
 	}
 
 	if b.Manager.Enabled() {
@@ -211,11 +216,24 @@ func (bt *Heartbeat) Run(b *beat.Beat) error {
 
 	defer bt.scheduler.Stop()
 
-	<-bt.done
+	// Wait until run_once ends or bt is being shutdown
+	waitMonitors.AddChan(bt.done)
+	waitMonitors.Wait()
 
-	if err != nil {
-		logp.L().Errorf("could not write trace stop event: %s", err)
+	// Checks if on shutdown it should wait for all events to be published
+	if bt.config.RunOnce {
+		waitPublished.Add(withLog(pipelineWrapper.Wait,
+			// Wait for either timeout or all events having been ACKed by outputs.
+			"Continue shutdown: All enqueued events being published."))
+		if bt.config.RunOnceTimeout > 0 {
+			logp.Info("Shutdown output timer started. Waiting for max %v.", bt.config.RunOnceTimeout)
+			waitPublished.Add(withLog(waitDuration(bt.config.RunOnceTimeout),
+				"Continue shutdown: Time out waiting for events being published."))
+		}
+	} else {
+		waitPublished.AddChan(bt.done)
 	}
+
 	logp.L().Info("Shutting down.")
 	return nil
 }
@@ -224,7 +242,7 @@ func (bt *Heartbeat) Run(b *beat.Beat) error {
 func (bt *Heartbeat) RunStaticMonitors(b *beat.Beat) (stop func(), err error) {
 	runners := make([]cfgfile.Runner, 0, len(bt.config.Monitors))
 	for _, cfg := range bt.config.Monitors {
-		created, err := bt.monitorFactory.Create(b.Publisher, cfg)
+		created, err := bt.monitorFactory.Create(bt.pipeline, cfg)
 		if err != nil {
 			if errors.Is(err, monitors.ErrMonitorDisabled) {
 				logp.L().Info("skipping disabled monitor: %s", err)

heartbeat/beater/pipeline_wrapper.go

@@ -0,0 +1,76 @@
+// Licensed to Elasticsearch B.V. under one or more contributor
+// license agreements. See the NOTICE file distributed with
+// this work for additional information regarding copyright
+// ownership. Elasticsearch B.V. licenses this file to you under
+// the Apache License, Version 2.0 (the "License"); you may
+// not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+package beater
+
+import (
+	"sync"
+
+	"github.com/elastic/beats/v7/libbeat/beat"
+	"github.com/elastic/beats/v7/libbeat/common/acker"
+	"github.com/elastic/beats/v7/libbeat/publisher/pipetool"
+	"github.com/elastic/elastic-agent-libs/logp"
+)
+
+type PipelineClientWrapper struct {
+	wg  sync.WaitGroup
+	log *logp.Logger
+}
+
+type wrappedClient struct {
+	wg     *sync.WaitGroup
+	client beat.Client
+}
+
+func withPipelineWrapper(pipeline beat.Pipeline, pw *PipelineClientWrapper) beat.Pipeline {
+	pipeline = pipetool.WithACKer(pipeline, acker.TrackingCounter(func(_, total int) {
+		pw.log.Debugf("ack callback receives with events count of %d", total)
+		pw.onACK(total)
+	}))
+
+	pipeline = pipetool.WithClientWrapper(pipeline, func(client beat.Client) beat.Client {
+		return &wrappedClient{
+			wg:     &pw.wg,
+			client: client,
+		}
+	})
+
+	return pipeline
+}
+
+func (c *wrappedClient) Publish(event beat.Event) {
+	c.wg.Add(1)
+	c.client.Publish(event)
+}
+
+func (c *wrappedClient) PublishAll(events []beat.Event) {
+	c.wg.Add(len(events))
+	c.client.PublishAll(events)
+}
+
+func (c *wrappedClient) Close() error {
+	return c.client.Close()
+}
+
+// Wait waits until we received a ACK for every events that were sent
+func (s *PipelineClientWrapper) Wait() {
+	s.wg.Wait()
+}
+
+func (s *PipelineClientWrapper) onACK(n int) {
+	s.wg.Add(-1 * n)
+}

heartbeat/beater/signalwait.go

@@ -0,0 +1,90 @@
+// Licensed to Elasticsearch B.V. under one or more contributor
+// license agreements. See the NOTICE file distributed with
+// this work for additional information regarding copyright
+// ownership. Elasticsearch B.V. licenses this file to you under
+// the Apache License, Version 2.0 (the "License"); you may
+// not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+package beater
+
+import (
+	"time"
+
+	"github.com/elastic/elastic-agent-libs/logp"
+)
+
+type signalWait struct {
+	count   int // number of potential 'alive' signals
+	signals chan struct{}
+}
+
+type signaler func()
+
+func newSignalWait() *signalWait {
+	return &signalWait{
+		signals: make(chan struct{}, 1),
+	}
+}
+
+func (s *signalWait) Wait() {
+	if s.count == 0 {
+		return
+	}
+
+	<-s.signals
+	s.count--
+}
+
+func (s *signalWait) Add(fn signaler) {
+	s.count++
+	go func() {
+		fn()
+		var v struct{}
+		s.signals <- v
+	}()
+}
+
+func (s *signalWait) AddChan(c <-chan struct{}) {
+	s.Add(waitChannel(c))
+}
+
+func (s *signalWait) AddTimer(t *time.Timer) {
+	s.Add(waitTimer(t))
+}
+
+func (s *signalWait) AddTimeout(d time.Duration) {
+	s.Add(waitDuration(d))
+}
+
+func (s *signalWait) Signal() {
+	s.Add(func() {})
+}
+
+func waitChannel(c <-chan struct{}) signaler {
+	return func() { <-c }
+}
+
+func waitTimer(t *time.Timer) signaler {
+	return func() { <-t.C }
+}
+
+func waitDuration(d time.Duration) signaler {
+	return waitTimer(time.NewTimer(d))
+}
+
+func withLog(s signaler, msg string) signaler {
+	return func() {
+		s()
+		logp.Info("%v", msg)
+	}
+}

heartbeat/config/config.go

@@ -41,6 +41,7 @@ type LocationWithID struct {
 // Config defines the structure of heartbeat.yml.
 type Config struct {
 	RunOnce        bool                 `config:"run_once"`
+	RunOnceTimeout time.Duration        `config:"run_once_timeout"`
 	Monitors       []*conf.C            `config:"monitors"`
 	ConfigMonitors *conf.C              `config:"config.monitors"`
 	Scheduler      Scheduler            `config:"scheduler"`