segmentio/kafka-go: 如何实现Kafka读写的极致性能与精准一次处理

dialer := &kafka.Dialer{
    Timeout:   10 * time.Second,
    DualStack: true,
    TLS:       &tls.Config{...tls config...},
}

w := kafka.Writer{
    Addr: kafka.TCP("localhost:9092", "localhost:9093", "localhost:9094"),
    Topic:   "topic-A",
    Balancer: &kafka.Hash{},
    Dialer:         dialer,
}

dialer := &kafka.Dialer{
    Timeout:   10 * time.Second,
    DualStack: true,
    TLS:       &tls.Config{...tls config...},
}

r := kafka.NewReader(kafka.ReaderConfig{
    Brokers:        []string{"localhost:9092", "localhost:9093", "localhost:9094"},
    GroupID:        "consumer-group-id",
    Topic:          "topic-A",
    Dialer:         dialer,
})

mechanism := plain.Mechanism{
    Username: "username",
    Password: "password",
}

mechanism, err := scram.Mechanism(scram.SHA512, "username", "password")
if err != nil {
    panic(err)
}

dialer := &kafka.Dialer{
    Timeout:       10 * time.Second,
    DualStack:     true,
    SASLMechanism: mechanism,
}

conn, err := dialer.DialContext(ctx, "tcp", "localhost:9093")

func consumeWithFetchMessage() {
    ctx := context.Background()
    for {
        // 从 Kafka 中获取下一条消息
        m, err := kafkaReader.FetchMessage(ctx)
        if err != nil {
            log.Printf("获取消息时出错: %v", err)
            break
        }
        // 打印消息内容
        log.Printf("消息: %s, 偏移量: %d", string(m.Value), m.Offset)
        // 处理消息 (在这里可以进行你的业务逻辑)
        // 手动提交偏移量
        if err := kafkaReader.CommitMessages(ctx, m); err != nil {
            log.Printf("提交偏移量时出错: %v", err)
        }
    }
}

writerConfig := kafka.WriterConfig{
		Brokers:      kc.Brokers,
		BatchSize:    kc.BatchSize,
		BatchBytes:   kc.BatchBytes,
		BatchTimeout: kc.BatchTimeout,
		MaxAttempts:  kc.MaxRetryTimes,
		WriteTimeout: kc.WriteTimeout,
		RequiredAcks: kc.RequiredAcks,
}

readConfig := kafka.ReaderConfig{
		Topic:                 topic,
		Brokers:               []string{"localhost:9092", "localhost:9093", "localhost:9094"},
		GroupID:               "consumer-group-id",
		StartOffset:           kc.StartOffset,
		QueueCapacity:         kc.QueueCapacity,
		MinBytes:              kc.MinM * 1024 * 1024,
		MaxBytes:              kc.MaxM * 1024 * 1024,
		MaxWait:               time.Duration(kc.MaxWaitMs) * time.Millisecond,
		CommitInterval:        time.Duration(kc.CommitIntervalSec) * time.Second,
		HeartbeatInterval:     time.Duration(kc.HeartbeatInterval) * time.Second,
		SessionTimeout:        time.Duration(kc.SessionTimeoutSec) * time.Second,
		RebalanceTimeout:      time.Duration(kc.RebalanceTimeoutSec) * time.Second,
		JoinGroupBackoff:      time.Duration(kc.JoinGroupBackoffSec) * time.Second,
		RetentionTime:         time.Duration(kc.RetentionTimeSec) * time.Second,
		MaxAttempts:           kc.MaxAttempts,
		OffsetOutOfRangeError: kc.OffsetOutOfRangeError,
}

func consumeAndProcess(consumer *kafka.Reader) {
    for {
        msg, err := consumer.ReadMessage(context.Background())
        if err != nil {
            log.Fatal("failed to read message:", err)
        }

        // 处理消息
        processMessage(msg)

        // 提交偏移量
        if err := consumer.CommitMessages(context.Background(), msg); err != nil {
            log.Fatal("failed to commit message:", err)
        }
    }
}

func processMessage(msg kafka.Message) {
    // 这里添加消息处理逻辑
    log.Printf("message at topic/partition/offset %v/%v/%v: %s = %s\n", msg.Topic, msg.Partition, msg.Offset, string(msg.Key), string(msg.Value))
}

func saveOffsetToDB(topic string, partition int32, offset int64) {
    // 这里添加将偏移量保存到数据库的逻辑
}

func loadOffsetFromDB(topic string, partition int32) int64 {
    // 这里添加从数据库加载偏移量的逻辑
    return 0 // 示例返回值
}

type Balancer interface {
	// Balance receives a message and a set of available partitions and
	// returns the partition number that the message should be routed to.
	//
	// An application should refrain from using a balancer to manage multiple
	// sets of partitions (from different topics for examples), use one balancer
	// instance for each partition set, so the balancer can detect when the
	// partitions change and assume that the kafka topic has been rebalanced.
	Balance(msg Message, partitions ...int) (partition int)
}


func (rr *RoundRobin) Balance(msg Message, partitions ...int) int {
	return rr.balance(partitions)
}

func (rr *RoundRobin) balance(partitions []int) int {
	length := uint32(len(partitions))
	offset := atomic.AddUint32(&rr.offset, 1) - 1
	return partitions[offset%length]
}

// When sending synchronously and the writer's batch size is configured to be
// greater than 1, this method blocks until either a full batch can be assembled
// or the batch timeout is reached.  The batch size and timeouts are evaluated
// per partition, so the choice of Balancer can also influence the flushing
// behavior.  For example, the Hash balancer will require on average N * batch
// size messages to trigger a flush where N is the number of partitions.  The
// best way to achieve good batching behavior is to share one Writer amongst
// multiple go routines.
func (w *Writer) WriteMessages(ctx context.Context, msgs ...Message) error {
  ...
  // We use int32 here to half the memory footprint (compared to using int
	// on 64 bits architectures). We map lists of the message indexes instead
	// of the message values for the same reason, int32 is 4 bytes, vs a full
	// Message value which is 100+ bytes and contains pointers and contributes
	// to increasing GC work.
	assignments := make(map[topicPartition][]int32)

	for i, msg := range msgs {
		topic, err := w.chooseTopic(msg)
		if err != nil {
			return err
		}

		numPartitions, err := w.partitions(ctx, topic)
		if err != nil {
			return err
		}

		partition := balancer.Balance(msg, loadCachedPartitions(numPartitions)...)

		key := topicPartition{
			topic:     topic,
			partition: int32(partition),
		}
		//先根据msg和balance机制选择某个topic下的partition
		assignments[key] = append(assignments[key], int32(i))
	}

  batches := w.batchMessages(msgs, assignments)
	if w.Async {
		return nil
	}

	done := ctx.Done()
	hasErrors := false
	for batch := range batches {
		select {
		case <-done:
			return ctx.Err()
		case <-batch.done:
			if batch.err != nil {
				hasErrors = true
			}
		}
	}
  ...
}

func (w *Writer) batchMessages(messages []Message, assignments map[topicPartition][]int32) map[*writeBatch][]int32 {
	var batches map[*writeBatch][]int32
	if !w.Async {
		batches = make(map[*writeBatch][]int32, len(assignments))
	}

	w.mutex.Lock()
	defer w.mutex.Unlock()

	if w.writers == nil {
		w.writers = map[topicPartition]*partitionWriter{}
	}

  // key: topic/partition; indexes: msg id
	for key, indexes := range assignments {
		writer := w.writers[key]
		if writer == nil {
			writer = newPartitionWriter(w, key)
			w.writers[key] = writer
		}
    // 每一个partitionWriter都会维护一个batch, 当batch满时, 会flush到kafka
		wbatches := writer.writeMessages(messages, indexes)

		for batch, idxs := range wbatches {
			batches[batch] = idxs
		}
	}

	return batches
}

func newPartitionWriter(w *Writer, key topicPartition) *partitionWriter {
	writer := &partitionWriter{
		meta:  key,
		queue: newBatchQueue(10),
		w:     w,
	}
	// 启动一个携程写批
	w.spawn(writer.writeBatches)
	return writer
}

func (ptw *partitionWriter) writeBatches() {
	for {
		// 从队列中阻塞地获取一个batch
		batch := ptw.queue.Get()

		// The only time we can return nil is when the queue is closed
		// and empty. If the queue is closed that means
		// the Writer is closed so once we're here it's time to exit.
		if batch == nil {
			return
		}

		// 1. 调用 ptw.w.stats() 进行读写统计
		// 2. 写批次: ptw.w.produce(key, batch)
		// 3. close(b.done)
		ptw.writeBatch(batch)
	}
}

func newBatchQueue(initialSize int) batchQueue {
	bq := batchQueue{
		queue: make([]*writeBatch, 0, initialSize),
		mutex: &sync.Mutex{},
		cond:  &sync.Cond{},
	}

	bq.cond.L = bq.mutex

	return bq
}

func (b *batchQueue) Get() *writeBatch {
	b.cond.L.Lock()
	defer b.cond.L.Unlock()

	for len(b.queue) == 0 && !b.closed {
		b.cond.Wait()
	}

	if len(b.queue) == 0 {
		return nil
	}

	batch := b.queue[0]
	b.queue[0] = nil
	b.queue = b.queue[1:]

	return batch
}

func (b *batchQueue) Put(batch *writeBatch) bool {
	b.cond.L.Lock()
	defer b.cond.L.Unlock()
	defer b.cond.Broadcast()

	if b.closed {
		return false
	}
	b.queue = append(b.queue, batch)
	return true
}

func (ptw *partitionWriter) writeMessages(msgs []Message, indexes []int32) map[*writeBatch][]int32 {
	ptw.mutex.Lock()
	defer ptw.mutex.Unlock()

	batchSize := ptw.w.batchSize()
	batchBytes := ptw.w.batchBytes()

	var batches map[*writeBatch][]int32
	if !ptw.w.Async {
		batches = make(map[*writeBatch][]int32, 1)
	}

  //indexes: msg id
	for _, i := range indexes {
	assignMessage:
		batch := ptw.currBatch
		if batch == nil {
			batch = ptw.newWriteBatch()
			ptw.currBatch = batch
		}
		// 判断是否可以加入到当前batch
		if !batch.add(msgs[i], batchSize, batchBytes) {
			batch.trigger() // close(b.ready)
			ptw.queue.Put(batch)
			ptw.currBatch = nil
			goto assignMessage
		}
	  // 判断是否需要flush
		if batch.full(batchSize, batchBytes) {
			batch.trigger() 			// close(b.ready)
			ptw.queue.Put(batch)
			ptw.currBatch = nil
		}

		if !ptw.w.Async {
			batches[batch] = append(batches[batch], i)
		}
	}
	return batches
}

func (ptw *partitionWriter) newWriteBatch() *writeBatch {
	batch := newWriteBatch(time.Now(), ptw.w.batchTimeout())
	// 启动一个携程等待超时或者批次满了
	ptw.w.spawn(func() { ptw.awaitBatch(batch) })
	return batch
}

func (ptw *partitionWriter) awaitBatch(batch *writeBatch) {
	select {
	case <-batch.timer.C:
		ptw.mutex.Lock()
		// detach the batch from the writer if we're still attached
		// and queue for writing.
		// Only the current batch can expire, all previous batches were already written to the queue.
		// If writeMesseages locks pw.mutex after the timer fires but before this goroutine
		// can lock pw.mutex it will either have filled the batch and enqueued it which will mean
		// pw.currBatch != batch so we just move on.
		// Otherwise, we detach the batch from the ptWriter and enqueue it for writing.
		if ptw.currBatch == batch {
			ptw.queue.Put(batch)
			ptw.currBatch = nil
		}
		ptw.mutex.Unlock()
	case <-batch.ready:
		// The batch became full, it was removed from the ptwriter and its
		// ready channel was closed. We need to close the timer to avoid
		// having it leak until it expires.
		batch.timer.Stop()
	}
}