writer.go

package kafka

import (
	"context"
	"fmt"
	"io"
	"math/rand"
	"sort"
	"sync"
	"time"
)

// The Writer type provides the implementation of a producer of kafka messages
// that automatically distributes messages across partitions of a single topic
// using a configurable balancing policy.
//
// Instances of Writer are safe to use concurrently from multiple goroutines.
type Writer struct {
	config WriterConfig

	mutex  sync.RWMutex
	closed bool

	join sync.WaitGroup
	msgs chan writerMessage
	done chan struct{}

	// writer stats are all made of atomic values, no need for synchronization.
	// Use a pointer to ensure 64-bit alignment of the values.
	stats *writerStats
}

// WriterConfig is a configuration type used to create new instances of Writer.
type WriterConfig struct {
	// The list of brokers used to discover the partitions available on the
	// kafka cluster.
	//
	// This field is required, attempting to create a writer with an empty list
	// of brokers will panic.
	Brokers []string

	// The topic that the writer will produce messages to.
	//
	// This field is required, attempting to create a writer with an empty topic
	// will panic.
	Topic string

	// The dialer used by the writer to establish connections to the kafka
	// cluster.
	//
	// If nil, the default dialer is used instead.
	Dialer *Dialer

	// The balancer used to distribute messages across partitions.
	//
	// The default is to use a round-robin distribution.
	Balancer Balancer

	// Limit on how many attempts will be made to deliver a message.
	//
	// The default is to try at most 10 times.
	MaxAttempts int

	// A hint on the capacity of the writer's internal message queue.
	//
	// The default is to use a queue capacity of 100 messages.
	QueueCapacity int

	// Limit on how many messages will be buffered before being sent to a
	// partition.
	//
	// The default is to use a target batch size of 100 messages.
	BatchSize int

	// Time limit on how often incomplete message batches will be flushed to
	// kafka.
	//
	// The default is to flush at least every second.
	BatchTimeout time.Duration

	// Timeout for read operations performed by the Writer.
	//
	// Defaults to 10 seconds.
	ReadTimeout time.Duration

	// Timeout for write operation performed by the Writer.
	//
	// Defaults to 10 seconds.
	WriteTimeout time.Duration

	// This interval defines how often the list of partitions is refreshed from
	// kafka. It allows the writer to automatically handle when new partitions
	// are added to a topic.
	//
	// The default is to refresh partitions every 15 seconds.
	RebalanceInterval time.Duration

	// Number of acknowledges from partition replicas required before receiving
	// a response to a produce request (default to -1, which means to wait for
	// all replicas).
	RequiredAcks int

	// Setting this flag to true causes the WriteMessages method to never block.
	// It also means that errors are ignored since the caller will not receive
	// the returned value. Use this only if you don't care about guarantees of
	// whether the messages were written to kafka.
	Async bool

	// CompressionCodec set the codec to be used to compress Kafka messages.
	CompressionCodec int8

	// CompressionLevel set the compression level if the codec supports it.
	CompressionLevel int

	newPartitionWriter func(partition int, config WriterConfig, stats *writerStats) partitionWriter
}

// WriterStats is a data structure returned by a call to Writer.Stats that
// exposes details about the behavior of the writer.
type WriterStats struct {
	Dials      int64 `metric:"kafka.writer.dial.count"      type:"counter"`
	Writes     int64 `metric:"kafka.writer.write.count"     type:"counter"`
	Messages   int64 `metric:"kafka.writer.message.count"   type:"counter"`
	Bytes      int64 `metric:"kafka.writer.message.bytes"   type:"counter"`
	Rebalances int64 `metric:"kafka.writer.rebalance.count" type:"counter"`
	Errors     int64 `metric:"kafka.writer.error.count"     type:"counter"`

	DialTime  DurationStats `metric:"kafka.writer.dial.seconds"`
	WriteTime DurationStats `metric:"kafka.writer.write.seconds"`
	WaitTime  DurationStats `metric:"kafka.writer.wait.seconds"`
	Retries   SummaryStats  `metric:"kafka.writer.retries.count"`
	BatchSize SummaryStats  `metric:"kafka.writer.batch.size"`

	MaxAttempts       int64         `metric:"kafka.writer.attempts.max"       type:"gauge"`
	MaxBatchSize      int64         `metric:"kafka.writer.batch.max"          type:"gauge"`
	BatchTimeout      time.Duration `metric:"kafka.writer.batch.timeout"      type:"gauge"`
	ReadTimeout       time.Duration `metric:"kafka.writer.read.timeout"       type:"gauge"`
	WriteTimeout      time.Duration `metric:"kafka.writer.write.timeout"      type:"gauge"`
	RebalanceInterval time.Duration `metric:"kafka.writer.rebalance.interval" type:"gauge"`
	RequiredAcks      int64         `metric:"kafka.writer.acks.required"      type:"gauge"`
	Async             bool          `metric:"kafka.writer.async"              type:"gauge"`
	QueueLength       int64         `metric:"kafka.writer.queue.length"       type:"gauge"`
	QueueCapacity     int64         `metric:"kafka.writer.queue.capacity"     type:"gauge"`

	ClientID string `tag:"client_id"`
	Topic    string `tag:"topic"`
}

// writerStats is a struct that contains statistics on a writer.
//
// Since atomic is used to mutate the statistics the values must be 64-bit aligned.
// This is easily accomplished by always allocating this struct directly, (i.e. using a pointer to the struct).
// See https://golang.org/pkg/sync/atomic/#pkg-note-BUG
type writerStats struct {
	dials      counter
	writes     counter
	messages   counter
	bytes      counter
	rebalances counter
	errors     counter
	dialTime   summary
	writeTime  summary
	waitTime   summary
	retries    summary
	batchSize  summary
}

// NewWriter creates and returns a new Writer configured with config.
func NewWriter(config WriterConfig) *Writer {
	if len(config.Brokers) == 0 {
		panic("cannot create a kafka writer with an empty list of brokers")
	}

	if len(config.Topic) == 0 {
		panic("cannot create a kafka writer with an empty topic")
	}

	if config.Dialer == nil {
		config.Dialer = DefaultDialer
	}

	if config.Balancer == nil {
		config.Balancer = &RoundRobin{}
	}

	if config.newPartitionWriter == nil {
		config.newPartitionWriter = func(partition int, config WriterConfig, stats *writerStats) partitionWriter {
			return newWriter(partition, config, stats)
		}
	}

	if config.MaxAttempts == 0 {
		config.MaxAttempts = 10
	}

	if config.QueueCapacity == 0 {
		config.QueueCapacity = 100
	}

	if config.BatchSize == 0 {
		config.BatchSize = 100
	}

	if config.BatchTimeout == 0 {
		config.BatchTimeout = 1 * time.Second
	}

	if config.ReadTimeout == 0 {
		config.ReadTimeout = 10 * time.Second
	}

	if config.WriteTimeout == 0 {
		config.WriteTimeout = 10 * time.Second
	}

	if config.RebalanceInterval == 0 {
		config.RebalanceInterval = 15 * time.Second
	}

	if config.CompressionLevel == 0 {
		config.CompressionLevel = DefaultCompressionLevel
	}

	w := &Writer{
		config: config,
		msgs:   make(chan writerMessage, config.QueueCapacity),
		done:   make(chan struct{}),
		stats: &writerStats{
			dialTime:  makeSummary(),
			writeTime: makeSummary(),
			waitTime:  makeSummary(),
			retries:   makeSummary(),
		},
	}

	w.join.Add(1)
	go w.run()
	return w
}

// WriteMessages writes a batch of messages to the kafka topic configured on this
// writer.
//
// Unless the writer was configured to write messages asynchronously, the method
// blocks until all messages have been written, or until the maximum number of
// attempts was reached.
//
// When the method returns an error, there's no way to know yet which messages
// have succeeded of failed.
//
// The context passed as first argument may also be used to asynchronously
// cancel the operation. Note that in this case there are no guarantees made on
// whether messages were written to kafka. The program should assume that the
// whole batch failed and re-write the messages later (which could then cause
// duplicates).
func (w *Writer) WriteMessages(ctx context.Context, msgs ...Message) error {
	if len(msgs) == 0 {
		return nil
	}

	var res = make(chan error, len(msgs))
	var err error

	t0 := time.Now()

	for attempt := 0; attempt < w.config.MaxAttempts; attempt++ {
		w.mutex.RLock()

		if w.closed {
			w.mutex.RUnlock()
			return io.ErrClosedPipe
		}

		for _, msg := range msgs {
			msg.CompressionCodec = w.config.CompressionCodec
			msg.CompressionLevel = w.config.CompressionLevel

			select {
			case w.msgs <- writerMessage{
				msg: msg,
				res: res,
			}:
			case <-ctx.Done():
				w.mutex.RUnlock()
				return ctx.Err()
			}
		}

		w.mutex.RUnlock()

		if w.config.Async {
			break
		}

		var retry []Message

		for i := 0; i != len(msgs); i++ {
			select {
			case e := <-res:
				if e != nil {
					if we, ok := e.(*writerError); ok {
						w.stats.retries.observe(1)
						retry, err = append(retry, we.msg), we.err
					} else {
						err = e
					}
				}
			case <-ctx.Done():
				return ctx.Err()
			}
		}

		if msgs = retry; len(msgs) == 0 {
			break
		}

		timer := time.NewTimer(backoff(attempt+1, 100*time.Millisecond, 1*time.Second))
		select {
		case <-timer.C:
			// Only clear the error (so we retry the loop) if we have more retries, otherwise
			// we risk silencing the error.
			if attempt < w.config.MaxAttempts-1 {
				err = nil
			}
		case <-ctx.Done():
			err = ctx.Err()
		case <-w.done:
			err = io.ErrClosedPipe
		}
		timer.Stop()

		if err != nil {
			break
		}
	}

	t1 := time.Now()
	w.stats.writeTime.observeDuration(t1.Sub(t0))

	return err
}

// Stats returns a snapshot of the writer stats since the last time the method
// was called, or since the writer was created if it is called for the first
// time.
//
// A typical use of this method is to spawn a goroutine that will periodically
// call Stats on a kafka writer and report the metrics to a stats collection
// system.
func (w *Writer) Stats() WriterStats {
	return WriterStats{
		Dials:             w.stats.dials.snapshot(),
		Writes:            w.stats.writes.snapshot(),
		Messages:          w.stats.messages.snapshot(),
		Bytes:             w.stats.bytes.snapshot(),
		Rebalances:        w.stats.rebalances.snapshot(),
		Errors:            w.stats.errors.snapshot(),
		DialTime:          w.stats.dialTime.snapshotDuration(),
		WriteTime:         w.stats.writeTime.snapshotDuration(),
		WaitTime:          w.stats.waitTime.snapshotDuration(),
		Retries:           w.stats.retries.snapshot(),
		BatchSize:         w.stats.batchSize.snapshot(),
		MaxAttempts:       int64(w.config.MaxAttempts),
		MaxBatchSize:      int64(w.config.BatchSize),
		BatchTimeout:      w.config.BatchTimeout,
		ReadTimeout:       w.config.ReadTimeout,
		WriteTimeout:      w.config.WriteTimeout,
		RebalanceInterval: w.config.RebalanceInterval,
		RequiredAcks:      int64(w.config.RequiredAcks),
		Async:             w.config.Async,
		QueueLength:       int64(len(w.msgs)),
		QueueCapacity:     int64(cap(w.msgs)),
		ClientID:          w.config.Dialer.ClientID,
		Topic:             w.config.Topic,
	}
}

// Close flushes all buffered messages and closes the writer. The call to Close
// aborts any concurrent calls to WriteMessages, which then return with the
// io.ErrClosedPipe error.
func (w *Writer) Close() (err error) {
	w.mutex.Lock()

	if !w.closed {
		w.closed = true
		close(w.msgs)
		close(w.done)
	}

	w.mutex.Unlock()
	w.join.Wait()
	return
}

func (w *Writer) run() {
	defer w.join.Done()

	ticker := time.NewTicker(w.config.RebalanceInterval)
	defer ticker.Stop()

	var rebalance = true
	var writers = make(map[int]partitionWriter)
	var partitions []int
	var err error

	for {
		if rebalance {
			w.stats.rebalances.observe(1)
			rebalance = false

			var newPartitions []int
			var oldPartitions = partitions

			if newPartitions, err = w.partitions(); err == nil {
				for _, partition := range diffp(oldPartitions, newPartitions) {
					w.close(writers[partition])
					delete(writers, partition)
				}

				for _, partition := range diffp(newPartitions, oldPartitions) {
					writers[partition] = w.open(partition)
				}

				partitions = newPartitions
			}
		}

		select {
		case wm, ok := <-w.msgs:
			if !ok {
				for _, writer := range writers {
					w.close(writer)
				}
				return
			}

			if len(partitions) != 0 {
				selectedPartition := w.config.Balancer.Balance(wm.msg, partitions...)
				writers[selectedPartition].messages() <- wm
			} else {
				// No partitions were found because the topic doesn't exist.
				if err == nil {
					err = fmt.Errorf("failed to find any partitions for topic %s", w.config.Topic)
				}

				wm.res <- &writerError{msg: wm.msg, err: err}
			}

		case <-ticker.C:
			rebalance = true
		}
	}
}

func (w *Writer) partitions() (partitions []int, err error) {
	for _, broker := range shuffledStrings(w.config.Brokers) {
		var conn *Conn
		var plist []Partition

		if conn, err = w.config.Dialer.Dial("tcp", broker); err != nil {
			continue
		}

		conn.SetReadDeadline(time.Now().Add(w.config.ReadTimeout))
		plist, err = conn.ReadPartitions(w.config.Topic)
		conn.Close()

		if err == nil {
			partitions = make([]int, len(plist))
			for i, p := range plist {
				partitions[i] = p.ID
			}
			break
		}
	}

	sort.Ints(partitions)
	return
}

func (w *Writer) open(partition int) partitionWriter {
	return w.config.newPartitionWriter(partition, w.config, w.stats)
}

func (w *Writer) close(writer partitionWriter) {
	w.join.Add(1)
	go func() {
		writer.close()
		w.join.Done()
	}()
}

func diffp(new []int, old []int) (diff []int) {
	for _, p := range new {
		if i := sort.SearchInts(old, p); i == len(old) || old[i] != p {
			diff = append(diff, p)
		}
	}
	return
}

type partitionWriter interface {
	messages() chan<- writerMessage
	close()
}

type writer struct {
	brokers      []string
	topic        string
	partition    int
	requiredAcks int
	batchSize    int
	batchTimeout time.Duration
	writeTimeout time.Duration
	dialer       *Dialer
	msgs         chan writerMessage
	join         sync.WaitGroup
	stats        *writerStats
}

func newWriter(partition int, config WriterConfig, stats *writerStats) *writer {
	w := &writer{
		brokers:      config.Brokers,
		topic:        config.Topic,
		partition:    partition,
		requiredAcks: config.RequiredAcks,
		batchSize:    config.BatchSize,
		batchTimeout: config.BatchTimeout,
		writeTimeout: config.WriteTimeout,
		dialer:       config.Dialer,
		msgs:         make(chan writerMessage, config.QueueCapacity),
		stats:        stats,
	}
	w.join.Add(1)
	go w.run()
	return w
}

func (w *writer) close() {
	close(w.msgs)
	w.join.Wait()
}

func (w *writer) messages() chan<- writerMessage {
	return w.msgs
}

func (w *writer) run() {
	defer w.join.Done()

	ticker := time.NewTicker(w.batchTimeout / 10)
	defer ticker.Stop()

	var conn *Conn
	var done bool
	var batch = make([]Message, 0, w.batchSize)
	var resch = make([](chan<- error), 0, w.batchSize)
	var lastFlushAt = time.Now()

	defer func() {
		if conn != nil {
			conn.Close()
		}
	}()

	for !done {
		var mustFlush bool

		select {
		case wm, ok := <-w.msgs:
			if !ok {
				done, mustFlush = true, true
			} else {
				batch = append(batch, wm.msg)
				resch = append(resch, wm.res)
				mustFlush = len(batch) >= w.batchSize
			}

		case now := <-ticker.C:
			mustFlush = now.Sub(lastFlushAt) > w.batchTimeout
		}

		if mustFlush {
			lastFlushAt = time.Now()

			if len(batch) == 0 {
				continue
			}

			var err error
			if conn, err = w.write(conn, batch, resch); err != nil {
				if conn != nil {
					conn.Close()
					conn = nil
				}
			}

			for i := range batch {
				batch[i] = Message{}
			}

			for i := range resch {
				resch[i] = nil
			}

			batch = batch[:0]
			resch = resch[:0]
		}
	}
}

func (w *writer) dial() (conn *Conn, err error) {
	for _, broker := range shuffledStrings(w.brokers) {
		t0 := time.Now()
		if conn, err = w.dialer.DialLeader(context.Background(), "tcp", broker, w.topic, w.partition); err == nil {
			t1 := time.Now()
			w.stats.dials.observe(1)
			w.stats.dialTime.observeDuration(t1.Sub(t0))
			conn.SetRequiredAcks(w.requiredAcks)
			break
		}
	}
	return
}

func (w *writer) write(conn *Conn, batch []Message, resch [](chan<- error)) (ret *Conn, err error) {
	w.stats.writes.observe(1)
	if conn == nil {
		if conn, err = w.dial(); err != nil {
			w.stats.errors.observe(1)
			for i, res := range resch {
				res <- &writerError{msg: batch[i], err: err}
			}
			return
		}
	}

	t0 := time.Now()
	conn.SetWriteDeadline(time.Now().Add(w.writeTimeout))

	if _, err = conn.WriteMessages(batch...); err != nil {
		w.stats.errors.observe(1)
		for i, res := range resch {
			res <- &writerError{msg: batch[i], err: err}
		}
	} else {
		for _, m := range batch {
			w.stats.messages.observe(1)
			w.stats.bytes.observe(int64(len(m.Key) + len(m.Value)))
		}
		for _, res := range resch {
			res <- nil
		}
	}

	t1 := time.Now()
	w.stats.waitTime.observeDuration(t1.Sub(t0))
	w.stats.batchSize.observe(int64(len(batch)))

	ret = conn
	return
}

type writerMessage struct {
	msg Message
	res chan<- error
}

type writerError struct {
	msg Message
	err error
}

func (e *writerError) Cause() error {
	return e.err
}

func (e *writerError) Error() string {
	return e.err.Error()
}

func (e *writerError) Temporary() bool {
	return isTemporary(e.err)
}

func (e *writerError) Timeout() bool {
	return isTimeout(e.err)
}

func shuffledStrings(list []string) []string {
	shuffledList := make([]string, len(list))
	copy(shuffledList, list)

	shufflerMutex.Lock()

	for i := range shuffledList {
		j := shuffler.Intn(i + 1)
		shuffledList[i], shuffledList[j] = shuffledList[j], shuffledList[i]
	}

	shufflerMutex.Unlock()
	return shuffledList
}

var (
	shufflerMutex = sync.Mutex{}
	shuffler      = rand.New(rand.NewSource(time.Now().Unix()))
)