mirror of
https://gitee.com/milvus-io/milvus.git
synced 2025-12-06 17:18:35 +08:00
fda720b880
issue: #33285 - add two grpc resolver (by session and by streaming coord assignment service) - add one grpc balancer (by serverID and roundrobin) - add lazy conn to avoid block by first service discovery - add some utility function for streaming service Signed-off-by: chyezh <chyezh@outlook.com>
36 lines
1.4 KiB
Go
36 lines
1.4 KiB
Go
package interceptor
|
|
|
|
import (
|
|
"context"
|
|
"strings"
|
|
|
|
"google.golang.org/grpc"
|
|
|
|
"github.com/milvus-io/milvus/internal/proto/streamingpb"
|
|
"github.com/milvus-io/milvus/internal/util/streamingutil/status"
|
|
)
|
|
|
|
// NewStreamingServiceUnaryClientInterceptor returns a new unary client interceptor for error handling.
|
|
func NewStreamingServiceUnaryClientInterceptor() grpc.UnaryClientInterceptor {
|
|
return func(ctx context.Context, method string, req, reply interface{}, cc *grpc.ClientConn, invoker grpc.UnaryInvoker, opts ...grpc.CallOption) error {
|
|
err := invoker(ctx, method, req, reply, cc, opts...)
|
|
if strings.HasPrefix(method, streamingpb.ServiceMethodPrefix) {
|
|
st := status.ConvertStreamingError(method, err)
|
|
return st
|
|
}
|
|
return err
|
|
}
|
|
}
|
|
|
|
// NewStreamingServiceStreamClientInterceptor returns a new stream client interceptor for error handling.
|
|
func NewStreamingServiceStreamClientInterceptor() grpc.StreamClientInterceptor {
|
|
return func(ctx context.Context, desc *grpc.StreamDesc, cc *grpc.ClientConn, method string, streamer grpc.Streamer, opts ...grpc.CallOption) (grpc.ClientStream, error) {
|
|
clientStream, err := streamer(ctx, desc, cc, method, opts...)
|
|
if strings.HasPrefix(method, streamingpb.ServiceMethodPrefix) {
|
|
e := status.ConvertStreamingError(method, err)
|
|
return status.NewClientStreamWrapper(method, clientStream), e
|
|
}
|
|
return clientStream, err
|
|
}
|
|
}
|