diff --git a/src/internal/trace/v2/batch.go b/src/internal/trace/v2/batch.go
index 899eb0f59b..d7afc06eec 100644
--- a/src/internal/trace/v2/batch.go
+++ b/src/internal/trace/v2/batch.go
@@ -5,7 +5,6 @@
 package trace
 
 import (
-	"bufio"
 	"bytes"
 	"encoding/binary"
 	"fmt"
@@ -43,7 +42,10 @@ func (b *batch) isFreqBatch() bool {
 }
 
 // readBatch reads the next full batch from r.
-func readBatch(r *bufio.Reader) (batch, uint64, error) {
+func readBatch(r interface {
+	io.Reader
+	io.ByteReader
+}) (batch, uint64, error) {
 	// Read batch header byte.
 	b, err := r.ReadByte()
 	if err != nil {