kgo: add support for KIP-951

This was more complicated than I was anticipating -- kgo is not designed
to have partition movement *outside* of the metadata loop.

Tested against producing to / consuming from kfake, moving a partition
every second.

Author: twmb

pkg/kgo/metadata.go

@@ -379,31 +379,8 @@ func (cl *Client) updateMetadata() (retryWhy multiUpdateWhy, err error) {
 		}
 	}
 
-	// Migrating a cursor requires stopping any consumer session. If we
-	// stop a session, we need to eventually re-start any offset listing or
-	// epoch loading that was stopped. Thus, we simply merge what we
-	// stopped into what we will reload.
-	var (
-		consumerSessionStopped bool
-		reloadOffsets          listOrEpochLoads
-		tpsPrior               *topicsPartitions
-	)
-	stopConsumerSession := func() {
-		if consumerSessionStopped {
-			return
-		}
-		consumerSessionStopped = true
-		loads, tps := cl.consumer.stopSession()
-		reloadOffsets.mergeFrom(loads)
-		tpsPrior = tps
-	}
-	defer func() {
-		if consumerSessionStopped {
-			session := cl.consumer.startNewSession(tpsPrior)
-			defer session.decWorker()
-			reloadOffsets.loadWithSession(session, "resuming reload offsets after session stopped for cursor migrating in metadata")
-		}
-	}()
+	css := &consumerSessionStopper{cl: cl}
+	defer css.maybeRestart()
 
 	var missingProduceTopics []*topicPartitions
 	for _, m := range []struct {
@@ -426,8 +403,7 @@ func (cl *Client) updateMetadata() (retryWhy multiUpdateWhy, err error) {
 				priorParts,
 				newParts,
 				m.isProduce,
-				&reloadOffsets,
-				stopConsumerSession,
+				css,
 				&retryWhy,
 			)
 		}
@@ -661,8 +637,7 @@ func (cl *Client) mergeTopicPartitions(
 	l *topicPartitions,
 	mt *metadataTopic,
 	isProduce bool,
-	reloadOffsets *listOrEpochLoads,
-	stopConsumerSession func(),
+	css *consumerSessionStopper,
 	retryWhy *multiUpdateWhy,
 ) {
 	lv := *l.load() // copy so our field writes do not collide with reads
@@ -853,11 +828,7 @@ func (cl *Client) mergeTopicPartitions(
 			if isProduce {
 				oldTP.migrateProductionTo(newTP) // migration clears failing state
 			} else {
-				oldTP.migrateCursorTo(
-					newTP,
-					reloadOffsets,
-					stopConsumerSession,
-				)
+				oldTP.migrateCursorTo(newTP, css)
 			}
 		}
 	}

pkg/kgo/sink.go

@@ -538,7 +538,7 @@ func (s *sink) handleReqClientErr(req *produceRequest, err error) {
 		if updateMeta {
 			s.cl.cfg.logger.Log(LogLevelInfo, "produce request failed, triggering metadata update", "broker", logID(s.nodeID), "err", err)
 		}
-		s.handleRetryBatches(req.batches, req.backoffSeq, updateMeta, false, "failed produce request triggered metadata update")
+		s.handleRetryBatches(req.batches, nil, req.backoffSeq, updateMeta, false, "failed produce request triggered metadata update")
 
 	case errors.Is(err, ErrClientClosed):
 		s.cl.failBufferedRecords(ErrClientClosed)
@@ -601,11 +601,13 @@ func (s *sink) handleReqResp(br *broker, req *produceRequest, resp kmsg.Response
 		return
 	}
 
+	var kmove kip951move
 	var reqRetry seqRecBatches // handled at the end
 
-	pr := resp.(*kmsg.ProduceResponse)
-	for _, rTopic := range pr.Topics {
-		topic := rTopic.Topic
+	kresp := resp.(*kmsg.ProduceResponse)
+	for i := range kresp.Topics {
+		rt := &kresp.Topics[i]
+		topic := rt.Topic
 		partitions, ok := req.batches[topic]
 		if !ok {
 			s.cl.cfg.logger.Log(LogLevelError, "broker erroneously replied with topic in produce request that we did not produce to", "broker", logID(s.nodeID), "topic", topic)
@@ -618,25 +620,26 @@ func (s *sink) handleReqResp(br *broker, req *produceRequest, resp kmsg.Response
 		}
 
 		tmetrics := req.metrics[topic]
-		for _, rPartition := range rTopic.Partitions {
-			partition := rPartition.Partition
+		for j := range rt.Partitions {
+			rp := &rt.Partitions[j]
+			partition := rp.Partition
 			batch, ok := partitions[partition]
 			if !ok {
-				s.cl.cfg.logger.Log(LogLevelError, "broker erroneously replied with partition in produce request that we did not produce to", "broker", logID(s.nodeID), "topic", rTopic.Topic, "partition", partition)
+				s.cl.cfg.logger.Log(LogLevelError, "broker erroneously replied with partition in produce request that we did not produce to", "broker", logID(s.nodeID), "topic", rt.Topic, "partition", partition)
 				delete(tmetrics, partition)
 				continue // should not hit this
 			}
 			delete(partitions, partition)
 
 			retry, didProduce := s.handleReqRespBatch(
 				b,
+				&kmove,
+				kresp,
 				topic,
-				partition,
+				rp,
 				batch,
 				req.producerID,
 				req.producerEpoch,
-				rPartition.BaseOffset,
-				rPartition.ErrorCode,
 			)
 			if retry {
 				reqRetry.addSeqBatch(topic, partition, batch)
@@ -660,22 +663,22 @@ func (s *sink) handleReqResp(br *broker, req *produceRequest, resp kmsg.Response
 
 	if len(req.batches) > 0 {
 		s.cl.cfg.logger.Log(LogLevelError, "broker did not reply to all topics / partitions in the produce request! reenqueuing missing partitions", "broker", logID(s.nodeID))
-		s.handleRetryBatches(req.batches, 0, true, false, "broker did not reply to all topics in produce request")
+		s.handleRetryBatches(req.batches, nil, 0, true, false, "broker did not reply to all topics in produce request")
 	}
 	if len(reqRetry) > 0 {
-		s.handleRetryBatches(reqRetry, 0, true, true, "produce request had retry batches")
+		s.handleRetryBatches(reqRetry, &kmove, 0, true, true, "produce request had retry batches")
 	}
 }
 
 func (s *sink) handleReqRespBatch(
 	b *bytes.Buffer,
+	kmove *kip951move,
+	resp *kmsg.ProduceResponse,
 	topic string,
-	partition int32,
+	rp *kmsg.ProduceResponseTopicPartition,
 	batch seqRecBatch,
 	producerID int64,
 	producerEpoch int16,
-	baseOffset int64,
-	errorCode int16,
 ) (retry, didProduce bool) {
 	batch.owner.mu.Lock()
 	defer batch.owner.mu.Unlock()
@@ -684,25 +687,25 @@ func (s *sink) handleReqRespBatch(
 
 	debug := b != nil
 	if debug {
-		fmt.Fprintf(b, "%d{", partition)
+		fmt.Fprintf(b, "%d{", rp.Partition)
 	}
 
 	// We only ever operate on the first batch in a record buf. Batches
 	// work sequentially; if this is not the first batch then an error
 	// happened and this later batch is no longer a part of a seq chain.
 	if !batch.isOwnersFirstBatch() {
 		if debug {
-			if err := kerr.ErrorForCode(errorCode); err == nil {
+			if err := kerr.ErrorForCode(rp.ErrorCode); err == nil {
 				if nrec > 0 {
-					fmt.Fprintf(b, "skipped@%d=>%d}, ", baseOffset, baseOffset+int64(nrec))
+					fmt.Fprintf(b, "skipped@%d=>%d}, ", rp.BaseOffset, rp.BaseOffset+int64(nrec))
 				} else {
-					fmt.Fprintf(b, "skipped@%d}, ", baseOffset)
+					fmt.Fprintf(b, "skipped@%d}, ", rp.BaseOffset)
 				}
 			} else {
 				if nrec > 0 {
-					fmt.Fprintf(b, "skipped@%d,%d(%s)}, ", baseOffset, nrec, err)
+					fmt.Fprintf(b, "skipped@%d,%d(%s)}, ", rp.BaseOffset, nrec, err)
 				} else {
-					fmt.Fprintf(b, "skipped@%d(%s)}, ", baseOffset, err)
+					fmt.Fprintf(b, "skipped@%d(%s)}, ", rp.BaseOffset, err)
 				}
 			}
 		}
@@ -713,7 +716,17 @@ func (s *sink) handleReqRespBatch(
 	// at this point re-enable failing from load errors.
 	batch.canFailFromLoadErrs = true
 
-	err := kerr.ErrorForCode(errorCode)
+	// By default, we assume we errored. Non-error updates this back
+	// to true.
+	batch.owner.okOnSink = false
+
+	if moving := kmove.maybeAddProducePartition(resp, rp, batch.owner); moving {
+		fmt.Fprintf(b, "move:%d:%d@%d,%d}, ", rp.CurrentLeader.LeaderID, rp.CurrentLeader.LeaderEpoch, rp.BaseOffset, nrec)
+		batch.owner.failing = true
+		return true, false
+	}
+
+	err := kerr.ErrorForCode(rp.ErrorCode)
 	failUnknown := batch.owner.checkUnknownFailLimit(err)
 	switch {
 	case kerr.IsRetriable(err) &&
@@ -722,7 +735,7 @@ func (s *sink) handleReqRespBatch(
 		batch.tries < s.cl.cfg.recordRetries:
 
 		if debug {
-			fmt.Fprintf(b, "retrying@%d,%d(%s)}, ", baseOffset, nrec, err)
+			fmt.Fprintf(b, "retrying@%d,%d(%s)}, ", rp.BaseOffset, nrec, err)
 		}
 		return true, false
 
@@ -769,21 +782,21 @@ func (s *sink) handleReqRespBatch(
 			s.cl.cfg.logger.Log(LogLevelInfo, "batch errored, failing the producer ID",
 				"broker", logID(s.nodeID),
 				"topic", topic,
-				"partition", partition,
+				"partition", rp.Partition,
 				"producer_id", producerID,
 				"producer_epoch", producerEpoch,
 				"err", err,
 			)
 			s.cl.failProducerID(producerID, producerEpoch, err)
 
-			s.cl.finishBatch(batch.recBatch, producerID, producerEpoch, partition, baseOffset, err)
+			s.cl.finishBatch(batch.recBatch, producerID, producerEpoch, rp.Partition, rp.BaseOffset, err)
 			if debug {
-				fmt.Fprintf(b, "fatal@%d,%d(%s)}, ", baseOffset, nrec, err)
+				fmt.Fprintf(b, "fatal@%d,%d(%s)}, ", rp.BaseOffset, nrec, err)
 			}
 			return false, false
 		}
 		if s.cl.cfg.onDataLoss != nil {
-			s.cl.cfg.onDataLoss(topic, partition)
+			s.cl.cfg.onDataLoss(topic, rp.Partition)
 		}
 
 		// For OOOSN, and UnknownProducerID
@@ -799,7 +812,7 @@ func (s *sink) handleReqRespBatch(
 		s.cl.cfg.logger.Log(LogLevelInfo, fmt.Sprintf("batch errored with %s, failing the producer ID and resetting all sequence numbers", err.(*kerr.Error).Message),
 			"broker", logID(s.nodeID),
 			"topic", topic,
-			"partition", partition,
+			"partition", rp.Partition,
 			"producer_id", producerID,
 			"producer_epoch", producerEpoch,
 			"err", err,
@@ -811,15 +824,15 @@ func (s *sink) handleReqRespBatch(
 		// will reset sequence numbers appropriately.
 		s.cl.failProducerID(producerID, producerEpoch, errReloadProducerID)
 		if debug {
-			fmt.Fprintf(b, "resetting@%d,%d(%s)}, ", baseOffset, nrec, err)
+			fmt.Fprintf(b, "resetting@%d,%d(%s)}, ", rp.BaseOffset, nrec, err)
 		}
 		return true, false
 
 	case err == kerr.DuplicateSequenceNumber: // ignorable, but we should not get
 		s.cl.cfg.logger.Log(LogLevelInfo, "received unexpected duplicate sequence number, ignoring and treating batch as successful",
 			"broker", logID(s.nodeID),
 			"topic", topic,
-			"partition", partition,
+			"partition", rp.Partition,
 		)
 		err = nil
 		fallthrough
@@ -828,22 +841,21 @@ func (s *sink) handleReqRespBatch(
 			s.cl.cfg.logger.Log(LogLevelInfo, "batch in a produce request failed",
 				"broker", logID(s.nodeID),
 				"topic", topic,
-				"partition", partition,
+				"partition", rp.Partition,
 				"err", err,
 				"err_is_retryable", kerr.IsRetriable(err),
 				"max_retries_reached", !failUnknown && batch.tries >= s.cl.cfg.recordRetries,
 			)
-			batch.owner.okOnSink = false
 		} else {
 			batch.owner.okOnSink = true
 		}
-		s.cl.finishBatch(batch.recBatch, producerID, producerEpoch, partition, baseOffset, err)
+		s.cl.finishBatch(batch.recBatch, producerID, producerEpoch, rp.Partition, rp.BaseOffset, err)
 		didProduce = err == nil
 		if debug {
 			if err != nil {
-				fmt.Fprintf(b, "err@%d,%d(%s)}, ", baseOffset, nrec, err)
+				fmt.Fprintf(b, "err@%d,%d(%s)}, ", rp.BaseOffset, nrec, err)
 			} else {
-				fmt.Fprintf(b, "%d=>%d}, ", baseOffset, baseOffset+int64(nrec))
+				fmt.Fprintf(b, "%d=>%d}, ", rp.BaseOffset, rp.BaseOffset+int64(nrec))
 			}
 		}
 	}
@@ -902,6 +914,7 @@ func (cl *Client) finishBatch(batch *recBatch, producerID int64, producerEpoch i
 // we fail it and anything after it.
 func (s *sink) handleRetryBatches(
 	retry seqRecBatches,
+	kmove *kip951move,
 	backoffSeq uint32,
 	updateMeta bool, // if we should maybe update the metadata
 	canFail bool, // if records can fail if they are at limits
@@ -911,7 +924,12 @@ func (s *sink) handleRetryBatches(
 	debug := logger.Level() >= LogLevelDebug
 	var needsMetaUpdate bool
 	var shouldBackoff bool
+	if kmove != nil {
+		defer kmove.maybeBeginMove(s.cl)
+	}
+	var numRetryBatches, numMoveBatches int
 	retry.eachOwnerLocked(func(batch seqRecBatch) {
+		numRetryBatches++
 		if !batch.isOwnersFirstBatch() {
 			if debug {
 				logger.Log(LogLevelDebug, "retry batch is not the first batch in the owner, skipping result",
@@ -931,6 +949,14 @@ func (s *sink) handleRetryBatches(
 
 		batch.owner.resetBatchDrainIdx()
 
+		// Now that the batch drain index is reset, if this retry is
+		// caused from a moving batch, return early. We do  not need
+		// to backoff nor do we need to trigger a metadata update.
+		if kmove.hasRecBuf(batch.owner) {
+			numMoveBatches++
+			return
+		}
+
 		// If our first batch (seq == 0) fails with unknown topic, we
 		// retry immediately. Kafka can reply with valid metadata
 		// immediately after a topic was created, before the leaders
@@ -975,7 +1001,7 @@ func (s *sink) handleRetryBatches(
 	// If neither of these cases are true, then we entered wanting a
 	// metadata update, but the batches either were not the first batch, or
 	// the batches were concurrently failed.
-	if shouldBackoff || !updateMeta {
+	if shouldBackoff || (!updateMeta && numRetryBatches != numMoveBatches) {
 		s.maybeTriggerBackoff(backoffSeq)
 		s.maybeDrain()
 	}
@@ -1047,7 +1073,7 @@ type recBuf struct {
 	recBufsIdx int
 
 	// A concurrent metadata update can move a recBuf from one sink to
-	// another wile requests are inflight on the original sink. We do not
+	// another while requests are inflight on the original sink. We do not
 	// want to allow new requests to start on the new sink until they all
 	// finish on the old, because with some pathological request order
 	// finishing, we would allow requests to finish out of order:
@@ -1885,7 +1911,7 @@ func (b *recBatch) tryBuffer(pr promisedRec, produceVersion, maxBatchBytes int32
 //////////////
 
 func (*produceRequest) Key() int16           { return 0 }
-func (*produceRequest) MaxVersion() int16    { return 9 }
+func (*produceRequest) MaxVersion() int16    { return 10 }
 func (p *produceRequest) SetVersion(v int16) { p.version = v }
 func (p *produceRequest) GetVersion() int16  { return p.version }
 func (p *produceRequest) IsFlexible() bool   { return p.version >= 9 }

pkg/kgo/source.go

@@ -811,6 +811,11 @@ func (s *source) fetch(consumerSession *consumerSession, doneFetch chan<- struct
 	}
 
 	// The logic below here should be relatively quick.
+	//
+	// Note that fetch runs entirely in the context of a consumer session.
+	// loopFetch does not return until this function does, meaning we
+	// cannot concurrently issue a second fetch for partitions that are
+	// being processed below.
 
 	deleteReqUsedOffset := func(topic string, partition int32) {
 		t := req.usedOffsets[topic]
@@ -937,7 +942,9 @@ func (s *source) handleReqResp(br *broker, req *fetchRequest, resp *kmsg.FetchRe
 		debugWhyStripped multiUpdateWhy
 		numErrsStripped  int
 		kip320           = s.cl.supportsOffsetForLeaderEpoch()
+		kmove            kip951move
 	)
+	defer kmove.maybeBeginMove(s.cl)
 
 	strip := func(t string, p int32, err error) {
 		numErrsStripped++
@@ -998,6 +1005,10 @@ func (s *source) handleReqResp(br *broker, req *fetchRequest, resp *kmsg.FetchRe
 
 			fp := partOffset.processRespPartition(br, rp, s.cl.decompressor, s.cl.cfg.hooks)
 			if fp.Err != nil {
+				if moving := kmove.maybeAddFetchPartition(resp, rp, partOffset.from); moving {
+					strip(topic, partition, fp.Err)
+					continue
+				}
 				updateWhy.add(topic, partition, fp.Err)
 			}
 
@@ -2026,7 +2037,7 @@ func (f *fetchRequest) MaxVersion() int16 {
 	if f.disableIDs || f.session.disableIDs {
 		return 12
 	}
-	return 15
+	return 16
 }
 func (f *fetchRequest) SetVersion(v int16) { f.version = v }
 func (f *fetchRequest) GetVersion() int16  { return f.version }