pkg/beholder: add Emitter.Close; track goroutines

Author: jmank88

pkg/beholder/beholdertest/beholder.go

@@ -169,6 +169,8 @@ type assertMessageEmitter struct {
 	msgs []beholder.Message
 }
 
+func (e *assertMessageEmitter) Close() error { return nil }
+
 func (e *assertMessageEmitter) Emit(_ context.Context, msg []byte, attrKVs ...any) error {
 	e.t.Helper()
 

pkg/beholder/chip_client.go

@@ -3,12 +3,15 @@ package beholder
 import (
 	"context"
 	"fmt"
+	"io"
+
 	"github.com/smartcontractkit/chainlink-common/pkg/chipingress"
 	"github.com/smartcontractkit/chainlink-common/pkg/chipingress/pb"
 )
 
 type ChipIngressClient interface {
 	RegisterSchema(ctx context.Context, schemas ...*pb.Schema) (map[string]int, error)
+	io.Closer
 }
 
 type chipIngressClient struct {
@@ -24,6 +27,7 @@ func NewChipIngressClient(client chipingress.Client) (ChipIngressClient, error)
 		client: client,
 	}, nil
 }
+func (sr *chipIngressClient) Close() error { return nil }
 
 // RegisterSchema registers one or more schemas with the Chip Ingress service. Returns a map of subject to version for each registered schema.
 func (sr *chipIngressClient) RegisterSchema(ctx context.Context, schemas ...*pb.Schema) (map[string]int, error) {

pkg/beholder/chip_ingress_emitter.go

@@ -21,6 +21,10 @@ func NewChipIngressEmitter(client chipingress.Client) (Emitter, error) {
 	return &ChipIngressEmitter{client: client}, nil
 }
 
+func (c *ChipIngressEmitter) Close() error {
+	return c.client.Close()
+}
+
 func (c *ChipIngressEmitter) Emit(ctx context.Context, body []byte, attrKVs ...any) error {
 
 	sourceDomain, entityType, err := ExtractSourceAndType(attrKVs...)

pkg/beholder/client.go

@@ -5,6 +5,7 @@ import (
 	"encoding/hex"
 	"errors"
 	"fmt"
+	"io"
 	"time"
 
 	"go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc"
@@ -30,6 +31,7 @@ import (
 type Emitter interface {
 	// Sends message with bytes and attributes to OTel Collector
 	Emit(ctx context.Context, body []byte, attrKVs ...any) error
+	io.Closer
 }
 
 type Client struct {
@@ -250,8 +252,14 @@ func NewGRPCClient(cfg Config, otlploggrpcNew otlploggrpcFactory) (*Client, erro
 
 // Closes all providers, flushes all data and stops all background processes
 func (c Client) Close() (err error) {
+	if c.Chip != nil {
+		err = errors.Join(err, c.Chip.Close())
+	}
+	if c.Emitter != nil {
+		err = errors.Join(err, c.Emitter.Close())
+	}
 	if c.OnClose != nil {
-		return c.OnClose()
+		err = errors.Join(err, c.OnClose())
 	}
 	return
 }

pkg/beholder/dual_source_emitter.go

@@ -2,19 +2,25 @@ package beholder
 
 import (
 	"context"
+	"errors"
 	"fmt"
+	"sync/atomic"
 
 	"github.com/smartcontractkit/chainlink-common/pkg/logger"
+	"github.com/smartcontractkit/chainlink-common/pkg/services"
 )
 
-// dualSourceEmitter emits both to chip ingress and to the otel collector
+// DualSourceEmitter emits both to chip ingress and to the otel collector
 // this is to help transition from sending custom messages via OTLP to instead use chip-ingress
 // we want to send to both during the transition period, then cutover to using
 // chipIngressEmitter only
 type DualSourceEmitter struct {
 	chipIngressEmitter   Emitter
 	otelCollectorEmitter Emitter
 	log                  logger.Logger
+	stopCh               services.StopChan
+	wg                   services.WaitGroup
+	closed               atomic.Bool
 }
 
 func NewDualSourceEmitter(chipIngressEmitter Emitter, otelCollectorEmitter Emitter) (Emitter, error) {
@@ -36,9 +42,19 @@ func NewDualSourceEmitter(chipIngressEmitter Emitter, otelCollectorEmitter Emitt
 		chipIngressEmitter:   chipIngressEmitter,
 		otelCollectorEmitter: otelCollectorEmitter,
 		log:                  logger,
+		stopCh:               make(services.StopChan),
 	}, nil
 }
 
+func (d *DualSourceEmitter) Close() error {
+	if wasClosed := d.closed.Swap(true); wasClosed {
+		return errors.New("already closed")
+	}
+	close(d.stopCh)
+	d.wg.Wait()
+	return errors.Join(d.chipIngressEmitter.Close(), d.otelCollectorEmitter.Close())
+}
+
 func (d *DualSourceEmitter) Emit(ctx context.Context, body []byte, attrKVs ...any) error {
 
 	// Emit via OTLP first
@@ -47,13 +63,21 @@ func (d *DualSourceEmitter) Emit(ctx context.Context, body []byte, attrKVs ...an
 	}
 
 	// Emit via chip ingress async
-	go func() {
-		if err := d.chipIngressEmitter.Emit(context.WithoutCancel(ctx), body, attrKVs...); err != nil {
+	if err := d.wg.TryAdd(1); err != nil {
+		return err
+	}
+	go func(ctx context.Context) {
+		defer d.wg.Done()
+		var cancel context.CancelFunc
+		ctx, cancel = d.stopCh.Ctx(ctx)
+		defer cancel()
+
+		if err := d.chipIngressEmitter.Emit(ctx, body, attrKVs...); err != nil {
 			// If the chip ingress emitter fails, we ONLY log the error
 			// because we still want to send the data to the OTLP collector and not cause disruption
 			d.log.Infof("failed to emit to chip ingress: %v", err)
 		}
-	}()
+	}(context.WithoutCancel(ctx))
 
 	return nil
 }

pkg/beholder/dual_source_emitter_test.go

@@ -5,9 +5,10 @@ import (
 	"fmt"
 	"testing"
 
-	"github.com/smartcontractkit/chainlink-common/pkg/beholder"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
+
+	"github.com/smartcontractkit/chainlink-common/pkg/beholder"
 )
 
 func TestNewDualSourceEmitter(t *testing.T) {
@@ -80,6 +81,8 @@ type mockEmitter struct {
 	emitFunc func(ctx context.Context, body []byte, attrKVs ...any) error
 }
 
+func (m *mockEmitter) Close() error { return nil }
+
 func (m *mockEmitter) Emit(ctx context.Context, body []byte, attrKVs ...any) error {
 	if m.emitFunc != nil {
 		return m.emitFunc(ctx, body, attrKVs...)

pkg/beholder/message_emitter.go

@@ -17,6 +17,8 @@ func NewMessageEmitter(logger otellog.Logger) Emitter {
 	}
 }
 
+func (e messageEmitter) Close() error { return nil }
+
 // Emits logs the message, but does not wait for the message to be processed.
 // Open question: what are pros/cons for using use map[]any vs use otellog.KeyValue
 func (e messageEmitter) Emit(ctx context.Context, body []byte, attrKVs ...any) error {

pkg/beholder/noop.go

@@ -99,6 +99,8 @@ func NewWriterClient(w io.Writer) (*Client, error) {
 
 type noopMessageEmitter struct{}
 
+func (e noopMessageEmitter) Close() error { return nil }
+
 func (noopMessageEmitter) Emit(ctx context.Context, body []byte, attrKVs ...any) error {
 	return nil
 }

pkg/capabilities/events/events_test.go

@@ -17,6 +17,8 @@ type testEmitter struct {
 	attrs   []any
 }
 
+func (t *testEmitter) Close() error { return nil }
+
 func (t *testEmitter) Emit(ctx context.Context, payload []byte, attrKVs ...any) error {
 	t.payload = payload
 	t.attrs = attrKVs