Alternative implementation for closing a beat.Client

CHANGELOG-developer.next.asciidoc

@@ -46,3 +46,4 @@ The list below covers the major changes between 7.0.0-rc2 and master only.
 - Add ClientFactory to TCP input source to add SplitFunc/NetworkFuncs per client. {pull}8543[8543]
 - Introduce beat.OutputChooses publisher mode. {pull}12996[12996]
 - Ensure that beat.Processor, beat.ProcessorList, and processors.ProcessorList are compatible and can be composed more easily. {pull}12996[12996]
+- Add support to close beat.Client via beat.CloseRef (a subset of context.Context). {pull}13031[13031]

libbeat/beat/pipeline.go

@@ -48,6 +48,8 @@ type ClientConfig struct {
 
 	Processing ProcessingConfig
 
+	CloseRef CloseRef
+
 	// WaitClose sets the maximum duration to wait on ACK, if client still has events
 	// active non-acknowledged events in the publisher pipeline.
 	// WaitClose is only effective if one of ACKCount, ACKEvents and ACKLastEvents
@@ -78,6 +80,13 @@ type ClientConfig struct {
 	ACKLastEvent func(interface{})
 }
 
+// CloseRef allows users to close the client asynchronously.
+// A CloseReg implements a subset of function required for context.Context.
+type CloseRef interface {
+	Done() <-chan struct{}
+	Err() error
+}
+
 // ProcessingConfig provides additional event processing settings a client can
 // pass to the publisher pipeline on Connect.
 type ProcessingConfig struct {

libbeat/publisher/pipeline/acker.go

@@ -30,6 +30,7 @@ import (
 // All pipeline and client ACK handling support is provided by acker instances.
 type acker interface {
 	close()
+	wait()
 	addEvent(event beat.Event, published bool) bool
 	ackEvents(int)
 }
@@ -40,6 +41,7 @@ type emptyACK struct{}
 var nilACKer acker = (*emptyACK)(nil)
 
 func (*emptyACK) close()                             {}
+func (*emptyACK) wait()                              {}
 func (*emptyACK) addEvent(_ beat.Event, _ bool) bool { return true }
 func (*emptyACK) ackEvents(_ int)                    {}
 
@@ -68,6 +70,7 @@ func newCountACK(pipeline *Pipeline, fn func(total, acked int)) *countACK {
 }
 
 func (a *countACK) close()                             {}
+func (a *countACK) wait()                              {}
 func (a *countACK) addEvent(_ beat.Event, _ bool) bool { return true }
 func (a *countACK) ackEvents(n int) {
 	if a.pipeline.ackActive.Load() {
@@ -220,6 +223,8 @@ func (a *gapCountACK) close() {
 	close(a.done)
 }
 
+func (a *gapCountACK) wait() {}
+
 func (a *gapCountACK) addEvent(_ beat.Event, published bool) bool {
 	// if gapList is empty and event is being dropped, forward drop event to ack
 	// loop worker:
@@ -313,9 +318,8 @@ func newBoundGapCountACK(
 	return a
 }
 
-func (a *boundGapCountACK) close() {
-	a.acker.close()
-}
+func (a *boundGapCountACK) close() { a.acker.close() }
+func (a *boundGapCountACK) wait()  { a.acker.wait() }
 
 func (a *boundGapCountACK) addEvent(event beat.Event, published bool) bool {
 	a.sema.inc()
@@ -361,9 +365,9 @@ func makeCountACK(pipeline *Pipeline, canDrop bool, sema *sema, fn func(int, int
 	return newCountACK(pipeline, fn)
 }
 
-func (a *eventDataACK) close() {
-	a.acker.close()
-}
+func (a *eventDataACK) close() { a.acker.close() }
+
+func (a *eventDataACK) wait() { a.acker.wait() }
 
 func (a *eventDataACK) addEvent(event beat.Event, published bool) bool {
 	a.mutex.Lock()
@@ -400,37 +404,57 @@ func (a *eventDataACK) onACK(total, acked int) {
 type waitACK struct {
 	acker acker
 
-	signal    chan struct{}
-	waitClose time.Duration
+	signalAll  chan struct{} // ack loop notifies `close` that all events have been acked
+	signalDone chan struct{} // shutdown handler telling `wait` that shutdown has been completed
+	waitClose  time.Duration
 
 	active atomic.Bool
 
 	// number of active events
 	events atomic.Uint64
+
+	afterClose func()
 }
 
-func newWaitACK(acker acker, timeout time.Duration) *waitACK {
+func newWaitACK(acker acker, timeout time.Duration, afterClose func()) *waitACK {
 	return &waitACK{
-		acker:     acker,
-		signal:    make(chan struct{}, 1),
-		waitClose: timeout,
-		active:    atomic.MakeBool(true),
+		acker:      acker,
+		signalAll:  make(chan struct{}, 1),
+		signalDone: make(chan struct{}),
+		waitClose:  timeout,
+		active:     atomic.MakeBool(true),
+		afterClose: afterClose,
 	}
 }
 
 func (a *waitACK) close() {
-	// TODO: wait for events
-
 	a.active.Store(false)
-	if a.events.Load() > 0 {
+
+	if a.events.Load() == 0 {
+		a.finishClose()
+		return
+	}
+
+	// start routine to propagate shutdown signals or timeouts to anyone
+	// being blocked in wait.
+	go func() {
+		defer a.finishClose()
+
 		select {
-		case <-a.signal:
+		case <-a.signalAll:
 		case <-time.After(a.waitClose):
 		}
-	}
+	}()
+}
 
-	// close the underlying acker upon exit
+func (a *waitACK) finishClose() {
 	a.acker.close()
+	a.afterClose()
+	close(a.signalDone)
+}
+
+func (a *waitACK) wait() {
+	<-a.signalDone
 }
 
 func (a *waitACK) addEvent(event beat.Event, published bool) bool {
@@ -454,6 +478,22 @@ func (a *waitACK) releaseEvents(n int) {
 
 	// send done signal, if close is waiting
 	if !a.active.Load() {
-		a.signal <- struct{}{}
+		a.signalAll <- struct{}{}
 	}
 }
+
+// closeACKer simply wraps any other acker. It calls a custom function after
+// the underlying acker has been closed.
+type closeACKer struct {
+	acker
+	afterClose func()
+}
+
+func newCloseACKer(a acker, fn func()) acker {
+	return &closeACKer{acker: a, afterClose: fn}
+}
+
+func (a closeACKer) close() {
+	a.acker.close()
+	a.afterClose()
+}

libbeat/publisher/pipeline/client.go

@@ -43,7 +43,11 @@ type client struct {
 	canDrop      bool
 	reportEvents bool
 
-	isOpen atomic.Bool
+	// Open state, signaling, and sync primitives for coordinating client Close.
+	isOpen    atomic.Bool   // set to false during shutdown, such that no new events will be accepted anymore.
+	closeOnce sync.Once     // closeOnce ensure that the client shutdown sequence is only executed once
+	closeRef  beat.CloseRef // extern closeRef for sending a signal that the client should be closed.
+	done      chan struct{} // the done channel will be closed if the closeReg gets closed, or Close is run.
 
 	eventer beat.ClientEventer
 }
@@ -135,23 +139,40 @@ func (c *client) publish(e beat.Event) {
 }
 
 func (c *client) Close() error {
-	// first stop ack handling. ACK handler might block (with timeout), waiting
+	log := c.logger()
+
+	// first stop ack handling. ACK handler might block on wait (with timeout), waiting
 	// for pending events to be ACKed.
+	c.doClose()
+	log.Debug("client: wait for acker to finish")
+	c.acker.wait()
+	log.Debug("client: acker shut down")
+	return nil
+}
 
-	log := c.logger()
+func (c *client) doClose() {
+	c.closeOnce.Do(func() {
+		close(c.done)
 
-	if !c.isOpen.Swap(false) {
-		return nil // closed or already closing
-	}
+		log := c.logger()
+
+		c.isOpen.Store(false)
+		c.onClosing()
 
-	c.onClosing()
+		log.Debug("client: closing acker")
+		c.acker.close() // this must trigger a direct/indirect call to 'unlink'
+	})
+}
 
-	log.Debug("client: closing acker")
-	c.acker.close()
+// unlink is the final step of closing a client. It must be executed only after
+// it is guaranteed that the underlying acker has been closed and will not
+// accept any new publish or ACK events.
+// This method is normally registered with the ACKer and triggered by it.
+func (c *client) unlink() {
+	log := c.logger()
 	log.Debug("client: done closing acker")
 
-	// finally disconnect client from broker
-	n := c.producer.Cancel()
+	n := c.producer.Cancel() // close connection to queue
 	log.Debugf("client: cancelled %v events", n)
 
 	if c.reportEvents {
@@ -162,7 +183,6 @@ func (c *client) Close() error {
 	}
 
 	c.onClosed()
-	return nil
 }
 
 func (c *client) logger() *logp.Logger {

libbeat/publisher/pipeline/client_ack.go

@@ -33,6 +33,7 @@ func (p *Pipeline) makeACKer(
 	canDrop bool,
 	cfg *beat.ClientConfig,
 	waitClose time.Duration,
+	afterClose func(),
 ) acker {
 	var (
 		bld   = p.ackBuilder
@@ -50,15 +51,16 @@ func (p *Pipeline) makeACKer(
 		acker = bld.createEventACKer(canDrop, sema, cb)
 	default:
 		if waitClose <= 0 {
-			return bld.createPipelineACKer(canDrop, sema)
+			acker = bld.createPipelineACKer(canDrop, sema)
+		} else {
+			acker = bld.createCountACKer(canDrop, sema, func(_ int) {})
 		}
-		acker = bld.createCountACKer(canDrop, sema, func(_ int) {})
 	}
 
 	if waitClose <= 0 {
-		return acker
+		return newCloseACKer(acker, afterClose)
 	}
-	return newWaitACK(acker, waitClose)
+	return newWaitACK(acker, waitClose, afterClose)
 }
 
 func lastEventACK(fn func(interface{})) func([]interface{}) {

libbeat/publisher/pipeline/client_test.go

@@ -0,0 +1,115 @@
+// 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 pipeline
+
+import (
+	"context"
+	"sync"
+	"testing"
+
+	"github.com/elastic/beats/libbeat/beat"
+	"github.com/elastic/beats/libbeat/logp"
+	"github.com/elastic/beats/libbeat/outputs"
+	"github.com/elastic/beats/libbeat/publisher/queue"
+	"github.com/elastic/beats/libbeat/tests/resources"
+)
+
+func TestClient(t *testing.T) {
+	makePipeline := func(settings Settings, qu queue.Queue) *Pipeline {
+		p, err := New(beat.Info{},
+			Monitors{},
+			func(_ queue.Eventer) (queue.Queue, error) {
+				return qu, nil
+			},
+			outputs.Group{},
+			settings,
+		)
+		if err != nil {
+			panic(err)
+		}
+
+		return p
+	}
+
+	t.Run("client close", func(t *testing.T) {
+		// Note: no asserts. If closing fails we have a deadlock, because Publish
+		// would block forever
+
+		cases := map[string]struct {
+			context bool
+			close   func(client beat.Client, cancel func())
+		}{
+			"close unblocks client without context": {
+				context: false,
+				close: func(client beat.Client, _ func()) {
+					client.Close()
+				},
+			},
+			"close unblocks client with context": {
+				context: true,
+				close: func(client beat.Client, _ func()) {
+					client.Close()
+				},
+			},
+			"context cancel unblocks client": {
+				context: true,
+				close: func(client beat.Client, cancel func()) {
+					cancel()
+				},
+			},
+		}
+
+		if testing.Verbose() {
+			logp.TestingSetup()
+		}
+
+		for name, test := range cases {
+			t.Run(name, func(t *testing.T) {
+				routinesChecker := resources.NewGoroutinesChecker()
+				defer routinesChecker.Check(t)
+
+				pipeline := makePipeline(Settings{}, makeBlockingQueue())
+				defer pipeline.Close()
+
+				var ctx context.Context
+				var cancel func()
+				if test.context {
+					ctx, cancel = context.WithCancel(context.Background())
+				}
+
+				client, err := pipeline.ConnectWith(beat.ClientConfig{
+					CloseRef: ctx,
+				})
+				if err != nil {
+					t.Fatal(err)
+				}
+				defer client.Close()
+
+				var wg sync.WaitGroup
+				wg.Add(1)
+				go func() {
+					defer wg.Done()
+					client.Publish(beat.Event{})
+				}()
+
+				test.close(client, cancel)
+				wg.Wait()
+			})
+		}
+	})
+}