client: create a sink & source for all partition replicas

Previously, if using preferred read replicas, if a replica is not the
leader for any partition, then the sink&source is never created. The
leader will say "consume from 3", and the client will not have 3 saved
because it only saves the leader.

We now create a sink&source for all replicas of a partition as well.
Now, when a leader says "consume from 3", we should have 3 saved because
it will be in the Replicas portion of a MetadataResponseTopicPartition.

This also tidies up some repeated locking code for sinks and sources.

pkg/kgo/client.go

@@ -92,6 +92,15 @@ type sinkAndSource struct {
  source *source
 }
 
+func (cl *Client) allSinksAndSources(fn func(sns sinkAndSource)) {
+ cl.sinksAndSourcesMu.Lock()
+ defer cl.sinksAndSourcesMu.Unlock()
+
+ for _, sns := range cl.sinksAndSources {
+ fn(sns)
+ }
+}
+
 type hostport struct {
  host string
  port int32

pkg/kgo/consumer.go

@@ -425,13 +425,9 @@ func (cl *Client) PauseFetchPartitions(topicPartitions map[string][]int32) map[s
 // paused. Resuming topics that are not currently paused is a per-topic no-op.
 // See the documentation on PauseTfetchTopics for more details.
 func (cl *Client) ResumeFetchTopics(topics ...string) {
- defer func() {
- cl.sinksAndSourcesMu.Lock()
- for _, sns := range cl.sinksAndSources {
- sns.source.maybeConsume()
- }
- cl.sinksAndSourcesMu.Unlock()
- }()
+ defer cl.allSinksAndSources(func(sns sinkAndSource) {
+ sns.source.maybeConsume()
+ })
 
  c := &cl.consumer
  c.pausedMu.Lock()
@@ -447,13 +443,9 @@ func (cl *Client) ResumeFetchTopics(topics ...string) {
 // per-topic no-op. See the documentation on PauseFetchPartitions for more
 // details.
 func (cl *Client) ResumeFetchPartitions(topicPartitions map[string][]int32) {
- defer func() {
- cl.sinksAndSourcesMu.Lock()
- for _, sns := range cl.sinksAndSources {
- sns.source.maybeConsume()
- }
- cl.sinksAndSourcesMu.Unlock()
- }()
+ defer cl.allSinksAndSources(func(sns sinkAndSource) {
+ sns.source.maybeConsume()
+ })
 
  c := &cl.consumer
  c.pausedMu.Lock()
@@ -1123,11 +1115,9 @@ func (c *consumer) stopSession() (listOrEpochLoads, *topicsPartitions) {
  // our num-fetches manager without worrying about a source trying to
  // register itself.
 
- c.cl.sinksAndSourcesMu.Lock()
- for _, sns := range c.cl.sinksAndSources {
+ c.cl.allSinksAndSources(func(sns sinkAndSource) {
  sns.source.session.reset()
- }
- c.cl.sinksAndSourcesMu.Unlock()
+ })
 
  // At this point, if we begin fetching anew, then the sources will not
  // be using stale fetch sessions.
@@ -1166,11 +1156,9 @@ func (c *consumer) startNewSession(tps *topicsPartitions) *consumerSession {
 
  c.sessionChangeMu.Unlock()
 
- c.cl.sinksAndSourcesMu.Lock()
- for _, sns := range c.cl.sinksAndSources {
+ c.cl.allSinksAndSources(func(sns sinkAndSource) {
  sns.source.maybeConsume()
- }
- c.cl.sinksAndSourcesMu.Unlock()
+ })
 
  // At this point, any source that was not consuming becauase it saw the
  // session was stopped has been notified to potentially start consuming

pkg/kgo/metadata.go

@@ -445,6 +445,17 @@ func (cl *Client) fetchTopicMetadata(all bool, reqTopics []string) (map[string]*
  }
  cl.sinksAndSources[p.leader] = sns
  }
+ for _, replica := range partMeta.Replicas {
+ if replica < 0 {
+ continue
+ }
+ if _, exists = cl.sinksAndSources[replica]; !exists {
+ cl.sinksAndSources[replica] = sinkAndSource{
+ sink: cl.newSink(replica),
+ source: cl.newSource(replica),
+ }
+ }
+ }
  cl.sinksAndSourcesMu.Unlock()
  p.records.sink = sns.sink
  p.cursor.source = sns.source