<commit_before>/*

kasper is a lightweight Kafka stream processing library.

 */

package kasper

import (
	"log"
	"github.com/Shopify/sarama"
	"time"
	"sync"
)

type TopicProcessor struct {
	config              *TopicProcessorConfig
	containerId         ContainerId
	client              sarama.Client
	producer            sarama.AsyncProducer
	offsetManager       sarama.OffsetManager
	partitionProcessors []*partitionProcessor
	inputTopics         []Topic
	partitions          []Partition
	shutdown            chan bool
	waitGroup           sync.WaitGroup
}

type MessageProcessor interface {
	Process(IncomingMessage, Sender, Coordinator)
}

// NewTopicProcessor creates a new TopicProcessor with the given config.
// It requires a factory function that creates MessageProcessor instances and a container id.
// The container id must be a number between 0 and config.ContainerCount - 1.
func NewTopicProcessor(config *TopicProcessorConfig, makeProcessor func() MessageProcessor, cid ContainerId) *TopicProcessor {
	// TODO: check all input topics are covered by a Serde
	// TODO: check all input partitions and make sure PartitionAssignment is valid
	// TODO: check cid is within [0, ContainerCount)
	inputTopics := config.InputTopics
	brokerList := config.BrokerList
	client, err := sarama.NewClient(brokerList, sarama.NewConfig())
	if err != nil {
		log.Fatal(err)
	}
	partitions := config.partitionsForContainer(cid)
	offsetManager, err := sarama.NewOffsetManagerFromClient(config.kafkaConsumerGroup(), client)
	if err != nil {
		log.Fatal(err)
	}
	partitionProcessors := make([]*partitionProcessor, len(partitions))
	requiredAcks := config.KasperConfig.RequiredAcks
	producer := mustSetupProducer(config.BrokerList, config.producerClientId(cid), requiredAcks)
	topicProcessor := TopicProcessor{
		config,
		cid,
		client,
		producer,
		offsetManager,
		partitionProcessors,
		inputTopics,
		partitions,
		make(chan bool),
		sync.WaitGroup{},
	}
	for i, partition := range partitions {
		processor := makeProcessor()
		partitionProcessors[i] = newPartitionProcessor(&topicProcessor, processor, partition)
	}
	return &topicProcessor
}

func (tp *TopicProcessor) Start() {
	tp.waitGroup.Add(1)
	go func() {
		defer tp.waitGroup.Done()
		tp.runLoop()
	}()
}

func (tp *TopicProcessor) runLoop() {
	consumerChan := tp.getConsumerMessagesChan()
	var tickerChan <-chan time.Time
	var ticker *time.Ticker

	if tp.config.markOffsetsAutomatically() {
		ticker = time.NewTicker(tp.config.AutoMarkOffsetsInterval)
		tickerChan = ticker.C
	} else {
		tickerChan = make(<-chan time.Time)
	}

	tp.waitGroup.Add(1)
	go func() {
		defer tp.waitGroup.Done()
		for err := range tp.producer.Errors() {
			tp.onProducerError(err)
		}
	}()

	for {
		select {
		case consumerMessage := <-consumerChan:
			pp := tp.partitionProcessors[consumerMessage.Partition]
			for {
				if pp.isReadyForMessage(consumerMessage) {
					producerMessages := pp.process(consumerMessage)
					for len(producerMessages) > 0 {
						select {
						case tp.producer.Input() <- producerMessages[0]:
							producerMessages = producerMessages[1:]
						case msg, more := <-tp.producer.Successes():
							tp.onProducerAck(msg, more)
						case <-tickerChan:
							tp.onTick()
						}
					}
					pp.onProcessCompleted()
					break
				} else {
					select {
					case msg, more := <-tp.producer.Successes():
						tp.onProducerAck(msg, more)
					case <-tickerChan:
						tp.onTick()
					}
				}
			}
		case msg, more := <-tp.producer.Successes():
			tp.onProducerAck(msg, more)
		case <-tickerChan:
			tp.onTick()
		case <-tp.shutdown:
			tp.onShutdown()
			return
		}
	}
}

func (tp *TopicProcessor) onShutdown() {
	for _, pp := range tp.partitionProcessors {
		pp.onShutdown()
	}
	tp.producer.Close()
	tp.client.Close()
}

func (tp *TopicProcessor) getConsumerMessagesChan() <-chan *sarama.ConsumerMessage {
	consumerMessagesChan := make(chan *sarama.ConsumerMessage)
	for _, ch := range tp.consumerMessageChannels() {
		tp.waitGroup.Add(1)
		go func(c <-chan *sarama.ConsumerMessage) {
			defer tp.waitGroup.Done()
			for msg := range c {
				consumerMessagesChan <- msg
			}
		}(ch)
	}
	return consumerMessagesChan
}

func (tp *TopicProcessor) onProducerError(error *sarama.ProducerError) {
	log.Fatal(error) /* FIXME Handle this gracefully with a retry count / backoff period */
}

func (tp *TopicProcessor) onTick() {
	for _, pp := range tp.partitionProcessors {
		pp.markOffsetsIfPossible()
	}
}

func (tp *TopicProcessor) consumerMessageChannels() []<-chan *sarama.ConsumerMessage {
	var chans []<-chan *sarama.ConsumerMessage
	for _, partitionProcessor := range tp.partitionProcessors {
		partitionChannels := partitionProcessor.consumerMessageChannels()
		for _, ch := range partitionChannels {
			chans = append(chans, ch)
		}
	}
	return chans
}

func mustSetupProducer(brokers []string, producerClientId string, requiredAcks sarama.RequiredAcks) sarama.AsyncProducer {
	saramaConfig := sarama.NewConfig()
	saramaConfig.ClientID = producerClientId
	saramaConfig.Producer.Return.Successes = true
	saramaConfig.Producer.Partitioner = sarama.NewManualPartitioner
	saramaConfig.Producer.RequiredAcks = requiredAcks

	producer, err := sarama.NewAsyncProducer(brokers, saramaConfig)
	if err != nil {
		log.Fatal(err)
	}

	return producer
}

func (tp *TopicProcessor) Shutdown() {
	tp.shutdown <- true
	tp.waitGroup.Wait()
}

func (tp *TopicProcessor) onProducerAck(producerMessage *sarama.ProducerMessage, more bool) {
	if !more {
		return
	}
	pp := tp.partitionProcessors[producerMessage.Partition]
	pp.onProducerAck(producerMessage)
}<commit_msg>Configure Sarama to use the default metrics registry<commit_after>/*

kasper is a lightweight Kafka stream processing library.

 */

package kasper

import (
	"log"
	"github.com/Shopify/sarama"
	"time"
	"sync"
	"github.com/rcrowley/go-metrics"
)

type TopicProcessor struct {
	config              *TopicProcessorConfig
	containerId         ContainerId
	client              sarama.Client
	producer            sarama.AsyncProducer
	offsetManager       sarama.OffsetManager
	partitionProcessors []*partitionProcessor
	inputTopics         []Topic
	partitions          []Partition
	shutdown            chan bool
	waitGroup           sync.WaitGroup
}

type MessageProcessor interface {
	Process(IncomingMessage, Sender, Coordinator)
}

// NewTopicProcessor creates a new TopicProcessor with the given config.
// It requires a factory function that creates MessageProcessor instances and a container id.
// The container id must be a number between 0 and config.ContainerCount - 1.
func NewTopicProcessor(config *TopicProcessorConfig, makeProcessor func() MessageProcessor, cid ContainerId) *TopicProcessor {
	// TODO: check all input topics are covered by a Serde
	// TODO: check all input partitions and make sure PartitionAssignment is valid
	// TODO: check cid is within [0, ContainerCount)
	inputTopics := config.InputTopics
	brokerList := config.BrokerList
	client, err := sarama.NewClient(brokerList, sarama.NewConfig())
	if err != nil {
		log.Fatal(err)
	}
	partitions := config.partitionsForContainer(cid)
	offsetManager, err := sarama.NewOffsetManagerFromClient(config.kafkaConsumerGroup(), client)
	if err != nil {
		log.Fatal(err)
	}
	partitionProcessors := make([]*partitionProcessor, len(partitions))
	requiredAcks := config.KasperConfig.RequiredAcks
	producer := mustSetupProducer(config.BrokerList, config.producerClientId(cid), requiredAcks)
	topicProcessor := TopicProcessor{
		config,
		cid,
		client,
		producer,
		offsetManager,
		partitionProcessors,
		inputTopics,
		partitions,
		make(chan bool),
		sync.WaitGroup{},
	}
	for i, partition := range partitions {
		processor := makeProcessor()
		partitionProcessors[i] = newPartitionProcessor(&topicProcessor, processor, partition)
	}
	return &topicProcessor
}

func (tp *TopicProcessor) Start() {
	tp.waitGroup.Add(1)
	go func() {
		defer tp.waitGroup.Done()
		tp.runLoop()
	}()
}

func (tp *TopicProcessor) runLoop() {
	consumerChan := tp.getConsumerMessagesChan()
	var tickerChan <-chan time.Time
	var ticker *time.Ticker

	if tp.config.markOffsetsAutomatically() {
		ticker = time.NewTicker(tp.config.AutoMarkOffsetsInterval)
		tickerChan = ticker.C
	} else {
		tickerChan = make(<-chan time.Time)
	}

	tp.waitGroup.Add(1)
	go func() {
		defer tp.waitGroup.Done()
		for err := range tp.producer.Errors() {
			tp.onProducerError(err)
		}
	}()

	for {
		select {
		case consumerMessage := <-consumerChan:
			pp := tp.partitionProcessors[consumerMessage.Partition]
			for {
				if pp.isReadyForMessage(consumerMessage) {
					producerMessages := pp.process(consumerMessage)
					for len(producerMessages) > 0 {
						select {
						case tp.producer.Input() <- producerMessages[0]:
							producerMessages = producerMessages[1:]
						case msg, more := <-tp.producer.Successes():
							tp.onProducerAck(msg, more)
						case <-tickerChan:
							tp.onTick()
						}
					}
					pp.onProcessCompleted()
					break
				} else {
					select {
					case msg, more := <-tp.producer.Successes():
						tp.onProducerAck(msg, more)
					case <-tickerChan:
						tp.onTick()
					}
				}
			}
		case msg, more := <-tp.producer.Successes():
			tp.onProducerAck(msg, more)
		case <-tickerChan:
			tp.onTick()
		case <-tp.shutdown:
			tp.onShutdown()
			return
		}
	}
}

func (tp *TopicProcessor) onShutdown() {
	for _, pp := range tp.partitionProcessors {
		pp.onShutdown()
	}
	tp.producer.Close()
	tp.client.Close()
}

func (tp *TopicProcessor) getConsumerMessagesChan() <-chan *sarama.ConsumerMessage {
	consumerMessagesChan := make(chan *sarama.ConsumerMessage)
	for _, ch := range tp.consumerMessageChannels() {
		tp.waitGroup.Add(1)
		go func(c <-chan *sarama.ConsumerMessage) {
			defer tp.waitGroup.Done()
			for msg := range c {
				consumerMessagesChan <- msg
			}
		}(ch)
	}
	return consumerMessagesChan
}

func (tp *TopicProcessor) onProducerError(error *sarama.ProducerError) {
	log.Fatal(error) /* FIXME Handle this gracefully with a retry count / backoff period */
}

func (tp *TopicProcessor) onTick() {
	for _, pp := range tp.partitionProcessors {
		pp.markOffsetsIfPossible()
	}
}

func (tp *TopicProcessor) consumerMessageChannels() []<-chan *sarama.ConsumerMessage {
	var chans []<-chan *sarama.ConsumerMessage
	for _, partitionProcessor := range tp.partitionProcessors {
		partitionChannels := partitionProcessor.consumerMessageChannels()
		for _, ch := range partitionChannels {
			chans = append(chans, ch)
		}
	}
	return chans
}

func mustSetupProducer(brokers []string, producerClientId string, requiredAcks sarama.RequiredAcks) sarama.AsyncProducer {
	saramaConfig := sarama.NewConfig()
	saramaConfig.ClientID = producerClientId
	saramaConfig.Producer.Return.Successes = true
	saramaConfig.Producer.Partitioner = sarama.NewManualPartitioner
	saramaConfig.Producer.RequiredAcks = requiredAcks
	saramaConfig.MetricRegistry = metrics.DefaultRegistry

	producer, err := sarama.NewAsyncProducer(brokers, saramaConfig)
	if err != nil {
		log.Fatal(err)
	}

	return producer
}

func (tp *TopicProcessor) Shutdown() {
	tp.shutdown <- true
	tp.waitGroup.Wait()
}

func (tp *TopicProcessor) onProducerAck(producerMessage *sarama.ProducerMessage, more bool) {
	if !more {
		return
	}
	pp := tp.partitionProcessors[producerMessage.Partition]
	pp.onProducerAck(producerMessage)
}
{'fraction_non_alphanumeric': 0.10851962682626298, 'fraction_numerical': 0.0015842281288505544, 'mean_word_length': 4.005726872246696, 'pattern_counts': {'":': 0, '<': 37, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3698008', 'n_tokens_mistral': 4068, 'n_tokens_neox': 3608, 'n_words': 832}