From dfcd19ac9ff5410b92081398a07df80874ff4c8a Mon Sep 17 00:00:00 2001 From: Maxim Kolosov Date: Fri, 2 Oct 2026 17:50:13 -0600 Subject: [PATCH] Reject only the unroutable partitions in a ProduceSync call. Port of #99 from release-0.1 to main. Co-authored-by: Cursor --- README.md | 2 + docs/internal/tracing.md | 5 +- pkg/wgo/client.go | 121 ++++++--- pkg/wgo/client_test.go | 498 +++++++++++++++++++++++++++++++--- pkg/wgo/produce_hooks_test.go | 75 ++++- pkg/wgo/route_records_test.go | 312 +++++++++++++++++++++ 6 files changed, 936 insertions(+), 77 deletions(-) create mode 100644 pkg/wgo/route_records_test.go diff --git a/README.md b/README.md index c76f043..7703563 100644 --- a/README.md +++ b/README.md @@ -51,6 +51,8 @@ For every record the client asks a `PartitionAssignmentStrategy` for an ordered A partition's leader can briefly go missing from that map during a Metadata refresh, even though nothing is actually wrong. When the topic is known, the client picks a live agent for that partition instead of treating it as unroutable — any agent can serve any partition. That call succeeds without waiting. The refresh that excluded the leader asks for another fetch on its own, even when no Produce is in flight. Further exclusions double the wait from `OnDemandMetadataRefreshInterval` up to `MetadataRefreshInterval`. Repeated routing misses climb the same way. While on-demand refreshes keep being requested, the periodic refresh does not run, so the gap stays at the ceiling until they stop. A refresh that excludes nobody stops the follow-ups, and the next periodic refresh puts the wait back. A topic is known when another partition still has a leader, and also when this refresh listed the topic's partitions and kept none of the leaders it named. A partition Metadata reported with no leader (`Leader` below 0) is left unroutable, including when a sibling partition still has one, so that produce fails and still gets an on-demand refresh. A topic that has never appeared in Metadata, or that came back with an error, is left alone for the same reason. +`Produce` rejects only the record whose partition has no candidate. `ProduceSync` does the same within one call: a partition with a candidate is buffered and the call waits for it, and a partition with no candidate is rejected on its own and is not sent. That rejection does not fail the other records. + ### Buffering: linger by destination agent, not by partition Records are buffered through a `ClusterRecordBuffer`, which bins them by the destination agent picked at routing time, then through a per-agent `AgentRecordBuffer`, which applies a configurable linger window before flushing. Each flush ships one Produce request to one agent carrying batches for as many partitions as the buffer accumulated. diff --git a/docs/internal/tracing.md b/docs/internal/tracing.md index e31580a..7b774ce 100644 --- a/docs/internal/tracing.md +++ b/docs/internal/tracing.md @@ -30,7 +30,10 @@ franz-go uses: - **Unbuffered** fires before the caller observes a record's outcome. For `Produce` the promise is wrapped so the hook runs just before it, mirroring franz-go (unbuffered hook, then promise). For `ProduceSync` the hooks fire, in input order, on the calling goroutine — - after it has finalized every result. + after it has finalized every result. A call that rejects some records for routing and + produces the rest still fires unbuffered once per record, at that return. Spans for the + rejected records stay open until the accepted records finish. Rejected records do not + get an earlier unbuffered hook. Callers pass their tracer via the existing `WithHooks`, exactly as for a franz-go client, so this needs no new API and no OpenTelemetry dependency in the client — it relies only on diff --git a/pkg/wgo/client.go b/pkg/wgo/client.go index 46c858c..64b5363 100644 --- a/pkg/wgo/client.go +++ b/pkg/wgo/client.go @@ -290,40 +290,57 @@ func (c *WarpstreamClient) ProduceSync(ctx context.Context, records []*kgo.Recor } wg.Add(len(okRecords)) + // Last input index for each pointer. Completions write that slot only. indexOf := make(map[*kgo.Record]int, len(okRecords)) for _, idx := range okIndices { indexOf[records[idx]] = idx } + // A repeated pointer shares that slot's outcome. Copy it to the other + // positions on return, before the unbuffered hooks above. + if len(indexOf) < len(okRecords) { + defer func() { + for _, idx := range okIndices { + canon := indexOf[records[idx]] + if idx != canon { + results[idx] = results[canon] + } + } + }() + } + write := func(recs []*kgo.Record, err error) { + for _, r := range recs { + results[indexOf[r]] = kgo.ProduceResult{Record: r, Err: err} + wg.Done() + } + } - routed, err := c.routeRecords(okRecords, func(groupRecords []*kgo.Record) func(ProduceResult) { + routed, rejected := c.routeRecords(okRecords, func(groupRecords []*kgo.Record) func(ProduceResult) { return perPartitionDone(groupRecords[0].Topic, groupRecords[0].Partition, groupRecords, func(err error) { if err != nil { // Post-dispatch failure, resolved uniformly for the whole // partition group; pre-dispatch rejections never reach here. c.metrics.produceRecordsFailedTotal.Add(float64(len(groupRecords))) } - for _, r := range groupRecords { - results[indexOf[r]] = kgo.ProduceResult{Record: r, Err: err} - wg.Done() - } + write(groupRecords, err) }) }) - if err != nil { - // One record had no known candidate. Fail the whole batch - // uniformly: every ok record gets the same error. - c.metrics.produceRecordsRejectedTotal.WithLabelValues(produceRejectedNoAgentAssigned).Add(float64(len(okIndices))) - for _, i := range okIndices { - results[i] = kgo.ProduceResult{Record: records[i], Err: err} + + if len(rejected) > 0 { + for _, rg := range rejected { + c.metrics.produceRecordsRejectedTotal.WithLabelValues(produceRejectedNoAgentAssigned).Add(float64(len(rg.records))) + write(rg.records, rg.err) } + } + if len(routed) == 0 { return results } - // Stamp each record's produce time only after routing succeeds, so a failed - // produce leaves the caller's records unchanged. A single now keeps records - // buffered together on one produce timestamp. Mirrors franz-go's bufferRecord. + // Stamp only accepted records with unset timestamps, using one shared now. now := time.Now() - for _, r := range okRecords { - ensureRecordTimestamp(r, now) + for _, g := range routed { + for _, r := range g.item.records { + ensureRecordTimestamp(r, now) + } } c.buffer.MultiAdd(ctx, routed) @@ -540,43 +557,71 @@ func (c *WarpstreamClient) waitRefreshCooldown(delay, elapsed time.Duration) { } } -// routeRecords groups records by (topic, partition), stamps each group with -// its initial destination NodeID and mints the per-group done callback. -// Returns an error if any record's partition has no known candidate. -func (c *WarpstreamClient) routeRecords(records []*kgo.Record, doneFor func(groupRecords []*kgo.Record) func(ProduceResult)) ([]promised[routedTopicPartitionRecords], error) { +type rejectedTopicPartitionRecords struct { + topicPartitionRecords + err error +} + +// routeRecords routes each partition once. A partition with no candidate is +// returned unsent. The first miss requests a metadata refresh. +func (c *WarpstreamClient) routeRecords(records []*kgo.Record, doneFor func(groupRecords []*kgo.Record) func(ProduceResult)) ([]promised[routedTopicPartitionRecords], []rejectedTopicPartitionRecords) { groups := make(map[topicPartition]*promised[routedTopicPartitionRecords]) order := make([]topicPartition, 0) + var rejectedByKey map[topicPartition]int + var rejected []rejectedTopicPartitionRecords + for _, r := range records { key := topicPartition{topic: r.Topic, partition: r.Partition} - g, ok := groups[key] - if !ok { - cands := c.demoter.Candidates(r.Topic, r.Partition, 1) - if len(cands) == 0 { + if g, ok := groups[key]; ok { + g.item.records = append(g.item.records, r) + continue + } + if rejectedByKey != nil { + if i, ok := rejectedByKey[key]; ok { + rejected[i].records = append(rejected[i].records, r) + continue + } + } + + cands := c.demoter.Candidates(r.Topic, r.Partition, 1) + if len(cands) == 0 { + if rejectedByKey == nil { + rejectedByKey = make(map[topicPartition]int) c.triggerRefresh() - return nil, fmt.Errorf("no agent assigned for topic %q partition %d", r.Topic, r.Partition) } - g = &promised[routedTopicPartitionRecords]{ - item: routedTopicPartitionRecords{ - topicPartitionRecords: topicPartitionRecords{ - topic: r.Topic, - partition: r.Partition, - }, - nodeID: cands[0].NodeID, - nodeState: cands[0].State, + rejectedByKey[key] = len(rejected) + rejected = append(rejected, rejectedTopicPartitionRecords{ + topicPartitionRecords: topicPartitionRecords{ + topic: r.Topic, + partition: r.Partition, + records: []*kgo.Record{r}, }, - } - groups[key] = g - order = append(order, key) + err: fmt.Errorf("no agent assigned for topic %q partition %d", r.Topic, r.Partition), + }) + continue + } + + groups[key] = &promised[routedTopicPartitionRecords]{ + item: routedTopicPartitionRecords{ + topicPartitionRecords: topicPartitionRecords{ + topic: r.Topic, + partition: r.Partition, + records: []*kgo.Record{r}, + }, + nodeID: cands[0].NodeID, + nodeState: cands[0].State, + }, } - g.item.records = append(g.item.records, r) + order = append(order, key) } + out := make([]promised[routedTopicPartitionRecords], 0, len(order)) for _, key := range order { g := groups[key] g.done = doneFor(g.item.records) out = append(out, *g) } - return out, nil + return out, rejected } // routeRecord is the single-record specialisation of routeRecords: it diff --git a/pkg/wgo/client_test.go b/pkg/wgo/client_test.go index b7b2895..32ea564 100644 --- a/pkg/wgo/client_test.go +++ b/pkg/wgo/client_test.go @@ -313,27 +313,50 @@ func TestWarpstreamClient_ProduceSync(t *testing.T) { }) }) - t.Run("multi-record batch with an unroutable partition counts every ok record rejected", func(t *testing.T) { + t.Run("unroutable partition rejects only its own records", func(t *testing.T) { synctest.Test(t, func(t *testing.T) { - c, _, _, _ := newTestWarpstreamClient(t, topic, 1) + c, _, clusterAddr, vnet := newTestWarpstreamClient(t, topic, 2) - // One record targets an unknown topic: routeRecords fails the whole - // batch uniformly, so every ok record is counted under no_agent_assigned. + oversize := make([]byte, 2<<20) records := []*kgo.Record{ {Topic: topic, Partition: 0, Value: []byte("a"), Timestamp: time.Now()}, {Topic: "does-not-exist", Partition: 0, Value: []byte("b"), Timestamp: time.Now()}, - {Topic: topic, Partition: 0, Value: []byte("c"), Timestamp: time.Now()}, + {Topic: topic, Partition: 1, Value: oversize, Timestamp: time.Now()}, + {Topic: topic, Partition: 1, Value: []byte("c"), Timestamp: time.Now()}, } results := c.ProduceSync(t.Context(), records) - require.Len(t, results, 3) - for i := range results { - assert.ErrorContains(t, results[i].Err, "no agent assigned") + require.Len(t, results, 4) + assert.NoError(t, results[0].Err) + assert.ErrorContains(t, results[1].Err, "no agent assigned") + assert.ErrorIs(t, results[2].Err, kerr.MessageTooLarge) + assert.NoError(t, results[3].Err) + for i, r := range results { + assert.Same(t, records[i], r.Record) } - assert.Equal(t, float64(len(records)), testutil.ToFloat64(c.metrics.produceRecordsRejectedTotal.WithLabelValues(produceRejectedNoAgentAssigned))) + assert.Equal(t, float64(1), testutil.ToFloat64(c.metrics.produceRecordsRejectedTotal.WithLabelValues(produceRejectedNoAgentAssigned))) + assert.Equal(t, float64(1), testutil.ToFloat64(c.metrics.produceRecordsRejectedTotal.WithLabelValues(produceRejectedRecordTooLarge))) assert.Equal(t, float64(len(records)), testutil.ToFloat64(c.metrics.produceRecordsTotal)) - // All records are rejections, not failures. assert.Equal(t, float64(0), testutil.ToFloat64(c.metrics.produceRecordsFailedTotal)) + + consumer, err := kgo.NewClient( + kgo.SeedBrokers(clusterAddr), + kgo.Dialer(vnet.DialContext), + kgo.ConsumePartitions(map[string]map[int32]kgo.Offset{ + topic: {0: kgo.NewOffset().AtStart(), 1: kgo.NewOffset().AtStart()}, + }), + ) + require.NoError(t, err) + t.Cleanup(consumer.Close) + + fetches := consumer.PollFetches(t.Context()) + require.NoError(t, fetches.Err()) + require.Len(t, fetches.Records(), 2) + got := map[int32]string{} + for _, r := range fetches.Records() { + got[r.Partition] = string(r.Value) + } + assert.Equal(t, map[int32]string{0: "a", 1: "c"}, got) }) }) @@ -591,6 +614,379 @@ func TestWarpstreamClient_ProduceSync(t *testing.T) { } }) }) + + t.Run("all unroutable records return without producing", func(t *testing.T) { + synctest.Test(t, func(t *testing.T) { + c, _, _, _ := newTestWarpstreamClient(t, topic, 1) + records := make([]*kgo.Record, 5) + for i := range records { + records[i] = &kgo.Record{Topic: "does-not-exist", Partition: 0, Value: []byte{byte('a' + i)}} + } + + started := time.Now() + results := c.ProduceSync(t.Context(), records) + assert.Zero(t, time.Since(started)) + require.Len(t, results, len(records)) + for i, res := range results { + assert.ErrorContains(t, res.Err, "no agent assigned") + assert.Same(t, records[i], res.Record) + assert.True(t, records[i].Timestamp.IsZero()) + } + assert.Equal(t, float64(len(records)), testutil.ToFloat64(c.metrics.produceRecordsRejectedTotal.WithLabelValues(produceRejectedNoAgentAssigned))) + assert.Equal(t, float64(0), testutil.ToFloat64(c.metrics.produceRecordsFailedTotal)) + assert.Equal(t, int64(0), c.BufferedProduceRecords()) + }) + }) + + t.Run("repeated routable record pointer fills every input position", func(t *testing.T) { + synctest.Test(t, func(t *testing.T) { + c, _, _, _ := newTestWarpstreamClient(t, topic, 1) + r := &kgo.Record{Topic: topic, Partition: 0, Value: []byte("v"), Timestamp: time.Now()} + other := &kgo.Record{Topic: topic, Partition: 0, Value: []byte("other"), Timestamp: time.Now()} + + results := c.ProduceSync(t.Context(), []*kgo.Record{r, other, r, r}) + require.Len(t, results, 4) + assert.Same(t, r, results[0].Record) + assert.NoError(t, results[0].Err) + assert.Same(t, other, results[1].Record) + assert.NoError(t, results[1].Err) + assert.Same(t, r, results[2].Record) + assert.NoError(t, results[2].Err) + assert.Same(t, r, results[3].Record) + assert.NoError(t, results[3].Err) + assert.Equal(t, float64(0), testutil.ToFloat64(c.metrics.produceRecordsFailedTotal)) + }) + }) + + t.Run("repeated routable record pointer canceled fills every input position", func(t *testing.T) { + synctest.Test(t, func(t *testing.T) { + c, _, _, _ := newTestWarpstreamClient(t, topic, 1) + ctx, cancel := context.WithCancel(t.Context()) + cancel() + r := &kgo.Record{Topic: topic, Partition: 0, Value: []byte("v"), Timestamp: time.Now()} + + results := c.ProduceSync(ctx, []*kgo.Record{r, r}) + require.Len(t, results, 2) + for i, res := range results { + assert.Same(t, r, res.Record) + assert.ErrorIs(t, res.Err, context.Canceled, "result %d", i) + } + assert.Equal(t, float64(2), testutil.ToFloat64(c.metrics.produceRecordsFailedTotal)) + assert.Equal(t, int64(0), c.BufferedProduceRecords()) + }) + }) + + t.Run("repeated routable record pointer terminal failure fills every input position", func(t *testing.T) { + synctest.Test(t, func(t *testing.T) { + vnet := &kfake.VirtualNetwork{} + cluster, addr := testkafka.CreateCluster(t, 1, topic, testkafka.WithVirtualNetwork(vnet)) + c, err := NewWarpstreamClient(nil, prometheus.NewPedanticRegistry(), append( + testWarpstreamOpts(addr, topic), + WithDialer(vnet.DialContext), + WithHedgerMaxHedgeAgents(1), + WithHedgerMinHedgeDelay(time.Hour), + )...) + require.NoError(t, err) + t.Cleanup(c.Close) + + cluster.ControlKey(int16(kmsg.Produce), func(req kmsg.Request) (kmsg.Response, error, bool) { + cluster.KeepControl() + preq := req.(*kmsg.ProduceRequest) + presp := preq.ResponseKind().(*kmsg.ProduceResponse) + presp.Version = preq.Version + for _, rt := range preq.Topics { + out := kmsg.ProduceResponseTopic{Topic: rt.Topic} + if out.Topic == "" { + out.Topic = topic + } + for _, rp := range rt.Partitions { + out.Partitions = append(out.Partitions, kmsg.ProduceResponseTopicPartition{ + Partition: rp.Partition, + ErrorCode: kerr.NotLeaderForPartition.Code, + }) + } + presp.Topics = append(presp.Topics, out) + } + return presp, nil, true + }) + + r := &kgo.Record{Topic: topic, Partition: 0, Value: []byte("wire"), Timestamp: time.Now()} + results := c.ProduceSync(t.Context(), []*kgo.Record{r, r}) + require.Len(t, results, 2) + for i, res := range results { + assert.Same(t, r, res.Record) + assert.ErrorIs(t, res.Err, kgo.ErrRecordTimeout, "result %d", i) + } + assert.Equal(t, float64(2), testutil.ToFloat64(c.metrics.produceRecordsFailedTotal)) + }) + }) + + t.Run("repeated routable record pointer stays separate from a routing miss", func(t *testing.T) { + synctest.Test(t, func(t *testing.T) { + c, _, _, _ := newTestWarpstreamClient(t, topic, 1) + r := &kgo.Record{Topic: topic, Partition: 0, Value: []byte("ok"), Timestamp: time.Now()} + miss := &kgo.Record{Topic: "does-not-exist", Partition: 0, Value: []byte("m"), Timestamp: time.Now()} + + results := c.ProduceSync(t.Context(), []*kgo.Record{r, miss, r, miss, miss}) + require.Len(t, results, 5) + assert.Same(t, r, results[0].Record) + assert.NoError(t, results[0].Err) + assert.Same(t, miss, results[1].Record) + assert.ErrorContains(t, results[1].Err, "no agent assigned") + assert.Same(t, r, results[2].Record) + assert.NoError(t, results[2].Err) + assert.Same(t, miss, results[3].Record) + assert.ErrorContains(t, results[3].Err, "no agent assigned") + assert.Same(t, miss, results[4].Record) + assert.ErrorContains(t, results[4].Err, "no agent assigned") + assert.Equal(t, float64(0), testutil.ToFloat64(c.metrics.produceRecordsFailedTotal)) + assert.Equal(t, float64(3), testutil.ToFloat64(c.metrics.produceRecordsRejectedTotal.WithLabelValues(produceRejectedNoAgentAssigned))) + }) + }) + + t.Run("repeated unroutable record pointer fails every input position", func(t *testing.T) { + synctest.Test(t, func(t *testing.T) { + c, _, _, _ := newTestWarpstreamClient(t, topic, 1) + r := &kgo.Record{Topic: "does-not-exist", Partition: 0, Value: []byte("v"), Timestamp: time.Now()} + okRec := &kgo.Record{Topic: topic, Partition: 0, Value: []byte("ok"), Timestamp: time.Now()} + + repeated := c.ProduceSync(t.Context(), []*kgo.Record{r, r}) + require.Len(t, repeated, 2) + for i, res := range repeated { + assert.Same(t, r, res.Record) + assert.ErrorContains(t, res.Err, "no agent assigned", "result %d", i) + } + + mixed := c.ProduceSync(t.Context(), []*kgo.Record{r, okRec, r}) + require.Len(t, mixed, 3) + assert.Same(t, r, mixed[0].Record) + assert.ErrorContains(t, mixed[0].Err, "no agent assigned") + assert.Same(t, okRec, mixed[1].Record) + assert.NoError(t, mixed[1].Err) + assert.Same(t, r, mixed[2].Record) + assert.ErrorContains(t, mixed[2].Err, "no agent assigned") + assert.Equal(t, float64(4), testutil.ToFloat64(c.metrics.produceRecordsRejectedTotal.WithLabelValues(produceRejectedNoAgentAssigned))) + }) + }) + + t.Run("negative leader rejects only that partition", func(t *testing.T) { + synctest.Test(t, func(t *testing.T) { + c, _, clusterAddr, vnet := newTestWarpstreamClient(t, topic, 2) + cands := c.demoter.Candidates(topic, 0, 1) + require.Len(t, cands, 1) + leader := cands[0].NodeID + topicID, ok := c.pool.TopicID(topic) + require.True(t, ok) + + c.pool.state.Store(&poolState{ + agents: []int32{leader}, + topicIDs: map[string][16]byte{topic: topicID}, + strategy: newDefaultPartitionAssignmentStrategy([]int32{leader}, map[topicPartition]int32{ + {topic: topic, partition: 0}: leader, + }, nil, map[topicPartition]struct{}{ + {topic: topic, partition: 1}: {}, + }), + }) + + results := c.ProduceSync(t.Context(), []*kgo.Record{ + {Topic: topic, Partition: 0, Value: []byte("kept"), Timestamp: time.Now()}, + {Topic: topic, Partition: 1, Value: []byte("unnamed"), Timestamp: time.Now()}, + }) + require.Len(t, results, 2) + require.NoError(t, results[0].Err) + require.ErrorContains(t, results[1].Err, "no agent assigned") + assert.Equal(t, float64(1), testutil.ToFloat64(c.metrics.produceRecordsRejectedTotal.WithLabelValues(produceRejectedNoAgentAssigned))) + + consumer, err := kgo.NewClient( + kgo.SeedBrokers(clusterAddr), + kgo.Dialer(vnet.DialContext), + kgo.ConsumePartitions(map[string]map[int32]kgo.Offset{ + topic: {0: kgo.NewOffset().AtStart(), 1: kgo.NewOffset().AtStart()}, + }), + ) + require.NoError(t, err) + t.Cleanup(consumer.Close) + fetches := consumer.PollFetches(t.Context()) + require.NoError(t, fetches.Err()) + require.Len(t, fetches.Records(), 1) + assert.Equal(t, []byte("kept"), fetches.Records()[0].Value) + assert.Equal(t, int32(0), fetches.Records()[0].Partition) + }) + }) + + t.Run("empty agent pool rejects every record", func(t *testing.T) { + synctest.Test(t, func(t *testing.T) { + c, _, _, _ := newTestWarpstreamClient(t, topic, 1) + c.pool.state.Store(&poolState{ + strategy: newDefaultPartitionAssignmentStrategy(nil, nil, nil, nil), + }) + + results := c.ProduceSync(t.Context(), []*kgo.Record{ + {Topic: topic, Partition: 0, Value: []byte("a"), Timestamp: time.Now()}, + {Topic: "other", Partition: 0, Value: []byte("b"), Timestamp: time.Now()}, + }) + require.Len(t, results, 2) + assert.ErrorContains(t, results[0].Err, "no agent assigned") + assert.ErrorContains(t, results[1].Err, "no agent assigned") + assert.Equal(t, float64(2), testutil.ToFloat64(c.metrics.produceRecordsRejectedTotal.WithLabelValues(produceRejectedNoAgentAssigned))) + assert.Equal(t, int64(0), c.BufferedProduceRecords()) + }) + }) + + t.Run("routing miss stays separate from a failed accepted partition", func(t *testing.T) { + synctest.Test(t, func(t *testing.T) { + vnet := &kfake.VirtualNetwork{} + cluster, addr := testkafka.CreateCluster(t, 1, topic, testkafka.WithVirtualNetwork(vnet)) + c, err := NewWarpstreamClient(nil, prometheus.NewPedanticRegistry(), append( + testWarpstreamOpts(addr, topic), + WithDialer(vnet.DialContext), + WithHedgerMaxHedgeAgents(1), + WithHedgerMinHedgeDelay(time.Hour), + )...) + require.NoError(t, err) + t.Cleanup(c.Close) + + cluster.ControlKey(int16(kmsg.Produce), func(req kmsg.Request) (kmsg.Response, error, bool) { + cluster.KeepControl() + preq := req.(*kmsg.ProduceRequest) + presp := preq.ResponseKind().(*kmsg.ProduceResponse) + presp.Version = preq.Version + for _, rt := range preq.Topics { + out := kmsg.ProduceResponseTopic{Topic: rt.Topic} + if out.Topic == "" { + out.Topic = topic + } + for _, rp := range rt.Partitions { + out.Partitions = append(out.Partitions, kmsg.ProduceResponseTopicPartition{ + Partition: rp.Partition, + ErrorCode: kerr.NotLeaderForPartition.Code, + }) + } + presp.Topics = append(presp.Topics, out) + } + return presp, nil, true + }) + + results := c.ProduceSync(t.Context(), []*kgo.Record{ + {Topic: topic, Partition: 0, Value: []byte("wire"), Timestamp: time.Now()}, + {Topic: "does-not-exist", Partition: 0, Value: []byte("miss"), Timestamp: time.Now()}, + }) + require.Len(t, results, 2) + require.ErrorIs(t, results[0].Err, kgo.ErrRecordTimeout) + assert.NotContains(t, results[0].Err.Error(), "no agent assigned") + require.ErrorContains(t, results[1].Err, "no agent assigned") + assert.Equal(t, float64(1), testutil.ToFloat64(c.metrics.produceRecordsFailedTotal)) + assert.Equal(t, float64(1), testutil.ToFloat64(c.metrics.produceRecordsRejectedTotal.WithLabelValues(produceRejectedNoAgentAssigned))) + }) + }) + + t.Run("accepted group split across flushes completes once beside a routing miss", func(t *testing.T) { + synctest.Test(t, func(t *testing.T) { + vnet := &kfake.VirtualNetwork{} + cluster, addr := testkafka.CreateCluster(t, 1, topic, testkafka.WithVirtualNetwork(vnet)) + c, err := NewWarpstreamClient(nil, prometheus.NewPedanticRegistry(), append( + testWarpstreamOpts(addr, topic), + WithDialer(vnet.DialContext), + WithBatchMaxBytes(512), + )...) + require.NoError(t, err) + t.Cleanup(c.Close) + + var requests atomic.Int32 + cluster.ControlKey(int16(kmsg.Produce), func(req kmsg.Request) (kmsg.Response, error, bool) { + cluster.KeepControl() + requests.Add(1) + return nil, nil, false + }) + + records := make([]*kgo.Record, 6) + for i := range 5 { + records[i] = &kgo.Record{Topic: topic, Partition: 0, Value: bytes.Repeat([]byte{byte('a' + i)}, 150)} + } + records[5] = &kgo.Record{Topic: "does-not-exist", Partition: 0, Value: []byte("miss")} + results := c.ProduceSync(t.Context(), records) + require.Len(t, results, 6) + for i := range 5 { + assert.NoError(t, results[i].Err) + } + assert.ErrorContains(t, results[5].Err, "no agent assigned") + assert.GreaterOrEqual(t, requests.Load(), int32(2)) + }) + }) + + t.Run("pre-canceled mixed call keeps the routing error", func(t *testing.T) { + synctest.Test(t, func(t *testing.T) { + c, _, _, _ := newTestWarpstreamClient(t, topic, 1) + ctx, cancel := context.WithCancel(t.Context()) + cancel() + miss := &kgo.Record{Topic: "does-not-exist", Partition: 0, Value: []byte("m")} + results := c.ProduceSync(ctx, []*kgo.Record{ + {Topic: topic, Partition: 0, Value: []byte("a")}, + miss, + }) + require.Len(t, results, 2) + require.ErrorIs(t, results[0].Err, context.Canceled) + require.ErrorContains(t, results[1].Err, "no agent assigned") + assert.True(t, miss.Timestamp.IsZero()) + assert.Equal(t, int64(0), c.BufferedProduceRecords()) + assert.Equal(t, float64(1), testutil.ToFloat64(c.metrics.produceRecordsFailedTotal)) + assert.Equal(t, float64(1), testutil.ToFloat64(c.metrics.produceRecordsRejectedTotal.WithLabelValues(produceRejectedNoAgentAssigned))) + }) + }) + + t.Run("mid-flight cancel keeps the routing error", func(t *testing.T) { + synctest.Test(t, func(t *testing.T) { + c, _, _, _ := newTestWarpstreamClient(t, topic, 1) + release := make(chan struct{}) + started := make(chan struct{}) + var once sync.Once + c.SetTestProduceResponseHook(func(ctx context.Context, _ int32, _ *kmsg.ProduceResponse, _ error) { + once.Do(func() { close(started) }) + select { + case <-release: + case <-ctx.Done(): + } + }) + + ctx, cancel := context.WithCancel(t.Context()) + var results kgo.ProduceResults + done := make(chan struct{}) + go func() { + results = c.ProduceSync(ctx, []*kgo.Record{ + {Topic: topic, Partition: 0, Value: []byte("a"), Timestamp: time.Now()}, + {Topic: "does-not-exist", Partition: 0, Value: []byte("m"), Timestamp: time.Now()}, + }) + close(done) + }() + <-started + cancel() + <-done + close(release) + synctest.Wait() + + require.Len(t, results, 2) + require.ErrorIs(t, results[0].Err, context.Canceled) + require.ErrorContains(t, results[1].Err, "no agent assigned") + assert.Equal(t, int64(0), c.BufferedProduceRecords()) + }) + }) + + t.Run("closed buffer fails accepted records and keeps routing errors", func(t *testing.T) { + synctest.Test(t, func(t *testing.T) { + c, _, _, _ := newTestWarpstreamClient(t, topic, 1) + c.buffer.Close() + + results := c.ProduceSync(t.Context(), []*kgo.Record{ + {Topic: topic, Partition: 0, Value: []byte("a"), Timestamp: time.Now()}, + {Topic: "does-not-exist", Partition: 0, Value: []byte("m"), Timestamp: time.Now()}, + }) + require.Len(t, results, 2) + require.ErrorIs(t, results[0].Err, errBufferClosed) + require.ErrorContains(t, results[1].Err, "no agent assigned") + assert.Equal(t, float64(1), testutil.ToFloat64(c.metrics.produceRecordsFailedTotal)) + assert.Equal(t, float64(1), testutil.ToFloat64(c.metrics.produceRecordsRejectedTotal.WithLabelValues(produceRejectedNoAgentAssigned))) + }) + }) } func TestWarpstreamClient_RoutingFailureLeavesTimestampUnstamped(t *testing.T) { @@ -614,6 +1010,32 @@ func TestWarpstreamClient_RoutingFailureLeavesTimestampUnstamped(t *testing.T) { }) }) + t.Run("ProduceSync mixed call stamps only accepted records", func(t *testing.T) { + synctest.Test(t, func(t *testing.T) { + c, _, _, _ := newTestWarpstreamClient(t, topic, 1) + miss := &kgo.Record{Topic: "does-not-exist", Partition: 0, Value: []byte("m")} + tooBig := &kgo.Record{Topic: topic, Partition: 0, Value: make([]byte, 2<<20)} + kept := time.Date(2024, 1, 2, 3, 4, 5, 0, time.UTC) + set := &kgo.Record{Topic: topic, Partition: 0, Value: []byte("set"), Timestamp: kept} + unset := &kgo.Record{Topic: topic, Partition: 0, Value: []byte("unset")} + unset2 := &kgo.Record{Topic: topic, Partition: 0, Value: []byte("unset2")} + + results := c.ProduceSync(t.Context(), []*kgo.Record{miss, tooBig, set, unset, unset2}) + require.Len(t, results, 5) + require.ErrorContains(t, results[0].Err, "no agent assigned") + require.ErrorIs(t, results[1].Err, kerr.MessageTooLarge) + require.NoError(t, results[2].Err) + require.NoError(t, results[3].Err) + require.NoError(t, results[4].Err) + + assert.True(t, miss.Timestamp.IsZero()) + assert.True(t, tooBig.Timestamp.IsZero()) + assert.True(t, set.Timestamp.Equal(kept)) + assert.False(t, unset.Timestamp.IsZero()) + assert.True(t, unset.Timestamp.Equal(unset2.Timestamp)) + }) + }) + t.Run("Produce", func(t *testing.T) { synctest.Test(t, func(t *testing.T) { c, _, _, _ := newTestWarpstreamClient(t, topic, 1) @@ -2151,35 +2573,39 @@ func TestWarpstreamClient_OnDemandRefreshBackoff(t *testing.T) { t.Run("close during a periodic fetch does not run a queued nudge", func(t *testing.T) { synctest.Test(t, func(t *testing.T) { - c, cluster, _, _ := newTestWarpstreamClient(t, topic, 1) - release := make(chan struct{}) - var blockNext atomic.Bool - blockNext.Store(true) - cluster.ControlKey(int16(kmsg.Metadata), func(req kmsg.Request) (kmsg.Response, error, bool) { - cluster.KeepControl() - // AgentPool.Refresh asks for every topic. kgo's own loop asks - // for brokers only, and blocking that request stalls the cluster. - if req.(*kmsg.MetadataRequest).Topics == nil && blockNext.CompareAndSwap(true, false) { - cluster.SleepControl(func() { <-release }) - } - return nil, nil, false - }) + // Forty passes. One pass still succeeds about half the time if the + // refreshCtx.Err() return after a periodic fetch is removed. + for i := range 40 { + c, cluster, _, _ := newTestWarpstreamClient(t, topic, 1) + release := make(chan struct{}) + var blockNext atomic.Bool + blockNext.Store(true) + cluster.ControlKey(int16(kmsg.Metadata), func(req kmsg.Request) (kmsg.Response, error, bool) { + cluster.KeepControl() + // AgentPool.Refresh asks for every topic. kgo's own loop asks + // for brokers only, and blocking that request stalls the cluster. + if req.(*kmsg.MetadataRequest).Topics == nil && blockNext.CompareAndSwap(true, false) { + cluster.SleepControl(func() { <-release }) + } + return nil, nil, false + }) - time.Sleep(10 * time.Second) - synctest.Wait() - // Queue a nudge while a Metadata request may be in flight. The count - // below is taken after that nudge has either run or is still waiting - // behind the blocked request. - c.triggerRefresh() - synctest.Wait() - before := onDemandRefreshes(c) + time.Sleep(10 * time.Second) + synctest.Wait() + // Queue a nudge while a Metadata request may be in flight. The count + // below is taken after that nudge has either run or is still waiting + // behind the blocked request. + c.triggerRefresh() + synctest.Wait() + before := onDemandRefreshes(c) - go c.Close() - synctest.Wait() - close(release) - synctest.Wait() + go c.Close() + synctest.Wait() + close(release) + synctest.Wait() - assert.Equal(t, before, onDemandRefreshes(c)) + assert.Equal(t, before, onDemandRefreshes(c), "iteration %d", i) + } }) }) } diff --git a/pkg/wgo/produce_hooks_test.go b/pkg/wgo/produce_hooks_test.go index 5a7e250..704ced6 100644 --- a/pkg/wgo/produce_hooks_test.go +++ b/pkg/wgo/produce_hooks_test.go @@ -161,16 +161,87 @@ func TestWarpstreamClient_ProduceRecordHooks(t *testing.T) { hook := &recordingHook{} c, _, _, _ := newTestWarpstreamClient(t, topic, 1, WithHooks(hook)) - // An unknown-topic record fails routing for the whole batch uniformly. + // The unknown topic is rejected on its own. The other record is produced. r1 := &kgo.Record{Topic: topic, Partition: 0, Value: []byte("a"), Timestamp: time.Now()} r2 := &kgo.Record{Topic: "does-not-exist", Partition: 0, Value: []byte("b"), Timestamp: time.Now()} results := c.ProduceSync(t.Context(), []*kgo.Record{r1, r2}) + require.NoError(t, results[0].Err) + require.ErrorContains(t, results[1].Err, "no agent assigned") + + events := hook.snapshot() + require.Equal(t, []string{"buffered", "buffered", "unbuffered", "unbuffered"}, kindsOf(events)) + assert.Same(t, r1, events[2].rec) + assert.Same(t, r2, events[3].rec) + assert.NoError(t, events[2].err) + assert.ErrorContains(t, events[3].err, "no agent assigned") + }) + }) + + t.Run("ProduceSync all rejected fires unbuffered once per record at return", func(t *testing.T) { + synctest.Test(t, func(t *testing.T) { + hook := &recordingHook{} + c, _, _, _ := newTestWarpstreamClient(t, topic, 1, WithHooks(hook)) + + r1 := &kgo.Record{Topic: "does-not-exist", Partition: 0, Value: []byte("a"), Timestamp: time.Now()} + r2 := &kgo.Record{Topic: "does-not-exist", Partition: 1, Value: []byte("b"), Timestamp: time.Now()} + results := c.ProduceSync(t.Context(), []*kgo.Record{r1, r2}) require.ErrorContains(t, results[0].Err, "no agent assigned") require.ErrorContains(t, results[1].Err, "no agent assigned") events := hook.snapshot() require.Equal(t, []string{"buffered", "buffered", "unbuffered", "unbuffered"}, kindsOf(events)) - assert.ErrorContains(t, events[2].err, "no agent assigned") + assert.Same(t, r1, events[2].rec) + assert.Same(t, r2, events[3].rec) + }) + }) + + t.Run("ProduceSync repeated record pointer reports each position's error", func(t *testing.T) { + synctest.Test(t, func(t *testing.T) { + hook := &recordingHook{} + c, _, _, _ := newTestWarpstreamClient(t, topic, 1, WithHooks(hook)) + + ctx, cancel := context.WithCancel(t.Context()) + cancel() + r := &kgo.Record{Topic: topic, Partition: 0, Value: []byte("v"), Timestamp: time.Now()} + miss := &kgo.Record{Topic: "does-not-exist", Partition: 0, Value: []byte("m"), Timestamp: time.Now()} + results := c.ProduceSync(ctx, []*kgo.Record{r, miss, r, miss}) + require.ErrorIs(t, results[0].Err, context.Canceled) + require.ErrorContains(t, results[1].Err, "no agent assigned") + require.ErrorIs(t, results[2].Err, context.Canceled) + require.ErrorContains(t, results[3].Err, "no agent assigned") + + events := hook.snapshot() + require.Equal(t, []string{ + "buffered", "buffered", "buffered", "buffered", + "unbuffered", "unbuffered", "unbuffered", "unbuffered", + }, kindsOf(events)) + assert.Same(t, r, events[4].rec) + assert.Same(t, miss, events[5].rec) + assert.Same(t, r, events[6].rec) + assert.Same(t, miss, events[7].rec) + assert.ErrorIs(t, events[4].err, context.Canceled) + assert.ErrorContains(t, events[5].err, "no agent assigned") + assert.ErrorIs(t, events[6].err, context.Canceled) + assert.ErrorContains(t, events[7].err, "no agent assigned") + }) + }) + + t.Run("ProduceSync pre-canceled mixed call keeps each record's error", func(t *testing.T) { + synctest.Test(t, func(t *testing.T) { + hook := &recordingHook{} + c, _, _, _ := newTestWarpstreamClient(t, topic, 1, WithHooks(hook)) + + ctx, cancel := context.WithCancel(t.Context()) + cancel() + routed := &kgo.Record{Topic: topic, Partition: 0, Value: []byte("a"), Timestamp: time.Now()} + miss := &kgo.Record{Topic: "does-not-exist", Partition: 0, Value: []byte("b"), Timestamp: time.Now()} + results := c.ProduceSync(ctx, []*kgo.Record{routed, miss}) + require.ErrorIs(t, results[0].Err, context.Canceled) + require.ErrorContains(t, results[1].Err, "no agent assigned") + + events := hook.snapshot() + require.Equal(t, []string{"buffered", "buffered", "unbuffered", "unbuffered"}, kindsOf(events)) + assert.ErrorIs(t, events[2].err, context.Canceled) assert.ErrorContains(t, events[3].err, "no agent assigned") }) }) diff --git a/pkg/wgo/route_records_test.go b/pkg/wgo/route_records_test.go new file mode 100644 index 0000000..245865f --- /dev/null +++ b/pkg/wgo/route_records_test.go @@ -0,0 +1,312 @@ +package wgo + +import ( + "testing" + "time" + + "github.com/prometheus/client_golang/prometheus" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "github.com/twmb/franz-go/pkg/kgo" +) + +func TestWarpstreamClient_RouteRecords(t *testing.T) { + const topic = "t" + + rec := func(topic string, partition int32, value string) *kgo.Record { + return &kgo.Record{Topic: topic, Partition: partition, Value: []byte(value)} + } + values := func(records []*kgo.Record) []string { + out := make([]string, len(records)) + for i, r := range records { + out[i] = string(r.Value) + } + return out + } + + t.Run("all routable preserves destination, probe state, and grouping", func(t *testing.T) { + strategy := &mockPartitionAssignmentStrategy{candidates: map[partitionKey][]Agent{ + {topic, 0}: {{NodeID: 7, State: AgentStateDemoted}}, + {topic, 1}: healthyAgents(8), + }} + nudge := make(chan struct{}, 4) + c := newRouteRecordsClient(strategy, nudge) + var accepted int + records := []*kgo.Record{rec(topic, 1, "a"), rec(topic, 0, "b"), rec(topic, 1, "c")} + + routed, rejected := c.routeRecords(records, countAccepted(&accepted)) + + require.Nil(t, rejected) + require.Len(t, routed, 2) + assert.Equal(t, int32(8), routed[0].item.nodeID) + assert.Equal(t, AgentStateHealthy, routed[0].item.nodeState) + assert.Equal(t, []string{"a", "c"}, values(routed[0].item.records)) + assert.Equal(t, int32(7), routed[1].item.nodeID) + assert.Equal(t, AgentStateDemoted, routed[1].item.nodeState) + assert.Equal(t, []string{"b"}, values(routed[1].item.records)) + assert.Equal(t, 2, accepted) + assert.Equal(t, 1, strategy.candidatesCalls(topic, 0)) + assert.Equal(t, 1, strategy.candidatesCalls(topic, 1)) + assert.Empty(t, nudge) + }) + + t.Run("interleaved topics keep input order and look up each partition once", func(t *testing.T) { + strategy := &mockPartitionAssignmentStrategy{candidates: map[partitionKey][]Agent{ + {"a", 0}: healthyAgents(1), + {"b", 1}: healthyAgents(2), + }} + c := newRouteRecordsClient(strategy, make(chan struct{}, 4)) + var accepted int + records := []*kgo.Record{ + rec("a", 0, "a1"), + rec("b", 0, "b1"), + rec("a", 0, "a2"), + rec("c", 0, "c1"), + rec("b", 0, "b2"), + rec("b", 1, "b3"), + } + + routed, rejected := c.routeRecords(records, countAccepted(&accepted)) + + require.Len(t, routed, 2) + assert.Equal(t, []string{"a1", "a2"}, values(routed[0].item.records)) + assert.Equal(t, "a", routed[0].item.topic) + assert.Equal(t, []string{"b3"}, values(routed[1].item.records)) + require.Len(t, rejected, 2) + assert.Equal(t, []string{"b1", "b2"}, values(rejected[0].records)) + assert.Equal(t, "b", rejected[0].topic) + assert.Equal(t, int32(0), rejected[0].partition) + assert.ErrorContains(t, rejected[0].err, `no agent assigned for topic "b" partition 0`) + assert.Equal(t, []string{"c1"}, values(rejected[1].records)) + assert.Equal(t, 2, accepted) + assert.Equal(t, 1, strategy.candidatesCalls("a", 0)) + assert.Equal(t, 1, strategy.candidatesCalls("b", 0)) + assert.Equal(t, 1, strategy.candidatesCalls("b", 1)) + assert.Equal(t, 1, strategy.candidatesCalls("c", 0)) + }) + + t.Run("all unroutable builds one group and no accepted callback", func(t *testing.T) { + nudge := make(chan struct{}, 4) + c := newRouteRecordsClient(&mockPartitionAssignmentStrategy{}, nudge) + var accepted int + records := []*kgo.Record{ + rec(topic, 3, "a"), rec(topic, 3, "b"), rec(topic, 3, "c"), + } + + routed, rejected := c.routeRecords(records, countAccepted(&accepted)) + + assert.Empty(t, routed) + require.Len(t, rejected, 1) + assert.Equal(t, []string{"a", "b", "c"}, values(rejected[0].records)) + assert.Equal(t, 0, accepted) + assert.Len(t, nudge, 1) + }) + + for _, place := range []struct { + name string + at int + }{{"front", 0}, {"middle", 1}, {"back", 2}} { + t.Run("one miss at the "+place.name, func(t *testing.T) { + strategy := &mockPartitionAssignmentStrategy{candidates: map[partitionKey][]Agent{ + {topic, 0}: healthyAgents(1), + {topic, 1}: healthyAgents(1), + {topic, 2}: healthyAgents(1), + }} + delete(strategy.candidates, partitionKey{topic, int32(place.at)}) + c := newRouteRecordsClient(strategy, make(chan struct{}, 4)) + records := []*kgo.Record{rec(topic, 0, "0"), rec(topic, 1, "1"), rec(topic, 2, "2")} + + routed, rejected := c.routeRecords(records, countAccepted(new(int))) + + require.Len(t, rejected, 1) + assert.Equal(t, int32(place.at), rejected[0].partition) + assert.Len(t, routed, 2) + for _, g := range routed { + assert.NotEqual(t, int32(place.at), g.item.partition) + } + }) + } + + t.Run("a repeated rejected partition is not looked up again", func(t *testing.T) { + strategy := &seqStrategy{answer: func(partition int32, call int) []Agent { + if partition == 1 && call == 1 { + return nil + } + if partition == 1 { + return healthyAgents(9) + } + if call == 1 { + return []Agent{{NodeID: 1, State: AgentStateDemoted}} + } + return healthyAgents(2) + }} + nudge := make(chan struct{}, 4) + c := newRouteRecordsClient(strategy, nudge) + records := []*kgo.Record{ + rec(topic, 0, "a"), + rec(topic, 1, "b"), + rec(topic, 0, "c"), + rec(topic, 1, "d"), + rec(topic, 2, "e"), + } + + routed, rejected := c.routeRecords(records, countAccepted(new(int))) + + require.Len(t, routed, 2) + assert.Equal(t, int32(1), routed[0].item.nodeID) + assert.Equal(t, AgentStateDemoted, routed[0].item.nodeState) + assert.Equal(t, []string{"a", "c"}, values(routed[0].item.records)) + assert.Equal(t, int32(2), routed[1].item.partition) + require.Len(t, rejected, 1) + assert.Equal(t, []string{"b", "d"}, values(rejected[0].records)) + assert.Equal(t, 1, strategy.calls[topicPartition{topic, 0}]) + assert.Equal(t, 1, strategy.calls[topicPartition{topic, 1}]) + assert.Equal(t, 1, strategy.calls[topicPartition{topic, 2}]) + assert.Len(t, nudge, 1) + }) + + t.Run("several misses nudge once", func(t *testing.T) { + nudge := make(chan struct{}, 8) + c := newRouteRecordsClient(&mockPartitionAssignmentStrategy{candidates: map[partitionKey][]Agent{ + {topic, 0}: healthyAgents(1), + }}, nudge) + records := []*kgo.Record{rec(topic, 1, "a"), rec(topic, 2, "b"), rec(topic, 1, "c")} + + _, rejected := c.routeRecords(records, countAccepted(new(int))) + + require.Len(t, rejected, 2) + assert.Len(t, nudge, 1) + }) + + t.Run("a full nudge channel does not block", func(t *testing.T) { + nudge := make(chan struct{}, 1) + nudge <- struct{}{} + c := newRouteRecordsClient(&mockPartitionAssignmentStrategy{}, nudge) + done := make(chan struct{}) + go func() { + c.routeRecords([]*kgo.Record{rec(topic, 0, "a"), rec(topic, 1, "b")}, countAccepted(new(int))) + close(done) + }() + select { + case <-done: + case <-time.After(2 * time.Second): + t.Fatal("routeRecords blocked on a full nudge channel") + } + assert.Len(t, nudge, 1) + }) + + t.Run("no leader, unknown topic, and empty pool are rejected", func(t *testing.T) { + agents := []int32{3} + strategy := newDefaultPartitionAssignmentStrategy(agents, map[topicPartition]int32{ + {topic: topic, partition: 0}: 3, + }, nil, map[topicPartition]struct{}{ + {topic: topic, partition: 1}: {}, + }) + c := newRouteRecordsClient(strategy, make(chan struct{}, 4)) + records := []*kgo.Record{ + rec(topic, 0, "leader"), + rec(topic, 1, "unnamed"), + rec("never-seen", 0, "unknown"), + rec(topic, 9, "named-missing"), + } + + routed, rejected := c.routeRecords(records, countAccepted(new(int))) + + require.Len(t, routed, 2) + assert.Equal(t, int32(0), routed[0].item.partition) + assert.Equal(t, int32(9), routed[1].item.partition) + assert.Equal(t, agents[0], routed[1].item.nodeID) + require.Len(t, rejected, 2) + assert.Equal(t, int32(1), rejected[0].partition) + assert.Equal(t, "never-seen", rejected[1].topic) + + empty := newRouteRecordsClient(newDefaultPartitionAssignmentStrategy(nil, nil, nil, nil), make(chan struct{}, 1)) + routed, rejected = empty.routeRecords([]*kgo.Record{rec(topic, 0, "a")}, countAccepted(new(int))) + assert.Empty(t, routed) + require.Len(t, rejected, 1) + }) +} + +func newRouteRecordsClient(strategy PartitionAssignmentStrategy, nudge chan struct{}) *WarpstreamClient { + return &WarpstreamClient{ + demoter: NewDemoter(strategy, noopAgentStatsTracker{}, HealthCheckConfig{}, DemoterConfig{}, nopLogger{}, prometheus.NewRegistry()), + refreshNowCh: nudge, + } +} + +func countAccepted(n *int) func([]*kgo.Record) func(ProduceResult) { + return func([]*kgo.Record) func(ProduceResult) { + *n++ + return func(ProduceResult) {} + } +} + +// seqStrategy changes its answer after the first call, so a repeat lookup shows up. +type seqStrategy struct { + answer func(partition int32, call int) []Agent + calls map[topicPartition]int +} + +func (s *seqStrategy) Candidates(topic string, partition int32, _ int) []Agent { + if s.calls == nil { + s.calls = map[topicPartition]int{} + } + key := topicPartition{topic: topic, partition: partition} + s.calls[key]++ + return s.answer(partition, s.calls[key]) +} + +func BenchmarkWarpstreamClient_RouteRecords(b *testing.B) { + const topic = "t" + cases := []struct { + name string + records int + partitions int + missEvery int + }{ + {"all routable 32x1", 32, 1, 0}, + {"all routable 32x32", 32, 32, 0}, + {"all routable 256x32", 256, 32, 0}, + {"miss at front 256", 256, 32, -1}, + {"several misses 256", 256, 32, 2}, + {"all unroutable 256", 256, 32, 1}, + {"many records one rejected partition", 256, 1, 1}, + {"interleaved 256", 256, 16, 2}, + } + for _, tc := range cases { + b.Run(tc.name, func(b *testing.B) { + records, strategy := benchRouteInputs(topic, tc.records, tc.partitions, tc.missEvery) + c := newRouteRecordsClient(strategy, make(chan struct{}, 1)) + doneFor := func([]*kgo.Record) func(ProduceResult) { return func(ProduceResult) {} } + b.ReportAllocs() + b.ResetTimer() + for range b.N { + benchRouted, benchRejected = c.routeRecords(records, doneFor) + } + }) + } +} + +var ( + benchRouted []promised[routedTopicPartitionRecords] + benchRejected []rejectedTopicPartitionRecords +) + +// missEvery: 0 routes all, 1 rejects all, -1 rejects the first record. +// Any other positive value rejects partitions divisible by it. +func benchRouteInputs(topic string, n, partitions, missEvery int) ([]*kgo.Record, *mockPartitionAssignmentStrategy) { + records := make([]*kgo.Record, n) + candidates := map[partitionKey][]Agent{} + for i := range n { + p := int32(i % partitions) + if missEvery < 0 && i == 0 { + p = int32(partitions + 1) + } + records[i] = &kgo.Record{Topic: topic, Partition: p, Value: []byte("v")} + reject := missEvery == 1 || (missEvery > 1 && int(p)%missEvery == 0) || (missEvery < 0 && i == 0) + if !reject { + candidates[partitionKey{topic, p}] = healthyAgents(1) + } + } + return records, &mockPartitionAssignmentStrategy{candidates: candidates} +}