Repository: knative-sandbox/eventing-kafka
Pull request #762 (2021-07-13 19:51:08)

[release-0.24] [Source] Optimize InitOffsets
This is an automated cherry-pick of #739

```release-note

```

google-cla[bot]: All (the pull request submitter and all commit authors) CLAs are signed, **but** one or more commits were authored or co-authored by someone other than the pull request submitter.

We need to confirm that all authors are ok with their commits being contributed to this project.  Please have them confirm that by leaving a comment that contains only `@googlebot I consent.` in this pull request.

*Note to project maintainer:* There may be cases where the author cannot leave a comment, or the comment is not properly detected as consent.  In those cases, you can manually confirm consent of the commit author(s), and set the `cla` label to `yes` (if enabled on your project).

ℹ️ **Googlers: [Go here](https://goto.google.com/prinfo/https%3A%2F%2Fgithub.com%2Fknative-sandbox%2Feventing-kafka%2Fpull%2F762) for more info**.

<!-- need_author_consent -->
knative-prow-robot: [APPROVALNOTIFIER] This PR is **NOT APPROVED**

This pull-request has been approved by: *<a href="https://github.com/knative-sandbox/eventing-kafka/pull/762#" title="Author self-approved">knative-prow-robot</a>*
To complete the [pull request process](https://git.k8s.io/community/contributors/guide/owners.md#the-code-review-process), please assign **phamilton** after the PR has been reviewed.
You can assign the PR to them by writing `/assign @phamilton` in a comment when ready.

The full list of commands accepted by this bot can be found [here](https://go.k8s.io/bot-commands?repo=knative-sandbox%2Feventing-kafka).

<details open>
Needs approval from an approver in each of these files:

- **[OWNERS](https://github.com/knative-sandbox/eventing-kafka/blob/release-0.24/OWNERS)**

Approvers can indicate their approval by writing `/approve` in a comment
Approvers can cancel approval by writing `/approve cancel` in a comment
</details>
<!-- META={"approvers":["phamilton"]} -->
knative-prow-robot: @knative-prow-robot: The following test **failed**, say `/retest` to rerun all failed tests or `/retest-required` to rerun all mandatory failed tests:

Test name | Commit | Details | Rerun command
--- | --- | --- | ---
pull-knative-sandbox-eventing-kafka-unit-tests | 9996cee818ed654eaa4ce95d641b7b2c16593959 | [link](https://prow.knative.dev/view/gs/knative-prow/pr-logs/pull/knative-sandbox_eventing-kafka/762/pull-knative-sandbox-eventing-kafka-unit-tests/1415036093350285312) | `/test pull-knative-sandbox-eventing-kafka-unit-tests`

[Full PR test history](https://gubernator.knative.dev/pr/knative-sandbox_eventing-kafka/762). [Your PR dashboard](https://gubernator.knative.dev/pr/knative-prow-robot).

<details>

Instructions for interacting with me using PR comments are available [here](https://git.k8s.io/community/contributors/guide/pull-requests.md).  If you have questions or suggestions related to my behavior, please file an issue against the [kubernetes/test-infra](https://github.com/kubernetes/test-infra/issues/new?title=Prow%20issue:) repository. I understand the commands that are listed [here](https://go.k8s.io/bot-commands).
</details>
<!-- test report -->
codecov[bot]: # [Codecov](https://codecov.io/gh/knative-sandbox/eventing-kafka/pull/762?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=knative-sandbox) Report
> :exclamation: No coverage uploaded for pull request base (`release-0.24@6d63a5f`). [Click here to learn what that means](https://docs.codecov.io/docs/error-reference?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=knative-sandbox#section-missing-base-commit).
> The diff coverage is `n/a`.

[![Impacted file tree graph](https://codecov.io/gh/knative-sandbox/eventing-kafka/pull/762/graphs/tree.svg?width=650&height=150&src=pr&token=W765CT93QC&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=knative-sandbox)](https://codecov.io/gh/knative-sandbox/eventing-kafka/pull/762?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=knative-sandbox)

```diff
@@               Coverage Diff               @@
##             release-0.24     #762   +/-   ##
===============================================
  Coverage                ?   75.56%           
===============================================
  Files                   ?      141           
  Lines                   ?     6336           
  Branches                ?        0           
===============================================
  Hits                    ?     4788           
  Misses                  ?     1322           
  Partials                ?      226           
```



------

[Continue to review full report at Codecov](https://codecov.io/gh/knative-sandbox/eventing-kafka/pull/762?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=knative-sandbox).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=knative-sandbox)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/knative-sandbox/eventing-kafka/pull/762?src=pr&el=footer&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=knative-sandbox). Last update [6d63a5f...9996cee](https://codecov.io/gh/knative-sandbox/eventing-kafka/pull/762?src=pr&el=lastupdated&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=knative-sandbox). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=knative-sandbox).

lionelvillard: /test pull-knative-sandbox-eventing-kafka-unit-tests
lionelvillard: @googlebot I consent
lionelvillard: /lgtm
/approve
knative-prow-robot: [APPROVALNOTIFIER] This PR is **APPROVED**

This pull-request has been approved by: *<a href="https://github.com/knative-sandbox/eventing-kafka/pull/762#" title="Author self-approved">knative-prow-robot</a>*, *<a href="https://github.com/knative-sandbox/eventing-kafka/pull/762#issuecomment-879389411" title="Approved">lionelvillard</a>*

The full list of commands accepted by this bot can be found [here](https://go.k8s.io/bot-commands?repo=knative-sandbox%2Feventing-kafka).

The pull request process is described [here](https://git.k8s.io/community/contributors/guide/owners.md#the-code-review-process)

<details >
Needs approval from an approver in each of these files:

- ~~[OWNERS](https://github.com/knative-sandbox/eventing-kafka/blob/release-0.24/OWNERS)~~ [lionelvillard]

Approvers can indicate their approval by writing `/approve` in a comment
Approvers can cancel approval by writing `/approve cancel` in a comment
</details>
<!-- META={"approvers":[]} -->

Patch:
diff --git a/pkg/common/kafka/sarama/batch_get_offsets.go b/pkg/common/kafka/sarama/batch_get_offsets.go
new file mode 100644
index 000000000..bb0dc16ca
--- /dev/null
+++ b/pkg/common/kafka/sarama/batch_get_offsets.go
@@ -0,0 +1,83 @@
+/*
+Copyright 2021 The Knative Authors
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package sarama
+
+import (
+	"github.com/Shopify/sarama"
+)
+
+// GetOffsets queries the cluster to get the most recent available offset at the
+// given time (in milliseconds) for the given topics and partition
+// Time should be OffsetOldest for the earliest available offset,
+// OffsetNewest for the offset of the message that will be produced next, or a time.
+//
+// See sarama.Client.GetOffset for getting the offset of a single topic/partition combination
+func GetOffsets(client sarama.Client, topicPartitions map[string][]int32, time int64) (map[string]map[int32]int64, error) {
+	if client.Closed() {
+		return nil, sarama.ErrClosedClient
+	}
+
+	version := int16(0)
+	if client.Config().Version.IsAtLeast(sarama.V0_10_1_0) {
+		version = 1
+	}
+
+	offsets := make(map[string]map[int32]int64)
+
+	// Step 1: build one OffsetRequest instance per broker.
+	requests := make(map[*sarama.Broker]*sarama.OffsetRequest)
+
+	for topic, partitions := range topicPartitions {
+		offsets[topic] = make(map[int32]int64)
+
+		for _, partitionID := range partitions {
+			broker, err := client.Leader(topic, partitionID)
+			if err != nil {
+				return nil, err
+			}
+
+			request, ok := requests[broker]
+			if !ok {
+				request = &sarama.OffsetRequest{Version: version}
+				requests[broker] = request
+			}
+
+			request.AddBlock(topic, partitionID, time, 1)
+		}
+	}
+
+	// Step 2: send requests, one per broker, and collect offsets
+	for broker, request := range requests {
+		response, err := broker.GetAvailableOffsets(request)
+
+		if err != nil {
+			return nil, err
+		}
+
+		for topic, blocks := range response.Blocks {
+			for partitionID, block := range blocks {
+				if block.Err != sarama.ErrNoError {
+					return nil, block.Err
+				}
+
+				offsets[topic][partitionID] = block.Offset
+			}
+		}
+	}
+
+	return offsets, nil
+}
diff --git a/pkg/source/client/offsets.go b/pkg/source/client/offsets.go
index 0f8a1d2fe..dac1574a0 100644
--- a/pkg/source/client/offsets.go
+++ b/pkg/source/client/offsets.go
@@ -23,6 +23,8 @@ import (
 	"github.com/Shopify/sarama"
 	"go.uber.org/zap"
 	"knative.dev/pkg/logging"
+
+	knsarama "knative.dev/eventing-kafka/pkg/common/kafka/sarama"
 )
 
 // We want to make sure that ALL consumer group offsets are set before marking
@@ -54,6 +56,12 @@ func InitOffsets(ctx context.Context, kafkaClient sarama.Client, topics []string
 		topicPartitions[topic] = partitions
 	}
 
+	// Fetch topic offsets
+	topicOffsets, err := knsarama.GetOffsets(kafkaClient, topicPartitions, sarama.OffsetNewest)
+	if err != nil {
+		return fmt.Errorf("failed to get the topic offsets: %w", err)
+	}
+
 	// Look for uninitialized offset (-1)
 	offsets, err := kafkaAdminClient.ListConsumerGroupOffsets(consumerGroup, topicPartitions)
 	if err != nil {
@@ -62,19 +70,25 @@ func InitOffsets(ctx context.Context, kafkaClient sarama.Client, topics []string
 
 	dirty := false
 	for topic, partitions := range offsets.Blocks {
-		for partition, block := range partitions {
+		for partitionID, block := range partitions {
 			if block.Offset == -1 { // not initialized?
-				// Fetch the newest offset in the topic/partition and set it in the consumer group
-				offset, err := kafkaClient.GetOffset(topic, partition, sarama.OffsetNewest)
-				if err != nil {
-					return fmt.Errorf("failed to get the offset for topic %s and partition %d: %w", topic, partition, err)
+				partitionsOffsets, ok := topicOffsets[topic]
+				if !ok {
+					// topic may have been deleted. ignore.
+					continue
+				}
+
+				offset, ok := partitionsOffsets[partitionID]
+				if !ok {
+					// partition may have been deleted. ignore.
+					continue
 				}
 
-				logging.FromContext(ctx).Infow("initializing offset", zap.String("topic", topic), zap.Int32("partition", partition), zap.Int64("offset", offset))
+				logging.FromContext(ctx).Infow("initializing offset", zap.String("topic", topic), zap.Int32("partition", partitionID), zap.Int64("offset", offset))
 
-				pm, err := offsetManager.ManagePartition(topic, partition)
+				pm, err := offsetManager.ManagePartition(topic, partitionID)
 				if err != nil {
-					return fmt.Errorf("failed to create the partition manager for topic %s and partition %d: %w", topic, partition, err)
+					return fmt.Errorf("failed to create the partition manager for topic %s and partition %d: %w", topic, partitionID, err)
 				}
 
 				pm.MarkOffset(offset, "")


Test patch:
diff --git a/pkg/common/kafka/sarama/batch_get_offsets_test.go b/pkg/common/kafka/sarama/batch_get_offsets_test.go
new file mode 100644
index 000000000..5ec05148c
--- /dev/null
+++ b/pkg/common/kafka/sarama/batch_get_offsets_test.go
@@ -0,0 +1,118 @@
+/*
+Copyright 2021 The Knative Authors
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package sarama
+
+import (
+	"testing"
+	"time"
+
+	"github.com/Shopify/sarama"
+)
+
+func TestGetOffset(t *testing.T) {
+	testCases := map[string]struct {
+		topicPartitions map[string][]int32
+		topicOffsets    map[string]map[int32]int64
+	}{
+		"no topics": {
+			topicPartitions: map[string][]int32{},
+			topicOffsets:    map[string]map[int32]int64{},
+		},
+		"one topic, one partition": {
+			topicPartitions: map[string][]int32{"my-topic": {0}},
+			topicOffsets: map[string]map[int32]int64{
+				"my-topic": {
+					0: 5,
+				},
+			},
+		},
+		"several topics, several partitions": {
+			topicPartitions: map[string][]int32{
+				"my-topic":   {0, 1},
+				"my-topic-2": {0, 1, 2},
+				"my-topic-3": {0, 1, 2, 3}},
+			topicOffsets: map[string]map[int32]int64{
+				"my-topic":   {0: 5, 1: 7},
+				"my-topic-2": {0: 5, 1: 7, 2: 9},
+				"my-topic-3": {0: 5, 1: 7, 2: 2, 3: 10},
+			},
+		},
+	}
+
+	for n, tc := range testCases {
+		t.Run(n, func(t *testing.T) {
+			broker := sarama.NewMockBroker(t, 1)
+			defer broker.Close()
+			offsetResponse := sarama.NewMockOffsetResponse(t).SetVersion(1)
+			for topic, partitions := range tc.topicOffsets {
+				for partition, offset := range partitions {
+					offsetResponse = offsetResponse.SetOffset(topic, partition, -1, offset)
+				}
+			}
+
+			metadataResponse := sarama.NewMockMetadataResponse(t).
+				SetController(broker.BrokerID()).
+				SetBroker(broker.Addr(), broker.BrokerID())
+
+			for topic, partitions := range tc.topicOffsets {
+				for partition := range partitions {
+					metadataResponse = metadataResponse.SetLeader(topic, partition, broker.BrokerID())
+				}
+			}
+
+			broker.SetHandlerByMap(map[string]sarama.MockResponse{
+				"OffsetRequest":   offsetResponse,
+				"MetadataRequest": metadataResponse,
+			})
+
+			config := sarama.NewConfig()
+			config.Version = sarama.MaxVersion
+
+			sc, err := sarama.NewClient([]string{broker.Addr()}, config)
+			if err != nil {
+				t.Errorf("unexpected error: %v", err)
+			}
+			defer sc.Close()
+
+			offsets, err := GetOffsets(sc, tc.topicPartitions, sarama.OffsetNewest)
+
+			if err != nil {
+				t.Errorf("unexpected error: %v", err)
+			}
+
+			if len(tc.topicOffsets) != len(offsets) {
+				t.Errorf("unexpected number of topics. wanted %d, got %d", len(tc.topicOffsets), len(offsets))
+			}
+
+			for topic, topicOffsets := range offsets {
+				if len(tc.topicOffsets[topic]) != len(topicOffsets) {
+					t.Errorf("unexpected number of partitions. wanted %d, got %d", len(tc.topicOffsets[topic]), len(topicOffsets))
+				}
+
+				for partitionID, offset := range topicOffsets {
+					if tc.topicOffsets[topic][partitionID] != offset {
+						t.Errorf("unexpected offset. wanted %d, got %d", tc.topicOffsets[topic][partitionID], offset)
+					}
+				}
+			}
+
+		})
+	}
+
+	// Seems like there is no way to check the broker has been really closed. So wait few seconds
+	time.Sleep(2 * time.Second)
+}