kgo: add AddConsumePartitions and RemoveConsumePartitions

There are some catches, but this isn't too bad.

Closes #475.

Author: twmb

pkg/kgo/consumer.go

@@ -749,6 +749,91 @@ func (cl *Client) AddConsumeTopics(topics ...string) {
 	cl.triggerUpdateMetadataNow("from AddConsumeTopics")
 }
 
+// AddConsumePartitions adds new partitions to be consumed at the given
+// offsets. This function works only for direct, non-regex consumers.
+func (cl *Client) AddConsumePartitions(partitions map[string]map[int32]Offset) {
+	c := &cl.consumer
+	if c.d == nil || cl.cfg.regex {
+		return
+	}
+	var topics []string
+	for t, ps := range partitions {
+		if len(ps) == 0 {
+			delete(partitions, t)
+			continue
+		}
+		topics = append(topics, t)
+	}
+	if len(partitions) == 0 {
+		return
+	}
+
+	c.mu.Lock()
+	defer c.mu.Unlock()
+
+	c.d.tps.storeTopics(topics)
+	for t, ps := range partitions {
+		if c.d.ps[t] == nil {
+			c.d.ps[t] = make(map[int32]Offset)
+		}
+		for p, o := range ps {
+			c.d.m.add(t, p)
+			c.d.ps[t][p] = o
+		}
+	}
+	cl.triggerUpdateMetadataNow("from AddConsumePartitions")
+}
+
+// RemoveConsumePartitions removes partitions from being consumed. This
+// function works only for direct, non-regex consumers.
+//
+// This method does not purge the concept of any topics from the client -- if
+// you remove all partitions from a topic that was being consumed, metadata
+// fetches will still occur for the topic. If you want to remove the topic
+// entirely, use PurgeTopicsFromClient.
+//
+// If you specified ConsumeTopics and this function removes all partitions for
+// a topic, the topic will no longer be consumed.
+func (cl *Client) RemoveConsumePartitions(partitions map[string][]int32) {
+	c := &cl.consumer
+	if c.d == nil || cl.cfg.regex {
+		return
+	}
+	for t, ps := range partitions {
+		if len(ps) == 0 {
+			delete(partitions, t)
+			continue
+		}
+	}
+	if len(partitions) == 0 {
+		return
+	}
+
+	c.mu.Lock()
+	defer c.mu.Unlock()
+
+	removeOffsets := make(map[string]map[int32]Offset, len(partitions))
+	for t, ps := range partitions {
+		removePartitionOffsets := make(map[int32]Offset, len(ps))
+		for _, p := range ps {
+			removePartitionOffsets[p] = Offset{}
+		}
+		removeOffsets[t] = removePartitionOffsets
+	}
+
+	c.assignPartitions(removeOffsets, assignInvalidateMatching, c.d.tps, fmt.Sprintf("remove of %v requested", partitions))
+	for t, ps := range partitions {
+		for _, p := range ps {
+			c.d.using.remove(t, p)
+			c.d.m.remove(t, p)
+			delete(c.d.ps[t], p)
+		}
+		if len(c.d.ps[t]) == 0 {
+			delete(c.d.ps, t)
+		}
+	}
+}
+
 // assignHow controls how assignPartitions operates.
 type assignHow int8
 

pkg/kgo/consumer_direct.go

@@ -2,10 +2,11 @@ package kgo
 
 type directConsumer struct {
 	cfg    *cfg
-	tps    *topicsPartitions // data for topics that the user assigned
-	using  mtmps             // topics we are currently using
-	m      mtmps             // mirrors cfg.topics and cfg.partitions, but can change with Purge or Add
-	reSeen map[string]bool   // topics we evaluated against regex, and whether we want them or not
+	tps    *topicsPartitions           // data for topics that the user assigned
+	using  mtmps                       // topics we are currently using
+	m      mtmps                       // mirrors cfg.topics and cfg.partitions, but can change with Purge or Add
+	ps     map[string]map[int32]Offset // mirrors cfg.partitions, changed in Purge or Add
+	reSeen map[string]bool             // topics we evaluated against regex, and whether we want them or not
 }
 
 func (c *consumer) initDirect() {
@@ -15,6 +16,7 @@ func (c *consumer) initDirect() {
 		reSeen: make(map[string]bool),
 		using:  make(mtmps),
 		m:      make(mtmps),
+		ps:     make(map[string]map[int32]Offset),
 	}
 	c.d = d
 
@@ -28,6 +30,11 @@ func (c *consumer) initDirect() {
 		for partition := range partitions {
 			d.m.add(topic, partition)
 		}
+		p := make(map[int32]Offset, len(partitions))
+		for partition, offset := range partitions {
+			p[partition] = offset
+		}
+		d.ps[topic] = p
 	}
 	for topic := range d.cfg.topics {
 		topics = append(topics, topic)
@@ -106,7 +113,7 @@ func (d *directConsumer) findNewAssignments() map[string]map[int32]Offset {
 		// we set those. We only use partitions from topics that have
 		// not been purged.
 		for topic := range d.m {
-			for partition, offset := range d.cfg.partitions[topic] {
+			for partition, offset := range d.ps[topic] {
 				toUseTopic, exists := toUse[topic]
 				if !exists {
 					toUseTopic = make(map[int32]Offset, 10)

pkg/kgo/consumer_direct_test.go

@@ -3,6 +3,7 @@ package kgo
 import (
 	"context"
 	"fmt"
+	"sort"
 	"testing"
 	"time"
 )
@@ -199,3 +200,63 @@ func TestIssue434(t *testing.T) {
 		}
 	}
 }
+
+func TestAddRemovePartitions(t *testing.T) {
+	t.Parallel()
+
+	t1, cleanup := tmpTopicPartitions(t, 2)
+	defer cleanup()
+
+	cl, _ := NewClient(
+		getSeedBrokers(),
+		UnknownTopicRetries(-1),
+		RecordPartitioner(ManualPartitioner()),
+		FetchMaxWait(100*time.Millisecond),
+	)
+	defer cl.Close()
+
+	if err := cl.ProduceSync(context.Background(),
+		&Record{Topic: t1, Partition: 0, Value: []byte("v1")},
+		&Record{Topic: t1, Partition: 1, Value: []byte("v2")},
+		&Record{Topic: t1, Partition: 1, Value: []byte("v3")},
+	).FirstErr(); err != nil {
+		t.Fatal(err)
+	}
+
+	cl.AddConsumePartitions(map[string]map[int32]Offset{
+		t1: {0: NewOffset().At(0)},
+	})
+
+	recs := cl.PollFetches(context.Background()).Records()
+	if len(recs) != 1 || string(recs[0].Value) != "v1" {
+		t.Fatalf("expected to see v1, got %v", recs)
+	}
+
+	cl.RemoveConsumePartitions(map[string][]int32{
+		t1:   {0, 1, 2},
+		"t2": {0, 1, 2},
+	})
+
+	cl.AddConsumePartitions(map[string]map[int32]Offset{
+		t1: {
+			0: NewOffset().At(0),
+			1: NewOffset().At(1),
+		},
+	})
+
+	recs = recs[:0]
+	for len(recs) < 2 {
+		recs = append(recs, cl.PollFetches(context.Background()).Records()...)
+	}
+	if len(recs) > 2 {
+		t.Fatalf("expected to see 2 records, got %v", recs)
+	}
+
+	sort.Slice(recs, func(i, j int) bool {
+		return recs[i].Partition < recs[j].Partition
+	})
+
+	if string(recs[0].Value) != "v1" || string(recs[1].Value) != "v3" {
+		t.Fatalf("expected to see v1 and v2, got %v", recs)
+	}
+}

pkg/kgo/topics_and_partitions.go

@@ -108,6 +108,20 @@ func (m mtmps) onlyt(t string) bool {
 	return exists && len(ps) == 0
 }
 
+func (m mtmps) remove(t string, p int32) {
+	if m == nil {
+		return
+	}
+	mps, exists := m[t]
+	if !exists {
+		return
+	}
+	delete(mps, p)
+	if len(mps) == 0 {
+		delete(m, t)
+	}
+}
+
 ////////////
 // PAUSED // -- types for pausing topics and partitions
 ////////////