Use atomic.Bool in Go 1.19+

Instead of using our own atomicBool type, let's use the standard
atomic.Bool type in versions of Go that contain it.

Author: PleasingFungus

pkg/kgo/atomic_maybe_work.go

@@ -2,27 +2,6 @@ package kgo
 
 import "sync/atomic"
 
-// a helper type for some places
-type atomicBool uint32
-
-func (b *atomicBool) set(v bool) {
-	if v {
-		atomic.StoreUint32((*uint32)(b), 1)
-	} else {
-		atomic.StoreUint32((*uint32)(b), 0)
-	}
-}
-
-func (b *atomicBool) get() bool { return atomic.LoadUint32((*uint32)(b)) == 1 }
-
-func (b *atomicBool) swap(v bool) bool {
-	var swap uint32
-	if v {
-		swap = 1
-	}
-	return atomic.SwapUint32((*uint32)(b), swap) == 1
-}
-
 const (
 	stateUnstarted = iota
 	stateWorking

pkg/kgo/consumer_group.go

@@ -371,7 +371,7 @@ func (g *groupConsumer) manage() {
 			g.lastAssigned = nil
 			g.fetching = nil
 
-			g.leader.set(false)
+			g.leader.Store(false)
 			g.resetExternal()
 		}
 
@@ -992,13 +992,13 @@ func (g *groupConsumer) rejoin(why string) {
 // for group cancelation to return early.
 func (g *groupConsumer) joinAndSync(joinWhy string) error {
 	g.cfg.logger.Log(LogLevelInfo, "joining group", "group", g.cfg.group)
-	g.leader.set(false)
+	g.leader.Store(false)
 	g.getAndResetExternalRejoin()
 	defer func() {
 		// If we are not leader, we clear any tracking of external
 		// topics from when we were previously leader, since tracking
 		// these is just a waste.
-		if !g.leader.get() {
+		if !g.leader.Load() {
 			g.resetExternal()
 		}
 	}()
@@ -1169,7 +1169,7 @@ func (g *groupConsumer) handleJoinResp(resp *kmsg.JoinGroupResponse) (restart bo
 	leader := resp.LeaderID == resp.MemberID
 	leaderNoPlan := !leader && resp.Version <= 8 && g.cfg.instanceID != nil && strings.HasPrefix(resp.LeaderID, *g.cfg.instanceID+"-")
 	if leader {
-		g.leader.set(true)
+		g.leader.Store(true)
 		g.cfg.logger.Log(LogLevelInfo, "joined, balancing group",
 			"group", g.cfg.group,
 			"member_id", g.memberID,
@@ -1180,7 +1180,7 @@ func (g *groupConsumer) handleJoinResp(resp *kmsg.JoinGroupResponse) (restart bo
 		)
 		plan, err = g.balanceGroup(protocol, resp.Members, resp.SkipAssignment)
 	} else if leaderNoPlan {
-		g.leader.set(true)
+		g.leader.Store(true)
 		g.cfg.logger.Log(LogLevelInfo, "joined as leader but unable to balance group due to KIP-345 limitations",
 			"group", g.cfg.group,
 			"member_id", g.memberID,
@@ -1259,8 +1259,8 @@ func (g *groupConsumer) getAndResetExternalRejoin() bool {
 	if e == nil {
 		return false
 	}
-	defer e.rejoin.set(false)
-	return e.rejoin.get()
+	defer e.rejoin.Store(false)
+	return e.rejoin.Load()
 }
 
 // Runs fn over a load, not copy, of our map.
@@ -1307,7 +1307,7 @@ func (g *groupExternal) updateLatest(meta map[string]*metadataTopic) {
 			}
 		}
 		if rejoin {
-			g.rejoin.set(true)
+			g.rejoin.Store(true)
 		}
 	})
 }
@@ -1662,7 +1662,7 @@ func (g *groupConsumer) findNewAssignments() {
 		}
 	}
 
-	externalRejoin := g.leader.get() && g.getAndResetExternalRejoin()
+	externalRejoin := g.leader.Load() && g.getAndResetExternalRejoin()
 
 	if len(toChange) == 0 && !externalRejoin {
 		return
@@ -1687,7 +1687,7 @@ func (g *groupConsumer) findNewAssignments() {
 
 	if numNewTopics > 0 {
 		g.rejoin("rejoining because there are more topics to consume, our interests have changed")
-	} else if g.leader.get() {
+	} else if g.leader.Load() {
 		if len(toChange) > 0 {
 			g.rejoin("rejoining because we are the leader and noticed some topics have new partitions")
 		} else if externalRejoin {

pkg/kgo/go118.go

@@ -0,0 +1,26 @@
+//go:build !go1.19
+// +build !go1.19
+
+package kgo
+
+import "sync/atomic"
+
+type atomicBool uint32
+
+func (b *atomicBool) Store(v bool) {
+	if v {
+		atomic.StoreUint32((*uint32)(b), 1)
+	} else {
+		atomic.StoreUint32((*uint32)(b), 0)
+	}
+}
+
+func (b *atomicBool) Load() bool { return atomic.LoadUint32((*uint32)(b)) == 1 }
+
+func (b *atomicBool) Swap(v bool) bool {
+	var swap uint32
+	if v {
+		swap = 1
+	}
+	return atomic.SwapUint32((*uint32)(b), swap) == 1
+}

pkg/kgo/go119.go

@@ -0,0 +1,10 @@
+//go:build go1.19
+// +build go1.19
+
+package kgo
+
+import "sync/atomic"
+
+type atomicBool struct {
+	atomic.Bool
+}

pkg/kgo/sink.go

@@ -158,7 +158,7 @@ func (t *txnReqBuilder) add(rb *recBuf) {
 	if t.txnID == nil {
 		return
 	}
-	if rb.addedToTxn.swap(true) {
+	if rb.addedToTxn.Swap(true) {
 		return
 	}
 	if t.req == nil {
@@ -443,7 +443,7 @@ func (s *sink) doTxnReq(
 // inflight, and that it was not added to the txn and that we need to reset the
 // drain index.
 func (b *recBatch) removeFromTxn() {
-	b.owner.addedToTxn.set(false)
+	b.owner.addedToTxn.Store(false)
 	b.owner.resetBatchDrainIdx()
 	b.decInflight()
 }

pkg/kgo/txn.go

@@ -596,7 +596,7 @@ func (cl *Client) EndAndBeginTransaction(
 	var readd map[string][]int32
 	for topic, parts := range cl.producer.topics.load() {
 		for i, part := range parts.load().partitions {
-			if part.records.addedToTxn.swap(false) {
+			if part.records.addedToTxn.Swap(false) {
 				if how == EndBeginTxnUnsafe {
 					if readd == nil {
 						readd = make(map[string][]int32)
@@ -705,7 +705,7 @@ func (cl *Client) EndAndBeginTransaction(
 
 		for topic, parts := range cl.producer.topics.load() {
 			for i, part := range parts.load().partitions {
-				if part.records.addedToTxn.get() {
+				if part.records.addedToTxn.Load() {
 					readd[topic] = append(readd[topic], int32(i))
 				}
 			}
@@ -836,7 +836,7 @@ func (cl *Client) EndTransaction(ctx context.Context, commit TransactionEndTry)
 	// addedToTxn to false outside of any mutex.
 	for _, parts := range cl.producer.topics.load() {
 		for _, part := range parts.load().partitions {
-			anyAdded = part.records.addedToTxn.swap(false) || anyAdded
+			anyAdded = part.records.addedToTxn.Swap(false) || anyAdded
 		}
 	}