refactor: sink-worker lifecycle (#3647)

Author: chrisgacsal

app/common/kafka.go

@@ -132,3 +132,37 @@ func NewKafkaIngestNamespaceHandler(
 
 	return handler, nil
 }
+
+func NewKafkaConsumer(conf pkgkafka.ConsumerConfig, logger *slog.Logger) (*kafka.Consumer, func(), error) {
+	if err := conf.Validate(); err != nil {
+		return nil, nil, fmt.Errorf("invalid Kafka consumer configuration: %w", err)
+	}
+
+	consumerConfigMap, err := conf.AsConfigMap()
+	if err != nil {
+		return nil, nil, fmt.Errorf("failed to generate Kafka consumer configuration map: %w", err)
+	}
+
+	consumer, err := kafka.NewConsumer(&consumerConfigMap)
+	if err != nil {
+		return nil, nil, fmt.Errorf("failed to initialize Kafka consumer: %w", err)
+	}
+
+	kLogger := logger.WithGroup("kafka").WithGroup("consumer").With(
+		"group.id", conf.ConsumerGroupID,
+		"group.instance.id", conf.ConsumerGroupInstanceID,
+		"client.id", conf.ClientID,
+	)
+
+	// Enable Kafka client logging
+	// TODO: refactor ConsumeLogChannel to allow graceful shutdown
+	go pkgkafka.ConsumeLogChannel(consumer, kLogger)
+
+	closer := func() {
+		if err = consumer.Close(); err != nil {
+			kLogger.Error("failed to close Kafka consumer", slog.String("err", err.Error()))
+		}
+	}
+
+	return consumer, closer, nil
+}

app/common/openmeter_sinkworker.go

@@ -2,19 +2,37 @@ package common
 
 import (
 	"context"
+	"fmt"
 	"log/slog"
 
 	"github.com/ThreeDotsLabs/watermill/message"
+	"github.com/confluentinc/confluent-kafka-go/v2/kafka"
+	"github.com/google/wire"
 	"go.opentelemetry.io/otel/metric"
+	"go.opentelemetry.io/otel/trace"
 
 	"github.com/openmeterio/openmeter/app/config"
+	"github.com/openmeterio/openmeter/openmeter/dedupe"
+	"github.com/openmeterio/openmeter/openmeter/ingest/kafkaingest/topicresolver"
+	"github.com/openmeterio/openmeter/openmeter/meter"
+	"github.com/openmeterio/openmeter/openmeter/sink"
 	"github.com/openmeterio/openmeter/openmeter/sink/flushhandler"
 	"github.com/openmeterio/openmeter/openmeter/sink/flushhandler/ingestnotification"
+	"github.com/openmeterio/openmeter/openmeter/streaming"
 	watermillkafka "github.com/openmeterio/openmeter/openmeter/watermill/driver/kafka"
 	"github.com/openmeterio/openmeter/openmeter/watermill/eventbus"
 	pkgkafka "github.com/openmeterio/openmeter/pkg/kafka"
 )
 
+var Sink = wire.NewSet(
+	NewFlushHandler,
+	NewSinkWorkerPublisher,
+	NewSinkKafkaConsumer,
+	NewSinkDeduplicator,
+	NewSinkStorage,
+	NewSink,
+)
+
 // the sink-worker requires control over how the publisher is closed
 func NewSinkWorkerPublisher(
 	ctx context.Context,
@@ -27,13 +45,12 @@ func NewSinkWorkerPublisher(
 }
 
 func NewFlushHandler(
-	eventsConfig config.EventsConfiguration,
 	sinkConfig config.SinkConfiguration,
 	messagePublisher message.Publisher,
 	eventPublisher eventbus.Publisher,
 	logger *slog.Logger,
 	meter metric.Meter,
-) (flushhandler.FlushEventHandler, error) {
+) (flushhandler.FlushEventHandler, func(), error) {
 	flushHandlerMux := flushhandler.NewFlushEventHandlers()
 
 	// We should only close the producer once the ingest events are fully processed
@@ -48,12 +65,27 @@ func NewFlushHandler(
 		MaxEventsInBatch: sinkConfig.IngestNotifications.MaxEventsInBatch,
 	})
 	if err != nil {
-		return nil, err
+		return nil, nil, err
 	}
 
 	flushHandlerMux.AddHandler(ingestNotificationHandler)
 
-	return flushHandlerMux, nil
+	closer := func() {
+		logger.Info("shutting down flush success handlers")
+
+		if err = flushHandlerMux.Close(); err != nil {
+			logger.Error("failed to close flush success handler", slog.String("err", err.Error()))
+		}
+
+		drainCtx, cancel := context.WithTimeout(context.Background(), sinkConfig.DrainTimeout)
+		defer cancel()
+
+		if err = flushHandlerMux.WaitForDrain(drainCtx); err != nil {
+			logger.Error("failed to drain flush success handlers", slog.String("err", err.Error()))
+		}
+	}
+
+	return flushHandlerMux, closer, nil
 }
 
 func SinkWorkerProvisionTopics(conf config.EventsConfiguration) []pkgkafka.TopicConfig {
@@ -64,3 +96,95 @@ func SinkWorkerProvisionTopics(conf config.EventsConfiguration) []pkgkafka.Topic
 		},
 	}
 }
+
+func NewSinkStorage(
+	streaming streaming.Connector,
+) (sink.Storage, error) {
+	return sink.NewClickhouseStorage(sink.ClickHouseStorageConfig{
+		Streaming: streaming,
+	})
+}
+
+func NewSinkDeduplicator(conf config.SinkConfiguration, logger *slog.Logger) (dedupe.Deduplicator, func(), error) {
+	closer := func() {}
+
+	if conf.Dedupe.Enabled {
+		deduplicator, err := conf.Dedupe.NewDeduplicator()
+		if err != nil {
+			return nil, nil, fmt.Errorf("failed to initialize deduplicator: %w", err)
+		}
+
+		closer = func() {
+			err = deduplicator.Close()
+			if err != nil {
+				logger.Error("failed to close sink deduplicator", slog.String("err", err.Error()))
+			}
+		}
+
+		return deduplicator, closer, nil
+	}
+
+	return nil, closer, nil
+}
+
+type SinkKafkaConsumer = kafka.Consumer
+
+func NewSinkKafkaConsumer(conf config.SinkConfiguration, logger *slog.Logger) (*kafka.Consumer, func(), error) {
+	consumerConfig := conf.Kafka.AsConsumerConfig()
+
+	// Override the following Kafka consumer configuration parameters with hardcoded values
+	// as the Sink implementation relies on these to be set to a specific value.
+	consumerConfig.EnableAutoCommit = true
+	consumerConfig.EnableAutoOffsetStore = false
+	// Used when offset retention resets the offset. In this case we want to consume from the latest offset
+	// as everything before should be already processed.
+	consumerConfig.AutoOffsetReset = "latest"
+
+	return NewKafkaConsumer(consumerConfig, logger)
+}
+
+func NewSink(
+	conf config.SinkConfiguration,
+	logger *slog.Logger,
+	metricMeter metric.Meter,
+	tracer trace.Tracer,
+	kafkaConsumer *SinkKafkaConsumer,
+	sinkStorage sink.Storage,
+	deduplicator dedupe.Deduplicator,
+	meterService meter.Service,
+	topicResolver *topicresolver.NamespacedTopicResolver,
+	flushHandler flushhandler.FlushEventHandler,
+) (*sink.Sink, func(), error) {
+	s, err := sink.NewSink(sink.SinkConfig{
+		Logger:                  logger,
+		Tracer:                  tracer,
+		MetricMeter:             metricMeter,
+		Storage:                 sinkStorage,
+		Deduplicator:            deduplicator,
+		Consumer:                kafkaConsumer,
+		MinCommitCount:          conf.MinCommitCount,
+		MaxCommitWait:           conf.MaxCommitWait,
+		MaxPollTimeout:          conf.MaxPollTimeout,
+		NamespaceRefetch:        conf.NamespaceRefetch,
+		NamespaceRefetchTimeout: conf.NamespaceRefetchTimeout,
+		NamespaceTopicRegexp:    conf.NamespaceTopicRegexp,
+		FlushEventHandler:       flushHandler,
+		FlushSuccessTimeout:     conf.FlushSuccessTimeout,
+		DrainTimeout:            conf.DrainTimeout,
+		TopicResolver:           topicResolver,
+		MeterRefetchInterval:    conf.MeterRefetchInterval,
+		MeterService:            meterService,
+		LogDroppedEvents:        conf.LogDroppedEvents,
+	})
+	if err != nil {
+		return nil, nil, fmt.Errorf("failed to initialize sink: %w", err)
+	}
+
+	closer := func() {
+		if err = s.Close(); err != nil {
+			logger.Error("failed to close sink worker", slog.String("err", err.Error()))
+		}
+	}
+
+	return s, closer, nil
+}