kgo: add Buffered{Fetch,Produce}Bytes

twmb · twmb · commit c80d6f46e35b · 2023-09-20T16:53:29.000-06:00
This slows down fetching a little bit. If it is egregious, we can fix
the perf by tracking the size buffered when processing the fetch itself,
and then adding a new field to batch-untrack the size. That's left as an
exercise for a person that cares.

This is now done since the prior commit introduces buffered produce
bytes, and we may as well add it while fetching for both-sizes
consistency, and we may as well expose it.
diff --git a/pkg/kgo/consumer.go b/pkg/kgo/consumer.go
@@ -162,6 +162,7 @@ func (o Offset) At(at int64) Offset {
 
 type consumer struct {
 	bufferedRecords atomicI64
+	bufferedBytes   atomicI64
 
 	cl *Client
 
@@ -285,6 +286,13 @@ func (cl *Client) BufferedFetchRecords() int64 {
 	return cl.consumer.bufferedRecords.Load()
 }
 
+// BufferedFetchBytes returns the number of bytes currently buffered from
+// fetching within the client. This is the sum of all keys, values, and header
+// keys/values. See the related [BufferedFetchRecords] for more information.
+func (cl *Client) BufferedFetchBytes() int64 {
+	return cl.consumer.bufferedBytes.Load()
+}
+
 type usedCursors map[*cursor]struct{}
 
 func (u *usedCursors) use(c *cursor) {
diff --git a/pkg/kgo/producer.go b/pkg/kgo/producer.go
@@ -88,6 +88,13 @@ func (cl *Client) BufferedProduceRecords() int64 {
 	return cl.producer.bufferedRecords.Load()
 }
 
+// BufferedProduceBytes returns the number of bytes currently buffered for
+// producing within the client. This is the sum of all keys, values, and header
+// keys/values. See the related [BufferedProduceRecords] for more information.
+func (cl *Client) BufferedProduceBytes() int64 {
+	return cl.producer.bufferedBytes.Load()
+}
+
 type unknownTopicProduces struct {
 	buffered []promisedRec
 	wait     chan error // retryable errors
diff --git a/pkg/kgo/source.go b/pkg/kgo/source.go
@@ -323,16 +323,23 @@ func (s *source) hook(f *Fetch, buffered, polled bool) {
 	})
 
 	var nrecs int
+	var nbytes int64
 	for i := range f.Topics {
 		t := &f.Topics[i]
 		for j := range t.Partitions {
-			nrecs += len(t.Partitions[j].Records)
+			p := &t.Partitions[j]
+			nrecs += len(p.Records)
+			for k := range p.Records {
+				nbytes += p.Records[k].userSize()
+			}
 		}
 	}
 	if buffered {
 		s.cl.consumer.bufferedRecords.Add(int64(nrecs))
+		s.cl.consumer.bufferedBytes.Add(nbytes)
 	} else {
 		s.cl.consumer.bufferedRecords.Add(-int64(nrecs))
+		s.cl.consumer.bufferedBytes.Add(-nbytes)
 	}
 }
 

Original file line number	Diff line number	Diff line change
`@@ -323,16 +323,23 @@ func (s source) hook(f Fetch, buffered, polled bool) {`
`323`	`323`	`})`
`324`	`324`
`325`	`325`	`var nrecs int`
	`326`	`+ var nbytes int64`
`326`	`327`	`for i := range f.Topics {`
`327`	`328`	`t := &f.Topics[i]`
`328`	`329`	`for j := range t.Partitions {`
`329`		`- nrecs += len(t.Partitions[j].Records)`
	`330`	`+ p := &t.Partitions[j]`
	`331`	`+ nrecs += len(p.Records)`
	`332`	`+ for k := range p.Records {`
	`333`	`+ nbytes += p.Records[k].userSize()`
	`334`	`+ }`
`330`	`335`	`}`
`331`	`336`	`}`
`332`	`337`	`if buffered {`
`333`	`338`	`s.cl.consumer.bufferedRecords.Add(int64(nrecs))`
	`339`	`+ s.cl.consumer.bufferedBytes.Add(nbytes)`
`334`	`340`	`} else {`
`335`	`341`	`s.cl.consumer.bufferedRecords.Add(-int64(nrecs))`
	`342`	`+ s.cl.consumer.bufferedBytes.Add(-nbytes)`
`336`	`343`	`}`
`337`	`344`	`}`
`338`	`345`