Repository: Shopify/sarama
Pull request #1822 (2020-10-29 11:10:46)

KIP-392: Allow consumers to fetch from closest replica
Support for [KIP-392](https://cwiki.apache.org/confluence/display/KAFKA/KIP-392%3A+Allow+consumers+to+fetch+from+closest+replica) in Sarama consumers.

This is a continuation of #1696 with changes mentioned [here](https://github.com/Shopify/sarama/pull/1696#issuecomment-703320392).

Co-authored-by: "Deepak.Neralla" <Deepak.Neralla@mixpanel.com>

danp: I signed the CLA.

Patch:
diff --git a/client.go b/client.go
index f007338..6127ed7 100644
--- a/client.go
+++ b/client.go
@@ -29,6 +29,9 @@ type Client interface {
 	// Brokers returns the current set of active brokers as retrieved from cluster metadata.
 	Brokers() []*Broker
 
+	// Broker returns the active Broker if available for the broker ID.
+	Broker(brokerID int32) (*Broker, error)
+
 	// Topics returns the set of available topics as retrieved from cluster metadata.
 	Topics() ([]string, error)
 
@@ -198,6 +201,17 @@ func (client *client) Brokers() []*Broker {
 	return brokers
 }
 
+func (client *client) Broker(brokerID int32) (*Broker, error) {
+	client.lock.RLock()
+	defer client.lock.RUnlock()
+	broker, ok := client.brokers[brokerID]
+	if !ok {
+		return nil, ErrBrokerNotFound
+	}
+	_ = broker.Open(client.conf)
+	return broker, nil
+}
+
 func (client *client) InitProducerID() (*InitProducerIDResponse, error) {
 	var err error
 	for broker := client.any(); broker != nil; broker = client.any() {
diff --git a/consumer.go b/consumer.go
index fbdbff2..9bd8d18 100644
--- a/consumer.go
+++ b/consumer.go
@@ -303,6 +303,8 @@ type partitionConsumer struct {
 	errors   chan *ConsumerError
 	feeder   chan *FetchResponse
 
+	preferredReadReplica int32
+
 	trigger, dying chan none
 	closeOnce      sync.Once
 	topic          string
@@ -363,18 +365,29 @@ func (child *partitionConsumer) dispatcher() {
 	close(child.feeder)
 }
 
+func (child *partitionConsumer) preferredBroker() (*Broker, error) {
+	if child.preferredReadReplica >= 0 {
+		broker, err := child.consumer.client.Broker(child.preferredReadReplica)
+		if err == nil {
+			return broker, nil
+		}
+	}
+
+	// if prefered replica cannot be found fallback to leader
+	return child.consumer.client.Leader(child.topic, child.partition)
+}
+
 func (child *partitionConsumer) dispatch() error {
 	if err := child.consumer.client.RefreshMetadata(child.topic); err != nil {
 		return err
 	}
 
-	var leader *Broker
-	var err error
-	if leader, err = child.consumer.client.Leader(child.topic, child.partition); err != nil {
+	broker, err := child.preferredBroker()
+	if err != nil {
 		return err
 	}
 
-	child.broker = child.consumer.refBrokerConsumer(leader)
+	child.broker = child.consumer.refBrokerConsumer(broker)
 
 	child.broker.input <- child
 
@@ -593,6 +606,8 @@ func (child *partitionConsumer) parseResponse(response *FetchResponse) ([]*Consu
 
 	consumerBatchSizeMetric.Update(int64(nRecs))
 
+	child.preferredReadReplica = block.PreferredReadReplica
+
 	if nRecs == 0 {
 		partialTrailingMessage, err := block.isPartial()
 		if err != nil {
@@ -823,9 +838,19 @@ func (bc *brokerConsumer) handleResponses() {
 		result := child.responseResult
 		child.responseResult = nil
 
+		if result == nil {
+			if child.preferredReadReplica >= 0 && bc.broker.ID() != child.preferredReadReplica {
+				// not an error but needs redispatching to consume from prefered replica
+				child.trigger <- none{}
+				delete(bc.subscriptions, child)
+			}
+			continue
+		}
+
+		// Discard any replica preference.
+		child.preferredReadReplica = -1
+
 		switch result {
-		case nil:
-			// no-op
 		case errTimedOut:
 			Logger.Printf("consumer/broker/%d abandoned subscription to %s/%d because consuming was taking too long\n",
 				bc.broker.ID(), child.topic, child.partition)
diff --git a/errors.go b/errors.go
index ca621b0..5781c1c 100644
--- a/errors.go
+++ b/errors.go
@@ -9,6 +9,9 @@ import (
 // or otherwise failed to respond.
 var ErrOutOfBrokers = errors.New("kafka: client has run out of available brokers to talk to (Is your cluster reachable?)")
 
+// ErrBrokerNotFound is the error returned when there's no broker found for the requested ID.
+var ErrBrokerNotFound = errors.New("kafka: broker for ID is not found")
+
 // ErrClosedClient is the error returned when a method is called on a client that has been closed.
 var ErrClosedClient = errors.New("kafka: tried to use a client that was closed")
 
diff --git a/fetch_response.go b/fetch_response.go
index ca6d788..54b8828 100644
--- a/fetch_response.go
+++ b/fetch_response.go
@@ -89,6 +89,8 @@ func (b *FetchResponseBlock) decode(pd packetDecoder, version int16) (err error)
 		if err != nil {
 			return err
 		}
+	} else {
+		b.PreferredReadReplica = -1
 	}
 
 	recordsSize, err := pd.getInt32()


Test patch:
diff --git a/client_test.go b/client_test.go
index 797e2ac..512b094 100644
--- a/client_test.go
+++ b/client_test.go
@@ -576,6 +576,42 @@ func TestClientRefreshMetadataBrokerOffline(t *testing.T) {
 	}
 }
 
+func TestClientGetBroker(t *testing.T) {
+	seedBroker := NewMockBroker(t, 1)
+	leader := NewMockBroker(t, 5)
+
+	metadataResponse1 := new(MetadataResponse)
+	metadataResponse1.AddBroker(leader.Addr(), leader.BrokerID())
+	metadataResponse1.AddBroker(seedBroker.Addr(), seedBroker.BrokerID())
+	seedBroker.Returns(metadataResponse1)
+
+	client, err := NewClient([]string{seedBroker.Addr()}, NewTestConfig())
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	broker, err := client.Broker(leader.BrokerID())
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if broker.Addr() != leader.Addr() {
+		t.Errorf("Expected broker to have address %s, found %s", leader.Addr(), broker.Addr())
+	}
+
+	metadataResponse2 := new(MetadataResponse)
+	metadataResponse2.AddBroker(seedBroker.Addr(), seedBroker.BrokerID())
+	seedBroker.Returns(metadataResponse2)
+
+	if err := client.RefreshMetadata(); err != nil {
+		t.Error(err)
+	}
+	broker, err = client.Broker(leader.BrokerID())
+	if err != ErrBrokerNotFound {
+		t.Errorf("Expected Broker(brokerID) to return %v found %v", ErrBrokerNotFound, err)
+	}
+}
+
 func TestClientResurrectDeadSeeds(t *testing.T) {
 	initialSeed := NewMockBroker(t, 0)
 	emptyMetadata := new(MetadataResponse)
diff --git a/consumer_test.go b/consumer_test.go
index 9b48d8c..4c8aee1 100644
--- a/consumer_test.go
+++ b/consumer_test.go
@@ -619,6 +619,280 @@ func TestConsumeMessageWithSessionIDs(t *testing.T) {
 	}
 }
 
+func TestConsumeMessagesFromReadReplica(t *testing.T) {
+	// Given
+	fetchResponse1 := &FetchResponse{Version: 11}
+	fetchResponse1.AddMessage("my_topic", 0, nil, testMsg, 1)
+	fetchResponse1.AddMessage("my_topic", 0, nil, testMsg, 2)
+	block1 := fetchResponse1.GetBlock("my_topic", 0)
+	block1.PreferredReadReplica = -1
+
+	fetchResponse2 := &FetchResponse{Version: 11}
+	// Create a block with no records.
+	block2 := fetchResponse1.getOrCreateBlock("my_topic", 0)
+	block2.PreferredReadReplica = 1
+
+	fetchResponse3 := &FetchResponse{Version: 11}
+	fetchResponse3.AddMessage("my_topic", 0, nil, testMsg, 3)
+	fetchResponse3.AddMessage("my_topic", 0, nil, testMsg, 4)
+	block3 := fetchResponse3.GetBlock("my_topic", 0)
+	block3.PreferredReadReplica = -1
+
+	fetchResponse4 := &FetchResponse{Version: 11}
+	fetchResponse4.AddMessage("my_topic", 0, nil, testMsg, 5)
+	fetchResponse4.AddMessage("my_topic", 0, nil, testMsg, 6)
+	block4 := fetchResponse4.GetBlock("my_topic", 0)
+	block4.PreferredReadReplica = -1
+
+	cfg := NewConfig()
+	cfg.Version = V2_3_0_0
+	cfg.RackID = "consumer_rack"
+
+	leader := NewMockBroker(t, 0)
+	broker0 := NewMockBroker(t, 1)
+
+	leader.SetHandlerByMap(map[string]MockResponse{
+		"MetadataRequest": NewMockMetadataResponse(t).
+			SetBroker(broker0.Addr(), broker0.BrokerID()).
+			SetBroker(leader.Addr(), leader.BrokerID()).
+			SetLeader("my_topic", 0, leader.BrokerID()),
+		"OffsetRequest": NewMockOffsetResponse(t).
+			SetVersion(1).
+			SetOffset("my_topic", 0, OffsetNewest, 1234).
+			SetOffset("my_topic", 0, OffsetOldest, 0),
+		"FetchRequest": NewMockSequence(fetchResponse1, fetchResponse2),
+	})
+
+	broker0.SetHandlerByMap(map[string]MockResponse{
+		"MetadataRequest": NewMockMetadataResponse(t).
+			SetBroker(broker0.Addr(), broker0.BrokerID()).
+			SetBroker(leader.Addr(), leader.BrokerID()).
+			SetLeader("my_topic", 0, leader.BrokerID()),
+		"OffsetRequest": NewMockOffsetResponse(t).
+			SetVersion(1).
+			SetOffset("my_topic", 0, OffsetNewest, 1234).
+			SetOffset("my_topic", 0, OffsetOldest, 0),
+		"FetchRequest": NewMockSequence(fetchResponse3, fetchResponse4),
+	})
+
+	master, err := NewConsumer([]string{broker0.Addr()}, cfg)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	// When
+	consumer, err := master.ConsumePartition("my_topic", 0, 1)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	assertMessageOffset(t, <-consumer.Messages(), 1)
+	assertMessageOffset(t, <-consumer.Messages(), 2)
+	assertMessageOffset(t, <-consumer.Messages(), 3)
+	assertMessageOffset(t, <-consumer.Messages(), 4)
+
+	safeClose(t, consumer)
+	safeClose(t, master)
+	broker0.Close()
+	leader.Close()
+}
+
+func TestConsumeMessagesFromReadReplicaLeaderFallback(t *testing.T) {
+	// Given
+	fetchResponse1 := &FetchResponse{Version: 11}
+	fetchResponse1.AddMessage("my_topic", 0, nil, testMsg, 1)
+	fetchResponse1.AddMessage("my_topic", 0, nil, testMsg, 2)
+	block1 := fetchResponse1.GetBlock("my_topic", 0)
+	block1.PreferredReadReplica = 5 // Does not exist.
+
+	fetchResponse2 := &FetchResponse{Version: 11}
+	fetchResponse2.AddMessage("my_topic", 0, nil, testMsg, 3)
+	fetchResponse2.AddMessage("my_topic", 0, nil, testMsg, 4)
+	block2 := fetchResponse2.GetBlock("my_topic", 0)
+	block2.PreferredReadReplica = -1
+
+	cfg := NewConfig()
+	cfg.Version = V2_3_0_0
+	cfg.RackID = "consumer_rack"
+
+	leader := NewMockBroker(t, 0)
+
+	leader.SetHandlerByMap(map[string]MockResponse{
+		"MetadataRequest": NewMockMetadataResponse(t).
+			SetBroker(leader.Addr(), leader.BrokerID()).
+			SetLeader("my_topic", 0, leader.BrokerID()),
+		"OffsetRequest": NewMockOffsetResponse(t).
+			SetVersion(1).
+			SetOffset("my_topic", 0, OffsetNewest, 1234).
+			SetOffset("my_topic", 0, OffsetOldest, 0),
+		"FetchRequest": NewMockSequence(fetchResponse1, fetchResponse2),
+	})
+
+	master, err := NewConsumer([]string{leader.Addr()}, cfg)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	// When
+	consumer, err := master.ConsumePartition("my_topic", 0, 1)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	assertMessageOffset(t, <-consumer.Messages(), 1)
+	assertMessageOffset(t, <-consumer.Messages(), 2)
+	assertMessageOffset(t, <-consumer.Messages(), 3)
+	assertMessageOffset(t, <-consumer.Messages(), 4)
+
+	safeClose(t, consumer)
+	safeClose(t, master)
+	leader.Close()
+}
+
+func TestConsumeMessagesFromReadReplicaErrorReplicaNotAvailable(t *testing.T) {
+	// Given
+	fetchResponse1 := &FetchResponse{Version: 11}
+	block1 := fetchResponse1.getOrCreateBlock("my_topic", 0)
+	block1.PreferredReadReplica = 1
+
+	fetchResponse2 := &FetchResponse{Version: 11}
+	fetchResponse2.AddMessage("my_topic", 0, nil, testMsg, 1)
+	fetchResponse2.AddMessage("my_topic", 0, nil, testMsg, 2)
+	block2 := fetchResponse2.GetBlock("my_topic", 0)
+	block2.PreferredReadReplica = -1
+
+	fetchResponse3 := &FetchResponse{Version: 11}
+	fetchResponse3.AddError("my_topic", 0, ErrReplicaNotAvailable)
+
+	fetchResponse4 := &FetchResponse{Version: 11}
+	fetchResponse4.AddMessage("my_topic", 0, nil, testMsg, 3)
+	fetchResponse4.AddMessage("my_topic", 0, nil, testMsg, 4)
+
+	cfg := NewConfig()
+	cfg.Version = V2_3_0_0
+	cfg.RackID = "consumer_rack"
+
+	leader := NewMockBroker(t, 0)
+	broker0 := NewMockBroker(t, 1)
+
+	leader.SetHandlerByMap(map[string]MockResponse{
+		"MetadataRequest": NewMockMetadataResponse(t).
+			SetBroker(broker0.Addr(), broker0.BrokerID()).
+			SetBroker(leader.Addr(), leader.BrokerID()).
+			SetLeader("my_topic", 0, leader.BrokerID()),
+		"OffsetRequest": NewMockOffsetResponse(t).
+			SetVersion(1).
+			SetOffset("my_topic", 0, OffsetNewest, 1234).
+			SetOffset("my_topic", 0, OffsetOldest, 0),
+		"FetchRequest": NewMockSequence(fetchResponse1, fetchResponse4),
+	})
+
+	broker0.SetHandlerByMap(map[string]MockResponse{
+		"MetadataRequest": NewMockMetadataResponse(t).
+			SetBroker(broker0.Addr(), broker0.BrokerID()).
+			SetBroker(leader.Addr(), leader.BrokerID()).
+			SetLeader("my_topic", 0, leader.BrokerID()),
+		"OffsetRequest": NewMockOffsetResponse(t).
+			SetVersion(1).
+			SetOffset("my_topic", 0, OffsetNewest, 1234).
+			SetOffset("my_topic", 0, OffsetOldest, 0),
+		"FetchRequest": NewMockSequence(fetchResponse2, fetchResponse3),
+	})
+
+	master, err := NewConsumer([]string{broker0.Addr()}, cfg)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	// When
+	consumer, err := master.ConsumePartition("my_topic", 0, 1)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	assertMessageOffset(t, <-consumer.Messages(), 1)
+	assertMessageOffset(t, <-consumer.Messages(), 2)
+	assertMessageOffset(t, <-consumer.Messages(), 3)
+	assertMessageOffset(t, <-consumer.Messages(), 4)
+
+	safeClose(t, consumer)
+	safeClose(t, master)
+	broker0.Close()
+	leader.Close()
+}
+
+func TestConsumeMessagesFromReadReplicaErrorUnknown(t *testing.T) {
+	// Given
+	fetchResponse1 := &FetchResponse{Version: 11}
+	block1 := fetchResponse1.getOrCreateBlock("my_topic", 0)
+	block1.PreferredReadReplica = 1
+
+	fetchResponse2 := &FetchResponse{Version: 11}
+	fetchResponse2.AddMessage("my_topic", 0, nil, testMsg, 1)
+	fetchResponse2.AddMessage("my_topic", 0, nil, testMsg, 2)
+	block2 := fetchResponse2.GetBlock("my_topic", 0)
+	block2.PreferredReadReplica = -1
+
+	fetchResponse3 := &FetchResponse{Version: 11}
+	fetchResponse3.AddError("my_topic", 0, ErrUnknown)
+
+	fetchResponse4 := &FetchResponse{Version: 11}
+	fetchResponse4.AddMessage("my_topic", 0, nil, testMsg, 3)
+	fetchResponse4.AddMessage("my_topic", 0, nil, testMsg, 4)
+
+	cfg := NewConfig()
+	cfg.Version = V2_3_0_0
+	cfg.RackID = "consumer_rack"
+
+	leader := NewMockBroker(t, 0)
+	broker0 := NewMockBroker(t, 1)
+
+	leader.SetHandlerByMap(map[string]MockResponse{
+		"MetadataRequest": NewMockMetadataResponse(t).
+			SetBroker(broker0.Addr(), broker0.BrokerID()).
+			SetBroker(leader.Addr(), leader.BrokerID()).
+			SetLeader("my_topic", 0, leader.BrokerID()),
+		"OffsetRequest": NewMockOffsetResponse(t).
+			SetVersion(1).
+			SetOffset("my_topic", 0, OffsetNewest, 1234).
+			SetOffset("my_topic", 0, OffsetOldest, 0),
+		"FetchRequest": NewMockSequence(fetchResponse1, fetchResponse4),
+	})
+
+	broker0.SetHandlerByMap(map[string]MockResponse{
+		"MetadataRequest": NewMockMetadataResponse(t).
+			SetBroker(broker0.Addr(), broker0.BrokerID()).
+			SetBroker(leader.Addr(), leader.BrokerID()).
+			SetLeader("my_topic", 0, leader.BrokerID()),
+		"OffsetRequest": NewMockOffsetResponse(t).
+			SetVersion(1).
+			SetOffset("my_topic", 0, OffsetNewest, 1234).
+			SetOffset("my_topic", 0, OffsetOldest, 0),
+		"FetchRequest": NewMockSequence(fetchResponse2, fetchResponse3),
+	})
+
+	master, err := NewConsumer([]string{broker0.Addr()}, cfg)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	// When
+	consumer, err := master.ConsumePartition("my_topic", 0, 1)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	assertMessageOffset(t, <-consumer.Messages(), 1)
+	assertMessageOffset(t, <-consumer.Messages(), 2)
+	assertMessageOffset(t, <-consumer.Messages(), 3)
+	assertMessageOffset(t, <-consumer.Messages(), 4)
+
+	safeClose(t, consumer)
+	safeClose(t, master)
+	broker0.Close()
+	leader.Close()
+}
+
 // It is fine if offsets of fetched messages are not sequential (although
 // strictly increasing!).
 func TestConsumerNonSequentialOffsets(t *testing.T) {
diff --git a/fetch_response_test.go b/fetch_response_test.go
index e471581..99ad769 100644
--- a/fetch_response_test.go
+++ b/fetch_response_test.go
@@ -195,6 +195,9 @@ func TestOneMessageFetchResponse(t *testing.T) {
 	if block.HighWaterMarkOffset != 0x10101010 {
 		t.Error("Decoding didn't produce correct high water mark offset.")
 	}
+	if block.PreferredReadReplica != -1 {
+		t.Error("Decoding didn't produce correct preferred read replica.")
+	}
 	partial, err := block.isPartial()
 	if err != nil {
 		t.Fatalf("Unexpected error: %v", err)
@@ -308,6 +311,9 @@ func TestOneRecordFetchResponse(t *testing.T) {
 	if block.HighWaterMarkOffset != 0x10101010 {
 		t.Error("Decoding didn't produce correct high water mark offset.")
 	}
+	if block.PreferredReadReplica != -1 {
+		t.Error("Decoding didn't produce correct preferred read replica.")
+	}
 	partial, err := block.isPartial()
 	if err != nil {
 		t.Fatalf("Unexpected error: %v", err)
@@ -354,6 +360,9 @@ func TestPartailFetchResponse(t *testing.T) {
 	if block.HighWaterMarkOffset != 0x10101010 {
 		t.Error("Decoding didn't produce correct high water mark offset.")
 	}
+	if block.PreferredReadReplica != -1 {
+		t.Error("Decoding didn't produce correct preferred read replica.")
+	}
 	partial, err := block.isPartial()
 	if err != nil {
 		t.Fatalf("Unexpected error: %v", err)
@@ -393,6 +402,9 @@ func TestOneMessageFetchResponseV4(t *testing.T) {
 	if block.HighWaterMarkOffset != 0x10101010 {
 		t.Error("Decoding didn't produce correct high water mark offset.")
 	}
+	if block.PreferredReadReplica != -1 {
+		t.Error("Decoding didn't produce correct preferred read replica.")
+	}
 	partial, err := block.isPartial()
 	if err != nil {
 		t.Fatalf("Unexpected error: %v", err)
diff --git a/functional_consumer_test.go b/functional_consumer_test.go
index 0ec9930..6364f5b 100644
--- a/functional_consumer_test.go
+++ b/functional_consumer_test.go
@@ -246,42 +246,41 @@ func produceMsgs(t *testing.T, clientVersions []KafkaVersion, codecs []Compressi
 func consumeMsgs(t *testing.T, clientVersions []KafkaVersion, producedMessages []*ProducerMessage) {
 	// Consume all produced messages with all client versions supported by the
 	// cluster.
-consumerVersionLoop:
 	for _, consVer := range clientVersions {
-		t.Logf("*** Consuming with client version %s\n", consVer)
-		// Create a partition consumer that should start from the first produced
-		// message.
-		consCfg := NewTestConfig()
-		consCfg.Version = consVer
-		c, err := NewConsumer(FunctionalTestEnv.KafkaBrokerAddrs, consCfg)
-		if err != nil {
-			t.Fatal(err)
-		}
-		defer safeClose(t, c)
-		pc, err := c.ConsumePartition("test.1", 0, producedMessages[0].Offset)
-		if err != nil {
-			t.Fatal(err)
-		}
-		defer safeClose(t, pc)
-
-		// Consume as many messages as there have been produced and make sure that
-		// order is preserved.
-		for i, prodMsg := range producedMessages {
-			select {
-			case consMsg := <-pc.Messages():
-				if consMsg.Offset != prodMsg.Offset {
-					t.Errorf("Consumed unexpected offset: version=%s, index=%d, want=%s, got=%s",
-						consVer, i, prodMsg2Str(prodMsg), consMsg2Str(consMsg))
-					continue consumerVersionLoop
-				}
-				if string(consMsg.Value) != string(prodMsg.Value.(StringEncoder)) {
-					t.Errorf("Consumed unexpected msg: version=%s, index=%d, want=%s, got=%s",
-						consVer, i, prodMsg2Str(prodMsg), consMsg2Str(consMsg))
-					continue consumerVersionLoop
+		t.Run(consVer.String(), func(t *testing.T) {
+			t.Logf("*** Consuming with client version %s\n", consVer)
+			// Create a partition consumer that should start from the first produced
+			// message.
+			consCfg := NewTestConfig()
+			consCfg.Version = consVer
+			c, err := NewConsumer(FunctionalTestEnv.KafkaBrokerAddrs, consCfg)
+			if err != nil {
+				t.Fatal(err)
+			}
+			defer safeClose(t, c)
+			pc, err := c.ConsumePartition("test.1", 0, producedMessages[0].Offset)
+			if err != nil {
+				t.Fatal(err)
+			}
+			defer safeClose(t, pc)
+
+			// Consume as many messages as there have been produced and make sure that
+			// order is preserved.
+			for i, prodMsg := range producedMessages {
+				select {
+				case consMsg := <-pc.Messages():
+					if consMsg.Offset != prodMsg.Offset {
+						t.Fatalf("Consumed unexpected offset: version=%s, index=%d, want=%s, got=%s",
+							consVer, i, prodMsg2Str(prodMsg), consMsg2Str(consMsg))
+					}
+					if string(consMsg.Value) != string(prodMsg.Value.(StringEncoder)) {
+						t.Fatalf("Consumed unexpected msg: version=%s, index=%d, want=%s, got=%s",
+							consVer, i, prodMsg2Str(prodMsg), consMsg2Str(consMsg))
+					}
+				case <-time.After(3 * time.Second):
+					t.Fatalf("Timeout waiting for: index=%d, offset=%d, msg=%s", i, prodMsg.Offset, prodMsg.Value)
 				}
-			case <-time.After(3 * time.Second):
-				t.Fatalf("Timeout waiting for: index=%d, offset=%d, msg=%s", i, prodMsg.Offset, prodMsg.Value)
 			}
-		}
+		})
 	}
 }