<commit_before>package output

import (
	"time"

	"github.com/Shopify/sarama"
	"github.com/funkygao/dbus/engine"
	"github.com/funkygao/dbus/model"
	"github.com/funkygao/dbus/plugins/input/myslave"
	"github.com/funkygao/gafka/zk"
	conf "github.com/funkygao/jsconf"
	log "github.com/funkygao/log4go"
)

type sentPos struct {
	Log string
	Pos uint32
}

type KafkaOutput struct {
	zone, cluster, topic string

	zkzone   *zk.ZkZone
	ack      sarama.RequiredAcks
	async    bool
	compress bool

	p  sarama.SyncProducer
	ap sarama.AsyncProducer

	sendMessage func(row *model.RowsEvent)

	pos *sentPos

	// FIXME should be shared with MysqlbinlogInput
	// currently, KafkaOutput MUST setup master_host/master_port to correctly checkpoint position
	myslave *myslave.MySlave
}

func (this *KafkaOutput) Init(config *conf.Conf) {
	this.zone = config.String("zone", "")
	this.cluster = config.String("cluster", "")
	this.topic = config.String("topic", "")
	if this.cluster == "" || this.zone == "" || this.topic == "" {
		panic("invalid configuration")
	}

	// ack is ignored in async mode
	this.ack = sarama.RequiredAcks(config.Int("ack", int(sarama.WaitForLocal)))
	this.async = config.Bool("async", false)
	if this.async {
		this.sendMessage = this.asyncSendMessage
	} else {
		this.sendMessage = this.syncSendMessage
	}
	this.compress = config.Bool("compress", false)
	this.zkzone = engine.Globals().GetOrRegisterZkzone(this.zone)

	this.pos = &sentPos{}
	if config.Bool("load_position", true) {
		this.initPosition()
	}
}

func (this *KafkaOutput) Run(r engine.OutputRunner, h engine.PluginHelper) error {
	this.myslave = engine.Globals().Registered("myslave").(*myslave.MySlave)

	if err := this.prepareProducer(); err != nil {
		return err
	}

	defer func() {
		var err error
		if this.async {
			err = this.ap.Close()
		} else {
			err = this.p.Close()
		}

		if err != nil {
			log.Error("[%s.%s.%s] %s", this.zone, this.cluster, this.topic, err)
		}
	}()

	for {
		select {
		case pack, ok := <-r.InChan():
			if !ok {
				return nil
			}

			row, ok := pack.Payload.(*model.RowsEvent)
			if !ok {
				log.Error("[%s.%s.%s] bad payload: %+v", this.zone, this.cluster, this.topic, pack.Payload)
				continue
			}

			// best effort to reduce dup message to kafka
			if row.Log > this.pos.Log ||
				(row.Log == this.pos.Log && row.Position > this.pos.Pos) {
				this.sendMessage(row)
			} else {
				log.Trace("[%s.%s.%s] skipped {%s}", this.zone, this.cluster, this.topic, row)
			}

			pack.Recycle()
		}
	}

	return nil
}

func (this *KafkaOutput) prepareProducer() error {
	cf := sarama.NewConfig()
	cf.ChannelBufferSize = 256 * 2 // default was 256
	if this.compress {
		cf.Producer.Compression = sarama.CompressionSnappy
	}

	zkcluster := this.zkzone.NewCluster(this.cluster)

	if !this.async {
		cf.Producer.RequiredAcks = this.ack
		p, err := sarama.NewSyncProducer(zkcluster.BrokerList(), cf)
		if err != nil {
			return err
		}

		this.p = p
		return nil
	}

	// async producer<commit_msg>ajust the log level<commit_after>package output

import (
	"time"

	"github.com/Shopify/sarama"
	"github.com/funkygao/dbus/engine"
	"github.com/funkygao/dbus/model"
	"github.com/funkygao/dbus/plugins/input/myslave"
	"github.com/funkygao/gafka/zk"
	conf "github.com/funkygao/jsconf"
	log "github.com/funkygao/log4go"
)

type sentPos struct {
	Log string
	Pos uint32
}

type KafkaOutput struct {
	zone, cluster, topic string

	zkzone   *zk.ZkZone
	ack      sarama.RequiredAcks
	async    bool
	compress bool

	p  sarama.SyncProducer
	ap sarama.AsyncProducer

	sendMessage func(row *model.RowsEvent)

	pos *sentPos

	// FIXME should be shared with MysqlbinlogInput
	// currently, KafkaOutput MUST setup master_host/master_port to correctly checkpoint position
	myslave *myslave.MySlave
}

func (this *KafkaOutput) Init(config *conf.Conf) {
	this.zone = config.String("zone", "")
	this.cluster = config.String("cluster", "")
	this.topic = config.String("topic", "")
	if this.cluster == "" || this.zone == "" || this.topic == "" {
		panic("invalid configuration")
	}

	// ack is ignored in async mode
	this.ack = sarama.RequiredAcks(config.Int("ack", int(sarama.WaitForLocal)))
	this.async = config.Bool("async", false)
	if this.async {
		this.sendMessage = this.asyncSendMessage
	} else {
		this.sendMessage = this.syncSendMessage
	}
	this.compress = config.Bool("compress", false)
	this.zkzone = engine.Globals().GetOrRegisterZkzone(this.zone)

	this.pos = &sentPos{}
	if config.Bool("load_position", true) {
		this.initPosition()
	}
}

func (this *KafkaOutput) Run(r engine.OutputRunner, h engine.PluginHelper) error {
	this.myslave = engine.Globals().Registered("myslave").(*myslave.MySlave)

	if err := this.prepareProducer(); err != nil {
		return err
	}

	defer func() {
		var err error
		if this.async {
			err = this.ap.Close()
		} else {
			err = this.p.Close()
		}

		if err != nil {
			log.Error("[%s.%s.%s] %s", this.zone, this.cluster, this.topic, err)
		}
	}()

	for {
		select {
		case pack, ok := <-r.InChan():
			if !ok {
				return nil
			}

			row, ok := pack.Payload.(*model.RowsEvent)
			if !ok {
				log.Error("[%s.%s.%s] bad payload: %+v", this.zone, this.cluster, this.topic, pack.Payload)
				continue
			}

			// best effort to reduce dup message to kafka
			if row.Log > this.pos.Log ||
				(row.Log == this.pos.Log && row.Position > this.pos.Pos) {
				this.sendMessage(row)
			} else {
				log.Debug("[%s.%s.%s] skipped {%s}", this.zone, this.cluster, this.topic, row)
			}

			pack.Recycle()
		}
	}

	return nil
}

func (this *KafkaOutput) prepareProducer() error {
	cf := sarama.NewConfig()
	cf.ChannelBufferSize = 256 * 2 // default was 256
	if this.compress {
		cf.Producer.Compression = sarama.CompressionSnappy
	}

	zkcluster := this.zkzone.NewCluster(this.cluster)

	if !this.async {
		cf.Producer.RequiredAcks = this.ack
		p, err := sarama.NewSyncProducer(zkcluster.BrokerList(), cf)
		if err != nil {
			return err
		}

		this.p = p
		return nil
	}

	// async producer
{'fraction_non_alphanumeric': 0.14809873248832556, 'fraction_numerical': 0.00333555703802535, 'mean_word_length': 3.468703427719821, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 7, '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': '217965', 'n_tokens_mistral': 2421, 'n_tokens_neox': 2194, 'n_words': 472}