Repository: Shopify/sarama (go)

Issue:
 balance. concurrent map writes
##### Versions

*Please specify real version numbers or git SHAs, not just "Latest" since that changes fairly regularly.*

| Sarama | Kafka | Go |
|--------|-------|----|
| v1.37.2 | 2.2.0 | 1.18 |

##### Configuration

``` go
	s.config = sarama.NewConfig()
	s.config.Version = sarama.V2_2_0_0
	s.config.ChannelBufferSize = 10000
	s.config.ClientID = "server"

	if s.tlsConfig != nil {
		s.config.Net.TLS.Enable = true
		s.config.Net.TLS.Config = s.tlsConfig
	}

	s.config.Net.SASL.User = s.username
	s.config.Net.SASL.Password = s.password

	s.config.Metadata.RefreshFrequency = time.Second * 30

	s.config.Consumer.Retry.BackoffFunc = func(retries int) time.Duration {
		if retries >= 3 {
			return time.Second * 3
		}
		return time.Second
	}
	s.config.Consumer.Fetch.Default = 1024 * 1024
	s.config.Consumer.MaxWaitTime = 100 * time.Millisecond
	s.config.Consumer.Return.Errors = false
	s.config.Consumer.Group.Rebalance.GroupStrategies = []sarama.BalanceStrategy{sarama.BalanceStrategySticky}
	s.config.Consumer.Offsets.AutoCommit.Enable = s.autoCommit
	s.config.Consumer.Offsets.Retention = 24 * time.Hour

	s.config.Consumer.Group.Session.Timeout = 15 * time.Second
	s.config.Consumer.Group.Heartbeat.Interval = 5 * time.Second
```

##### Logs

When filing an issue please provide logs from Sarama and Kafka if at all
possible. You can set `sarama.Logger` to a `log.Logger` to capture Sarama debug
output.

<details><summary>logs: CLICK ME</summary>
<p>

```
fatal error: concurrent map writes

goroutine 291 [running]:
runtime.throw({0x36fb26e?, 0xc000000d28?})
	/src/runtime/panic.go:992 +0x76 fp=0xc0025d37f0 sp=0xc0025d37c0 pc=0xb80576
runtime.mapassign(0x3368c20, 0xc002ea8f60, 0xc002d418a0?)
	/src/runtime/map.go:686 +0x5ab fp=0xc0025d3888 sp=0xc0025d37f0 pc=0xb50f6b
github.com/Shopify/sarama.(*partitionMovements).addPartitionMovementRecord(0x4f8e410, {{0xc002d418e0, 0x1d}, 0x2}, {{0xc0042973e0, 0x2b}, {0xc0042973b0, 0x2b}})
	/go/pkg/mod/github.com/!shopify/sarama@v1.37.2/balance_strategy.go:922 +0xba fp=0xc0025d3958 sp=0xc0025d3888 pc=0x16942fa
github.com/Shopify/sarama.(*partitionMovements).movePartition(0x4f8e410, {{0xc002d418e0, 0x1d}, 0x2}, {0xc0042973e0, 0x2b}, {0xc0042973b0, 0x2b})
	/go/pkg/mod/github.com/!shopify/sarama@v1.37.2/balance_strategy.go:952 +0x3ad fp=0xc0025d3a80 sp=0xc0025d3958 pc=0x1694a0d
github.com/Shopify/sarama.(*stickyBalanceStrategy).processPartitionMovement(0x4f8e410, {{0xc002d418e0, 0x1d}, 0x2}, {0xc0042973b0, 0x2b}, 0xc002d7c540, {0xc002d35d00, 0x2, 0x2}, ...)
	/go/pkg/mod/github.com/!shopify/sarama@v1.37.2/balance_strategy.go:564 +0x11b fp=0xc0025d3b80 sp=0xc0025d3a80 pc=0x168f19b
github.com/Shopify/sarama.(*stickyBalanceStrategy).reassignPartition(0x4f8e410, {{0xc002d418e0, 0x1d}, 0x2}, 0xc002d7c540, {0xc002d35d00, 0x2, 0x2}, 0xc0025d4b60, {0xc0042973b0, ...})
	/go/pkg/mod/github.com/!shopify/sarama@v1.37.2/balance_strategy.go:558 +0x191 fp=0xc0025d3c60 sp=0xc0025d3b80 pc=0x168efd1
github.com/Shopify/sarama.(*stickyBalanceStrategy).reassignPartitionToNewConsumer(0x4f8e410, {{0xc002d418e0, 0x1d}, 0x2}, 0xc002d7c540, {0xc002d35d00, 0x2, 0x2}, 0xc0025d4b60, 0xc0025d4b90)
	/go/pkg/mod/github.com/!shopify/sarama@v1.37.2/balance_strategy.go:547 +0x217 fp=0xc0025d3d58 sp=0xc0025d3c60 pc=0x168ed57
github.com/Shopify/sarama.(*stickyBalanceStrategy).performReassignments(0x4f8e410, {0xc003cdf260, 0x3, 0x4}, 0xc002d7c540, 0xc002d7c570, {0xc002d35d00, 0x2, 0x2}, 0xc0025d4b90, ...)
	/go/pkg/mod/github.com/!shopify/sarama@v1.37.2/balance_strategy.go:530 +0x825 fp=0xc0025d3f08 sp=0xc0025d3d58 pc=0x168ea65
github.com/Shopify/sarama.(*stickyBalanceStrategy).balance(0x4f8e410, 0xc002d7c540, 0xc002d7c570, {0xc003cdf260, 0x3, 0x4}, {0x0, 0x0, 0x0}, {0xc002d35d00, ...}, ...)
	/go/pkg/mod/github.com/!shopify/sarama@v1.37.2/balance_strategy.go:322 +0x77d fp=0xc0025d46b0 sp=0xc0025d3f08 pc=0x168d03d
github.com/Shopify/sarama.(*stickyBalanceStrategy).Plan(0x4f8e410, 0xc002d7c480, 0xc002d7c4b0)
	/go/pkg/mod/github.com/!shopify/sarama@v1.37.2/balance_strategy.go:248 +0xc14 fp=0xc0025d5338 sp=0xc0025d46b0 pc=0x168b794
github.com/Shopify/sarama.(*consumerGroup).balance(0xc0000b6840, {0x3db0870, 0x4f8e410}, 0xc002d7c480)
	/go/pkg/mod/github.com/!shopify/sarama@v1.37.2/consumer_group.go:540 +0x39c fp=0xc0025d5578 sp=0xc0025d5338 pc=0x16c58fc
github.com/Shopify/sarama.(*consumerGroup).newSession(0xc0000b6840, {0x3db8e38, 0xc00049ba00}, {0xc002e3fad0, 0x1, 0x1}, {0x3db0480, 0xc000144698}, 0x4)
	/go/pkg/mod/github.com/!shopify/sarama@v1.37.2/consumer_group.go:351 +0xb74 fp=0xc0025d5aa8 sp=0xc0025d5578 pc=0x16c3014
github.com/Shopify/sarama.(*consumerGroup).Consume(0xc0000b6840, {0x3db8e38, 0xc00049ba00}, {0xc002e3fad0, 0x1, 0x1}, {0x3db0480, 0xc000144698})
	/go/pkg/mod/github.com/!shopify/sarama@v1.37.2/consumer_group.go:206 +0x2d3 fp=0xc0025d5c18 sp=0xc0025d5aa8 pc=0x16c1ab3
```

</p>
</details>

##### Problem Description


Pull request description:
fix: prevent data race in balance strategy
For every `BalanceStrategy`, it may be shared across different consumer groups in one instance. Considering every built-in balance strategy has several fields in the struct as shown in the below snippet, data race issues may happen when several consumer group rebalances simultaneously.
```
type balanceStrategy struct {
	coreFn func(plan BalanceStrategyPlan, memberIDs []string, topic string, partitions []int32)
	name   string
}
type stickyBalanceStrategy struct {
	movements partitionMovements
}
```

Even worse, I think this issue could cause
- panic
- some topics/partitions are not subscribed by the corresponding consumers, because the corresponding fields in the struct may be modified by other consumer groups

To verify this issue, a unit test was written, which fails when -race detection is enabled. 
```
func Test_stickyBalanceStrategy_Plan_data_race(t *testing.T) {
	for i := 0; i < 1000; i++ {
		go func(bs BalanceStrategy) {
			members := map[string]ConsumerGroupMemberMetadata{
				"m1": {
					Version: 3,
					Topics:  []string{"topic"},
				},
			}
			topics := map[string][]int32{
				"topic": {0, 1, 2},
			}
			_, _ = bs.Plan(members, topics)
		}(BalanceStrategySticky)
	}
}
```

To avoid this problem, it may be best not to use a singleton for each and every of BalanceStrategy. Instead, newly-added functions like `NewBalanceStrategyRange() BalanceStrategy ` are recommended so that every time an exclusive variable is assigned to each consumer group.

Fixes #2422 

Patch:
diff --git a/balance_strategy.go b/balance_strategy.go
index cf6e398..ec0a057 100644
--- a/balance_strategy.go
+++ b/balance_strategy.go
@@ -57,7 +57,8 @@ type BalanceStrategy interface {
 
 // --------------------------------------------------------------------
 
-// BalanceStrategyRange is the default and assigns partitions as ranges to consumer group members.
+// NewBalanceStrategyRange returns a range balance strategy,
+// which is the default and assigns partitions as ranges to consumer group members.
 // This follows the same logic as
 // https://kafka.apache.org/31/javadoc/org/apache/kafka/clients/consumer/RangeAssignor.html
 //
@@ -65,27 +66,33 @@ type BalanceStrategy interface {
 //
 //	M1: {T1: [0, 1, 2], T2: [0, 1, 2]}
 //	M2: {T2: [3, 4, 5], T2: [3, 4, 5]}
-var BalanceStrategyRange = &balanceStrategy{
-	name: RangeBalanceStrategyName,
-	coreFn: func(plan BalanceStrategyPlan, memberIDs []string, topic string, partitions []int32) {
-		partitionsPerConsumer := len(partitions) / len(memberIDs)
-		consumersWithExtraPartition := len(partitions) % len(memberIDs)
-
-		sort.Strings(memberIDs)
-
-		for i, memberID := range memberIDs {
-			min := i*partitionsPerConsumer + int(math.Min(float64(consumersWithExtraPartition), float64(i)))
-			extra := 0
-			if i < consumersWithExtraPartition {
-				extra = 1
+func NewBalanceStrategyRange() BalanceStrategy {
+	return &balanceStrategy{
+		name: RangeBalanceStrategyName,
+		coreFn: func(plan BalanceStrategyPlan, memberIDs []string, topic string, partitions []int32) {
+			partitionsPerConsumer := len(partitions) / len(memberIDs)
+			consumersWithExtraPartition := len(partitions) % len(memberIDs)
+
+			sort.Strings(memberIDs)
+
+			for i, memberID := range memberIDs {
+				min := i*partitionsPerConsumer + int(math.Min(float64(consumersWithExtraPartition), float64(i)))
+				extra := 0
+				if i < consumersWithExtraPartition {
+					extra = 1
+				}
+				max := min + partitionsPerConsumer + extra
+				plan.Add(memberID, topic, partitions[min:max]...)
 			}
-			max := min + partitionsPerConsumer + extra
-			plan.Add(memberID, topic, partitions[min:max]...)
-		}
-	},
+		},
+	}
 }
 
-// BalanceStrategySticky assigns partitions to members with an attempt to preserve earlier assignments
+// Deprecated: use NewBalanceStrategyRange to avoid data race issue
+var BalanceStrategyRange = NewBalanceStrategyRange()
+
+// NewBalanceStrategySticky returns a sticky balance strategy,
+// which assigns partitions to members with an attempt to preserve earlier assignments
 // while maintain a balanced partition distribution.
 // Example with topic T with six partitions (0..5) and two members (M1, M2):
 //
@@ -97,7 +104,12 @@ var BalanceStrategyRange = &balanceStrategy{
 //	M1: {T: [0, 2]}
 //	M2: {T: [1, 3]}
 //	M3: {T: [4, 5]}
-var BalanceStrategySticky = &stickyBalanceStrategy{}
+func NewBalanceStrategySticky() BalanceStrategy {
+	return &stickyBalanceStrategy{}
+}
+
+// Deprecated: use NewBalanceStrategySticky to avoid data race issue
+var BalanceStrategySticky = NewBalanceStrategySticky()
 
 // --------------------------------------------------------------------
 
@@ -331,11 +343,17 @@ func (s *stickyBalanceStrategy) balance(currentAssignment map[string][]topicPart
 	}
 }
 
-// BalanceStrategyRoundRobin assigns partitions to members in alternating order.
+// NewBalanceStrategyRoundRobin returns a round-robin balance strategy,
+// which assigns partitions to members in alternating order.
 // For example, there are two topics (t0, t1) and two consumer (m0, m1), and each topic has three partitions (p0, p1, p2):
 // M0: [t0p0, t0p2, t1p1]
 // M1: [t0p1, t1p0, t1p2]
-var BalanceStrategyRoundRobin = new(roundRobinBalancer)
+func NewBalanceStrategyRoundRobin() BalanceStrategy {
+	return new(roundRobinBalancer)
+}
+
+// Deprecated: use NewBalanceStrategyRoundRobin to avoid data race issue
+var BalanceStrategyRoundRobin = NewBalanceStrategyRoundRobin()
 
 type roundRobinBalancer struct{}
 
diff --git a/config.go b/config.go
index b070344..1933af3 100644
--- a/config.go
+++ b/config.go
@@ -294,7 +294,7 @@ type Config struct {
 				Interval time.Duration
 			}
 			Rebalance struct {
-				// Strategy for allocating topic partitions to members (default BalanceStrategyRange)
+				// Strategy for allocating topic partitions to members.
 				// Deprecated: Strategy exists for historical compatibility
 				// and should not be used. Please use GroupStrategies.
 				Strategy BalanceStrategy
@@ -302,7 +302,7 @@ type Config struct {
 				// GroupStrategies is the priority-ordered list of client-side consumer group
 				// balancing strategies that will be offered to the coordinator. The first
 				// strategy that all group members support will be chosen by the leader.
-				// default: [BalanceStrategyRange]
+				// default: [ NewBalanceStrategyRange() ]
 				GroupStrategies []BalanceStrategy
 
 				// The maximum allowed time for each worker to join the group once a rebalance has begun.
@@ -539,7 +539,7 @@ func NewConfig() *Config {
 
 	c.Consumer.Group.Session.Timeout = 10 * time.Second
 	c.Consumer.Group.Heartbeat.Interval = 3 * time.Second
-	c.Consumer.Group.Rebalance.GroupStrategies = []BalanceStrategy{BalanceStrategyRange}
+	c.Consumer.Group.Rebalance.GroupStrategies = []BalanceStrategy{NewBalanceStrategyRange()}
 	c.Consumer.Group.Rebalance.Timeout = 60 * time.Second
 	c.Consumer.Group.Rebalance.Retry.Max = 4
 	c.Consumer.Group.Rebalance.Retry.Backoff = 2 * time.Second
diff --git a/examples/consumergroup/main.go b/examples/consumergroup/main.go
index d7d2f91..88606b5 100644
--- a/examples/consumergroup/main.go
+++ b/examples/consumergroup/main.go
@@ -70,11 +70,11 @@ func main() {
 
 	switch assignor {
 	case "sticky":
-		config.Consumer.Group.Rebalance.GroupStrategies = []sarama.BalanceStrategy{sarama.BalanceStrategySticky}
+		config.Consumer.Group.Rebalance.GroupStrategies = []sarama.BalanceStrategy{sarama.NewBalanceStrategySticky()}
 	case "roundrobin":
-		config.Consumer.Group.Rebalance.GroupStrategies = []sarama.BalanceStrategy{sarama.BalanceStrategyRoundRobin}
+		config.Consumer.Group.Rebalance.GroupStrategies = []sarama.BalanceStrategy{sarama.NewBalanceStrategyRoundRobin()}
 	case "range":
-		config.Consumer.Group.Rebalance.GroupStrategies = []sarama.BalanceStrategy{sarama.BalanceStrategyRange}
+		config.Consumer.Group.Rebalance.GroupStrategies = []sarama.BalanceStrategy{sarama.NewBalanceStrategyRange()}
 	default:
 		log.Panicf("Unrecognized consumer group partition assignor: %s", assignor)
 	}
diff --git a/examples/exactly_once/main.go b/examples/exactly_once/main.go
index 61fbc04..70207c8 100644
--- a/examples/exactly_once/main.go
+++ b/examples/exactly_once/main.go
@@ -76,7 +76,7 @@ func main() {
 
 	config.Consumer.IsolationLevel = sarama.ReadCommitted
 	config.Consumer.Offsets.AutoCommit.Enable = false
-	config.Consumer.Group.Rebalance.Strategy = sarama.BalanceStrategyRoundRobin
+	config.Consumer.Group.Rebalance.Strategy = sarama.NewBalanceStrategyRoundRobin()
 
 	if oldest {
 		config.Consumer.Offsets.Initial = sarama.OffsetOldest
diff --git a/mocks/consumer.go b/mocks/consumer.go
index 37a07cf..03283f8 100644
--- a/mocks/consumer.go
+++ b/mocks/consumer.go
@@ -343,7 +343,7 @@ func (pc *PartitionConsumer) Messages() <-chan *sarama.ConsumerMessage {
 }
 
 func (pc *PartitionConsumer) HighWaterMarkOffset() int64 {
-	return atomic.LoadInt64(&pc.highWaterMarkOffset) + 1
+	return atomic.LoadInt64(&pc.highWaterMarkOffset)
 }
 
 // Pause implements the Pause method from the sarama.PartitionConsumer interface.


Test patch:
diff --git a/balance_strategy_test.go b/balance_strategy_test.go
index 452a6a3..f397a90 100644
--- a/balance_strategy_test.go
+++ b/balance_strategy_test.go
@@ -73,7 +73,7 @@ func TestBalanceStrategyRange(t *testing.T) {
 		},
 	}
 
-	strategy := BalanceStrategyRange
+	strategy := NewBalanceStrategyRange()
 	if strategy.Name() != "range" {
 		t.Errorf("Unexpected stategy name\nexpected: range\nactual: %v", strategy.Name())
 	}
@@ -96,7 +96,7 @@ func TestBalanceStrategyRange(t *testing.T) {
 }
 
 func TestBalanceStrategyRangeAssignmentData(t *testing.T) {
-	strategy := BalanceStrategyRange
+	strategy := NewBalanceStrategyRange()
 
 	members := make(map[string]ConsumerGroupMemberMetadata, 2)
 	members["consumer1"] = ConsumerGroupMemberMetadata{
@@ -177,7 +177,7 @@ func TestBalanceStrategyRoundRobin(t *testing.T) {
 		},
 	}
 
-	strategy := BalanceStrategyRoundRobin
+	strategy := NewBalanceStrategyRoundRobin()
 	if strategy.Name() != "roundrobin" {
 		t.Errorf("Unexpected strategy name\nexpected: roundrobin\nactual: %v", strategy.Name())
 	}
@@ -284,7 +284,7 @@ func Test_deserializeTopicPartitionAssignment(t *testing.T) {
 }
 
 func TestBalanceStrategyRoundRobinAssignmentData(t *testing.T) {
-	strategy := BalanceStrategyRoundRobin
+	strategy := NewBalanceStrategyRoundRobin()
 
 	members := make(map[string]ConsumerGroupMemberMetadata, 2)
 	members["consumer1"] = ConsumerGroupMemberMetadata{
@@ -2094,6 +2094,23 @@ func Test_stickyBalanceStrategy_Plan_AssignmentData(t *testing.T) {
 	}
 }
 
+func Test_stickyBalanceStrategy_Plan_data_race(t *testing.T) {
+	for i := 0; i < 1000; i++ {
+		go func(bs BalanceStrategy) {
+			members := map[string]ConsumerGroupMemberMetadata{
+				"m1": {
+					Version: 3,
+					Topics:  []string{"topic"},
+				},
+			}
+			topics := map[string][]int32{
+				"topic": {0, 1, 2},
+			}
+			_, _ = bs.Plan(members, topics)
+		}(NewBalanceStrategySticky())
+	}
+}
+
 func BenchmarkStickAssignmentWithLargeNumberOfConsumersAndTopics(b *testing.B) {
 	s := &stickyBalanceStrategy{}
 	r := rand.New(rand.NewSource(time.Now().UnixNano()))
diff --git a/config_test.go b/config_test.go
index a97ca18..1f29ee2 100644
--- a/config_test.go
+++ b/config_test.go
@@ -548,7 +548,7 @@ func TestGroupInstanceIdAndVersionValidation(t *testing.T) {
 
 func TestConsumerGroupStrategyCompatibility(t *testing.T) {
 	config := NewTestConfig()
-	config.Consumer.Group.Rebalance.Strategy = BalanceStrategySticky
+	config.Consumer.Group.Rebalance.Strategy = NewBalanceStrategySticky()
 	if err := config.Validate(); err != nil {
 		t.Error("Expected passing config validation, got ", err)
 	}
diff --git a/functional_consumer_group_test.go b/functional_consumer_group_test.go
index 6be383f..f056661 100644
--- a/functional_consumer_group_test.go
+++ b/functional_consumer_group_test.go
@@ -552,7 +552,7 @@ func (m *testFuncConsumerGroupMember) loop(topics []string) {
 
 func newTestStatefulStrategy(t *testing.T) *testStatefulStrategy {
 	return &testStatefulStrategy{
-		BalanceStrategy: BalanceStrategyRange,
+		BalanceStrategy: NewBalanceStrategyRange(),
 		t:               t,
 	}
 }
diff --git a/mocks/consumer_test.go b/mocks/consumer_test.go
index 2aa9614..4399fef 100644
--- a/mocks/consumer_test.go
+++ b/mocks/consumer_test.go
@@ -99,7 +99,7 @@ func TestConsumerHandlesExpectationsPausingResuming(t *testing.T) {
 		t.Error("Expected sarama.ErrOutOfBrokers, found:", test0_err.Err)
 	}
 
-	if pc_test0.HighWaterMarkOffset() != 1 {
+	if pc_test0.HighWaterMarkOffset() != 0 {
 		t.Error("High water mark offset with value different from the expected: ", pc_test0.HighWaterMarkOffset())
 	}
 
@@ -112,7 +112,7 @@ func TestConsumerHandlesExpectationsPausingResuming(t *testing.T) {
 		t.Error("Message was not as expected:", test1_msg)
 	}
 
-	if pc_test1.HighWaterMarkOffset() != 2 {
+	if pc_test1.HighWaterMarkOffset() != 1 {
 		t.Error("High water mark offset with value different from the expected: ", pc_test1.HighWaterMarkOffset())
 	}
 
@@ -125,7 +125,7 @@ func TestConsumerHandlesExpectationsPausingResuming(t *testing.T) {
 		t.Error("Message was not as expected:", other0_msg)
 	}
 
-	if pc_other0.HighWaterMarkOffset() != AnyOffset+2 {
+	if pc_other0.HighWaterMarkOffset() != AnyOffset+1 {
 		t.Error("High water mark offset with value different from the expected: ", pc_other0.HighWaterMarkOffset())
 	}
 
@@ -140,7 +140,7 @@ func TestConsumerHandlesExpectationsPausingResuming(t *testing.T) {
 		t.Error("Message was not as expected:", test0_msg2)
 	}
 
-	if pc_test0.HighWaterMarkOffset() != 3 {
+	if pc_test0.HighWaterMarkOffset() != 2 {
 		t.Error("High water mark offset with value different from the expected: ", pc_test0.HighWaterMarkOffset())
 	}
 }
@@ -395,6 +395,11 @@ func TestConsumerOffsetsAreManagedCorrectlyWithSpecifiedOffset(t *testing.T) {
 	if len(trm.errors) != 0 {
 		t.Errorf("Expected to not report any errors, found: %v", trm.errors)
 	}
+
+	if pc.HighWaterMarkOffset() != message2.Offset+1 {
+		diff := pc.HighWaterMarkOffset() - message2.Offset
+		t.Errorf("Difference between highwatermarkoffset and last message offset greater than 1, got: %v", diff)
+	}
 }
 
 func TestConsumerInvalidConfiguration(t *testing.T) {