Files
wehub-resource-sync 498b235461
Build and test / Build and test AMD64 Ubuntu 22.04 (push) Failing after 0s
Publish Builder / amazonlinux2023 (push) Failing after 1s
Build and test / UT for Go (push) Has been skipped
Publish KRTE Images / KRTE (push) Failing after 1s
Build and test / Integration Test (push) Has been skipped
Build and test / Upload Code Coverage (push) Has been skipped
Publish Builder / rockylinux9 (push) Failing after 1s
Publish Builder / ubuntu22.04 (push) Failing after 0s
Publish Builder / ubuntu24.04 (push) Failing after 0s
Publish Gpu Builder / publish-gpu-builder (push) Failing after 1s
Publish Test Images / PyTest (push) Failing after 0s
Build and test / UT for Cpp (push) Has been cancelled
chore: import upstream snapshot with attribution
2026-07-13 12:31:17 +08:00

73 lines
1.7 KiB
Go

package msgstream
import (
"context"
"fmt"
"math/rand"
"sync"
"testing"
"github.com/stretchr/testify/assert"
"github.com/milvus-io/milvus/pkg/v3/mq/common"
"github.com/milvus-io/milvus/pkg/v3/mq/msgstream/mqwrapper"
"github.com/milvus-io/milvus/pkg/v3/util/funcutil"
)
func benchmarkProduceAndConsume(b *testing.B, mqClient mqwrapper.Client, cases [][]byte) {
topic := fmt.Sprintf("test_produce_and_consume_topic_%d", rand.Int31n(100000))
wg := sync.WaitGroup{}
wg.Add(2)
go func() {
defer wg.Done()
p, err := mqClient.CreateProducer(context.TODO(), common.ProducerOptions{
Topic: topic,
})
assert.NoError(b, err)
defer p.Close()
benchmarkMQProduce(b, p, cases)
}()
go func() {
defer wg.Done()
c, _ := mqClient.Subscribe(context.TODO(), mqwrapper.ConsumerOptions{
Topic: topic,
SubscriptionName: topic,
SubscriptionInitialPosition: common.SubscriptionPositionEarliest,
BufSize: 1024,
})
defer c.Close()
benchmarkMQConsume(b, c, cases)
}()
wg.Wait()
}
func benchmarkMQConsume(b *testing.B, c mqwrapper.Consumer, cases [][]byte) {
ch := c.Chan()
for range cases {
msg := <-ch
c.Ack(msg)
}
c.Close()
}
func benchmarkMQProduce(b *testing.B, p mqwrapper.Producer, cases [][]byte) {
for _, c := range cases {
p.Send(context.Background(), &common.ProducerMessage{
Payload: c,
})
}
}
func generateRandBytes(m int, n int) [][]byte {
letterRunes := funcutil.RandomBytes(2 * m)
cases := make([][]byte, 0, n)
for i := 0; i < n; i++ {
startOffset := rand.Intn(m)
endOffset := startOffset + m
cases = append(cases, letterRunes[startOffset:endOffset])
}
return cases
}