mirror of https://github.com/milvus-io/milvus.git
90 lines
2.2 KiB
Go
90 lines
2.2 KiB
Go
|
package accesslog
|
||
|
|
||
|
import (
|
||
|
"context"
|
||
|
"fmt"
|
||
|
"testing"
|
||
|
|
||
|
"google.golang.org/grpc"
|
||
|
"google.golang.org/grpc/metadata"
|
||
|
|
||
|
"github.com/milvus-io/milvus-proto/go-api/v2/commonpb"
|
||
|
"github.com/milvus-io/milvus-proto/go-api/v2/milvuspb"
|
||
|
"github.com/milvus-io/milvus/internal/proxy/connection"
|
||
|
"github.com/milvus-io/milvus/pkg/util"
|
||
|
"github.com/milvus-io/milvus/pkg/util/paramtable"
|
||
|
)
|
||
|
|
||
|
type TestData struct {
|
||
|
req, resp interface{}
|
||
|
err error
|
||
|
}
|
||
|
|
||
|
func genTestData(clientInfo *commonpb.ClientInfo, identifier int64) []*TestData {
|
||
|
ret := []*TestData{}
|
||
|
|
||
|
ret = append(ret, &TestData{
|
||
|
req: &milvuspb.QueryRequest{
|
||
|
CollectionName: "test1",
|
||
|
Expr: "pk >= 100",
|
||
|
},
|
||
|
resp: &milvuspb.QueryResults{
|
||
|
CollectionName: "test1",
|
||
|
},
|
||
|
err: nil,
|
||
|
})
|
||
|
|
||
|
ret = append(ret, &TestData{
|
||
|
req: &milvuspb.SearchRequest{
|
||
|
CollectionName: "test2",
|
||
|
Dsl: "pk <= 100",
|
||
|
},
|
||
|
resp: &milvuspb.SearchResults{
|
||
|
CollectionName: "test2",
|
||
|
},
|
||
|
err: nil,
|
||
|
})
|
||
|
|
||
|
ret = append(ret, &TestData{
|
||
|
req: &milvuspb.ConnectRequest{
|
||
|
ClientInfo: clientInfo,
|
||
|
},
|
||
|
resp: &milvuspb.ConnectResponse{
|
||
|
Identifier: identifier,
|
||
|
},
|
||
|
err: nil,
|
||
|
})
|
||
|
|
||
|
return ret
|
||
|
}
|
||
|
|
||
|
func BenchmarkAccesslog(b *testing.B) {
|
||
|
paramtable.Init()
|
||
|
Params := paramtable.Get()
|
||
|
Params.Save(Params.ProxyCfg.AccessLog.Enable.Key, "true")
|
||
|
Params.Save(Params.ProxyCfg.AccessLog.Filename.Key, "")
|
||
|
Params.Save(Params.CommonCfg.ClusterPrefix.Key, "in-test")
|
||
|
initAccessLogger(&Params.ProxyCfg.AccessLog, &Params.MinioCfg)
|
||
|
paramtable.Get().CommonCfg.ClusterPrefix.GetValue()
|
||
|
|
||
|
clientInfo := &commonpb.ClientInfo{
|
||
|
SdkType: "gotest",
|
||
|
SdkVersion: "testversion",
|
||
|
}
|
||
|
identifier := int64(11111)
|
||
|
md := metadata.MD{util.IdentifierKey: []string{fmt.Sprint(identifier)}}
|
||
|
ctx := metadata.NewIncomingContext(context.TODO(), md)
|
||
|
connection.GetManager().Register(ctx, identifier, clientInfo)
|
||
|
rpcInfo := &grpc.UnaryServerInfo{Server: nil, FullMethod: "testMethod"}
|
||
|
datas := genTestData(clientInfo, identifier)
|
||
|
|
||
|
b.ResetTimer()
|
||
|
for i := 0; i < b.N; i++ {
|
||
|
data := datas[i%len(datas)]
|
||
|
accessInfo := NewGrpcAccessInfo(ctx, rpcInfo, data.req)
|
||
|
accessInfo.UpdateCtx(ctx)
|
||
|
accessInfo.SetResult(data.resp, data.err)
|
||
|
accessInfo.Write()
|
||
|
}
|
||
|
}
|