kafka-go源码解析三（Reader）

概要

Reader是暴露给应用程序的接口，前一章提到的Consumer Group是集成在本类型中使用的。

之前提到的Consumer Group主要处理消费topic的相关metadata信息，如relabance，commit offset，heartbeat等。而Reader类主要负责从kafka brokers中拉取数据。Reader有两种使用模式，一是单topic单partition的情形，由Application自己去管理offset的信息；而是Consumer Group的情形。因此在代码中有兼容两种使用模式。

结构示意图

应用程序通过FetchMsg方法拉取数据，FetchMsg是Reader类型的一个核心方法，应用程序对应底层如何实现的无感知。

实际从kafka broker中拉取数据的是reader类型，reader和Reader之间的message传输是通过Msg Channel完成的。一个reader实例对应一个partition，因此如果该Reader类型需要同时处理3个partition的数据，它会关联到3个reader实例上。

而reader获取partition数据的方式是batch模式，因为这样效率更高。reader类型会单独解析出每条msg来，返回给上层应用使用。

代码

Reader类

类型

type Reader struct {
	config ReaderConfig // 配置参数

	msgs chan readerMessage // 和reader通信的channel

	// mutable fields of the reader (synchronized on the mutex)
	mutex   sync.Mutex
	join    sync.WaitGroup
	cancel  context.CancelFunc
	stop    context.CancelFunc
	done    chan struct{}
	commits chan commitRequest
	version int64 // version holds the generation of the spawned readers
	offset  int64
	lag     int64
	closed  bool

	runError chan error

	once  uint32
	stctx context.Context
	stats *readerStats // 统计类型
}

核型方法

1 读取消息，并提交commit

func (r *Reader) ReadMessage(ctx context.Context) (Message, error) {
	m, err := r.FetchMessage(ctx) // 拉msg
	if err != nil {
		return Message{}, err
	}

	if r.useConsumerGroup() { // 使用group模式拉取数据，会自动提交offset
		if err := r.CommitMessages(ctx, m); err != nil { // commit消息
			return Message{}, err
		}
	}

	return m, nil
}

2.1 读取消息

func (r *Reader) FetchMessage(ctx context.Context) (Message, error) {
	r.activateReadLag()

	for {
		r.mutex.Lock()

		if !r.closed && r.version == 0 { // 对于单topic单partition的情形，启动一个reader拉取数据
			r.start(r.getTopicPartitionOffset())
		}

		version := r.version
		r.mutex.Unlock()

		select {
		case <-ctx.Done():
			return Message{}, ctx.Err()

		case err := <-r.runError:
			return Message{}, err

		case m, ok := <-r.msgs:
			if !ok {
				return Message{}, io.EOF
			}

			if m.version >= version { // 只有单topic单partition的情形会出现大于的情况
				r.mutex.Lock()

				switch {
				case m.error != nil:
				case version == r.version: // Consumer Group的情形
					r.offset = m.message.Offset + 1
					r.lag = m.watermark - r.offset
				}

				r.mutex.Unlock()
				return m.message, m.error
			}
		}
	}
}

1.2 commit offset

func (r *Reader) CommitMessages(ctx context.Context, msgs ...Message) error {
	var errch <-chan error
	creq := commitRequest{ // 构造向coordinator commit的请求
		commits: makeCommits(msgs...),
	}

	if r.useSyncCommits() { // 处理同步的情况
		ch := make(chan error, 1)
		errch, creq.errch = ch, ch
	}

	select {
	case r.commits <- creq: // 由commitloop实现真实请求
	case <-ctx.Done():
		return ctx.Err()
	case <-r.stctx.Done():

		return io.ErrClosedPipe
	}

	if !r.useSyncCommits() { // 异步时直接返回
		return nil
	}

	select {
	case <-ctx.Done():
		return ctx.Err()
	case err := <-errch: // 同步时等待结果
		return err
	}
}

2 设置offset

func (r *Reader) SetOffset(offset int64) error {
	if r.useConsumerGroup() { // 只在单topic单partition模式下使用
		return errNotAvailableWithGroup
	}

	if r.closed {
		err = io.ErrClosedPipe
	} else if offset != r.offset { // 效率考虑
		r.offset = offset
		if r.version != 0 {
			r.start(r.getTopicPartitionOffset()) // 结束上一个版本version，开启新一个version
		}

		r.activateReadLag()
	}

	r.mutex.Unlock()
	return err
}

3 启动新一轮version中所有的readers

func (r *Reader) start(offsetsByPartition map[topicPartition]int64) {
	if r.closed {
		// don't start child reader if parent Reader is closed
		return
	}

	ctx, cancel := context.WithCancel(context.Background())

	r.cancel() // 等待上一个version所有的readers退出
	r.cancel = cancel
	r.version++ // 更新版本号

	r.join.Add(len(offsetsByPartition))
	for key, offset := range offsetsByPartition { // 遍历所有的partition，生成一对一的reader
		go func(ctx context.Context, key topicPartition, offset int64, join *sync.WaitGroup) {
			defer join.Done()

			(&reader{
				dialer:          r.config.Dialer,
				logger:          r.config.Logger,
				errorLogger:     r.config.ErrorLogger,
				brokers:         r.config.Brokers,
				topic:           key.topic,
				partition:       int(key.partition),
				minBytes:        r.config.MinBytes,
				maxBytes:        r.config.MaxBytes,
				maxWait:         r.config.MaxWait,
				backoffDelayMin: r.config.ReadBackoffMin,
				backoffDelayMax: r.config.ReadBackoffMax,
				version:         r.version,
				msgs:            r.msgs,
				stats:           r.stats,
				isolationLevel:  r.config.IsolationLevel,
				maxAttempts:     r.config.MaxAttempts,
			}).run(ctx, offset)
		}(ctx, key, offset, &r.join)
	}
}

4 主函数，后台运行

func (r *Reader) run(cg *ConsumerGroup) {
	for {
		var err error
		var gen *Generation
		for attempt := 1; attempt <= r.config.MaxAttempts; attempt++ {
			gen, err = cg.Next(r.stctx) // 获取generation
			if err == nil {
				break
			}
                       。。。// 异常处理
		}

		r.subscribe(gen.Assignments) // 处理分配的任务，见后

		gen.Start(func(ctx context.Context) {
			r.commitLoop(ctx, gen) // 后台提交commit的goroutine
		})
		gen.Start(func(ctx context.Context) { // 等待本轮generation结束
			select {
			case <-ctx.Done():
			case <-r.stctx.Done():
			}
			r.unsubscribe()
		})
	}
}

4.1 处理分配的任务

func (r *Reader) subscribe(allAssignments map[string][]PartitionAssignment) {
	offsets := make(map[topicPartition]int64) // 每个partition对应的offset
	for topic, assignments := range allAssignments {
		for _, assignment := range assignments {
			key := topicPartition{
				topic:     topic,
				partition: int32(assignment.ID),
			}
			offsets[key] = assignment.Offset
		}
	}

	r.start(offsets) // partition和reader一一对应，启动从broker拉取数据
}

4.2 后台提交offset

func (r *Reader) commitLoop(ctx context.Context, gen *Generation) {
	if r.config.CommitInterval == 0 { // 每条msg都单独提交
		r.commitLoopImmediate(ctx, gen)
	} else {
		r.commitLoopInterval(ctx, gen) // 按照时间间隔合并msg offset提交
	}
}

4.2.1 每条msg都单独提交

func (r *Reader) commitLoopImmediate(ctx context.Context, gen *Generation) {
	offsets := offsetStash{}

	for {
		select {
		case <-ctx.Done(): // 函数取消
			var errchs []chan<- error
			for hasCommits := true; hasCommits; {
				select {
				case req := <-r.commits: // 处理剩余消息的offset
					offsets.merge(req.commits)
					errchs = append(errchs, req.errch)
				default:
					hasCommits = false
				}
			}
			err := r.commitOffsetsWithRetry(gen, offsets, defaultCommitRetries) // 真实的请求函数
			for _, errch := range errchs {
				// NOTE : this will be a buffered channel and will not block.
				errch <- err
			}
			return

		case req := <-r.commits: // 收到channel中的一条提交
			offsets.merge(req.commits)
			req.errch <- r.commitOffsetsWithRetry(gen, offsets, defaultCommitRetries)
			offsets.reset()
		}
	}
}

4.2.2 按照时间间隔合并msg offset提交

func (r *Reader) commitLoopInterval(ctx context.Context, gen *Generation) {
	ticker := time.NewTicker(r.config.CommitInterval) // 定时器
	defer ticker.Stop()

	commit := func() { // 提交函数
		if err := r.commitOffsetsWithRetry(gen, offsets, defaultCommitRetries); err != nil {
			r.withErrorLogger(func(l Logger) { l.Printf(err.Error()) })
		} else {
			offsets.reset()
		}
	}

	for {
		select {
		case <-ctx.Done(): // 取消运行
			// drain the commit channel in order to prepare the final commit.
			for hasCommits := true; hasCommits; {
				select {
				case req := <-r.commits: // 处理剩余提交
					offsets.merge(req.commits)
				default:
					hasCommits = false
				}
			}
			commit()
			return

		case <-ticker.C: // 定时提交
			commit()

		case req := <-r.commits: // 合并offset
			offsets.merge(req.commits)
		}
	}
}

4.2.3 实际提交offset的函数

func (r *Reader) commitOffsetsWithRetry(gen *Generation, offsetStash offsetStash, retries int) (err error) {
	for attempt := 0; attempt < retries; attempt++ { // 添加重试机制
		if attempt != 0 {
			if !sleep(r.stctx, backoff(attempt, backoffDelayMin, backoffDelayMax)) {
				return
			}
		}

		if err = gen.CommitOffsets(offsetStash); err == nil {
			return
		}
	}

	return 
}

reader类

reader启动后台goroutine从broker中拉取数据，它和partition是一一对应的关系，拉取到的数据放入msg channel中，由Reader类进行消费。

类型与配置参数

type reader struct {
	dialer          *Dialer
	logger          Logger
	errorLogger     Logger
	brokers         []string
	topic           string
	partition       int
	minBytes        int
	maxBytes        int
	maxWait         time.Duration
	backoffDelayMin time.Duration
	backoffDelayMax time.Duration
	version         int64 // 版本号，与Reader类联合使用
	msgs            chan<- readerMessage
	stats           *readerStats
	isolationLevel  IsolationLevel // 读隔离策略
	maxAttempts     int
}

核心方法

1 拉取数据，通常后台运行

func (r *reader) run(ctx context.Context, offset int64) {
	// 死循环，即使没有最新的数据
	for attempt := 0; true; attempt++ {
		if attempt != 0 {
			if !sleep(ctx, backoff(attempt, r.backoffDelayMin, r.backoffDelayMax)) { // 如果长时间没有新数据，那么获取数据的时间间隔可能会比较长
				return
			}
		}

		conn, start, err := r.initialize(ctx, offset) // 初始化，获取leader partition的连接，以及偏移量
		switch err {
		case nil:
		。。。 // 处理各种异常情况

	readLoop:
		for { // 循环保证持续读取数据
			if !sleep(ctx, backoff(errcount, r.backoffDelayMin, r.backoffDelayMax)) {
				conn.Close()
				return
			}

			switch offset, err = r.read(ctx, offset, conn); err { // 读一个batch的数据
			case nil:
				errcount = 0
				continue
			。。。 // 处理各种异常情况
		}
	}
}

1.1 初始化，获取leader partition的连接，以及偏移量

func (r *reader) initialize(ctx context.Context, offset int64) (conn *Conn, start int64, err error) {
	for i := 0; i != len(r.brokers) && conn == nil; i++ {
		broker := r.brokers[i]
		conn, err = r.dialer.DialLeader(ctx, "tcp", broker, r.topic, r.partition) // 获取该partition leader的连接

		if first, last, err = r.readOffsets(conn); err != nil { // 读取offset参数
			conn.Close()
			conn = nil
			break
		}

		switch { // 选取偏移量
		case offset == FirstOffset:
			offset = first

		case offset == LastOffset:
			offset = last

		case offset < first:
			offset = first
		}

		r.withLogger(func(log Logger) {
			log.Printf("the kafka reader for partition %d of %s is seeking to offset %d", r.partition, r.topic, offset)
		})

		if start, err = conn.Seek(offset, SeekAbsolute); err != nil { // 设置当前connection的offset值
			conn.Close()
			conn = nil
			break
		}

		conn.SetDeadline(time.Time{})
	}

	return
}

1.2 读取msg

func (r *reader) read(ctx context.Context, offset int64, conn *Conn) (int64, error) {
	batch := conn.ReadBatchWith(ReadBatchConfig{ // 读取batch数据
		MinBytes:       r.minBytes,
		MaxBytes:       r.maxBytes,
		IsolationLevel: r.isolationLevel,
	})
	highWaterMark := batch.HighWaterMark()

	for {
		if msg, err = batch.ReadMessage(); err != nil { // 读取该batch中的一条msg
			batch.Close()
			break
		}

		if err = r.sendMessage(ctx, msg, highWaterMark); err != nil { // 将该msg放入channel中
			batch.Close()
			break
		}

		offset = msg.Offset + 1
	}
	return offset, err
}