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>
94 lines
2.7 KiB
Go
94 lines
2.7 KiB
Go
package lazygrpc
|
|
|
|
import (
|
|
"context"
|
|
|
|
"github.com/cenkalti/backoff/v4"
|
|
"github.com/cockroachdb/errors"
|
|
"go.uber.org/zap"
|
|
"google.golang.org/grpc"
|
|
|
|
"github.com/milvus-io/milvus/pkg/log"
|
|
"github.com/milvus-io/milvus/pkg/util/syncutil"
|
|
)
|
|
|
|
var ErrClosed = errors.New("lazy grpc conn closed")
|
|
|
|
// NewConn creates a new lazy grpc conn.
|
|
func NewConn(dialer func(ctx context.Context) (*grpc.ClientConn, error)) Conn {
|
|
conn := &connImpl{
|
|
initializationNotifier: syncutil.NewAsyncTaskNotifier[struct{}](),
|
|
conn: syncutil.NewFuture[*grpc.ClientConn](),
|
|
dialer: dialer,
|
|
}
|
|
go conn.initialize()
|
|
return conn
|
|
}
|
|
|
|
// Conn is a lazy grpc conn implementation.
|
|
// grpc.Dial operation will block until new grpc conn is created at least once.
|
|
// Conn will dial the underlying grpc conn asynchronously to avoid dependency cycle of milvus component when create grpc client.
|
|
// TODO: Remove in future if we can refactor the dependency cycle.
|
|
type Conn interface {
|
|
// GetConn will block until the grpc.ClientConn is ready to use.
|
|
// If the context is done, return immediately with the context.Canceled or Context.DeadlineExceeded error.
|
|
// Return ErrClosed if the lazy grpc conn is closed.
|
|
GetConn(ctx context.Context) (*grpc.ClientConn, error)
|
|
|
|
// Close closes the lazy grpc conn.
|
|
// Close the underlying grpc conn if it is already created.
|
|
Close()
|
|
}
|
|
|
|
type connImpl struct {
|
|
initializationNotifier *syncutil.AsyncTaskNotifier[struct{}]
|
|
conn *syncutil.Future[*grpc.ClientConn]
|
|
|
|
dialer func(ctx context.Context) (*grpc.ClientConn, error)
|
|
}
|
|
|
|
func (c *connImpl) initialize() {
|
|
defer c.initializationNotifier.Finish(struct{}{})
|
|
|
|
backoff.Retry(func() error {
|
|
conn, err := c.dialer(c.initializationNotifier.Context())
|
|
if err != nil {
|
|
if c.initializationNotifier.Context().Err() != nil {
|
|
log.Info("lazy grpc conn canceled", zap.Error(c.initializationNotifier.Context().Err()))
|
|
return nil
|
|
}
|
|
log.Warn("async dial failed, wait for retry...", zap.Error(err))
|
|
return err
|
|
}
|
|
c.conn.Set(conn)
|
|
return nil
|
|
}, backoff.NewExponentialBackOff())
|
|
}
|
|
|
|
func (c *connImpl) GetConn(ctx context.Context) (*grpc.ClientConn, error) {
|
|
// If the context is done, return immediately to perform a stable shutdown error after closing.
|
|
if c.initializationNotifier.Context().Err() != nil {
|
|
return nil, ErrClosed
|
|
}
|
|
|
|
select {
|
|
case <-ctx.Done():
|
|
return nil, ctx.Err()
|
|
case <-c.initializationNotifier.Context().Done():
|
|
return nil, ErrClosed
|
|
case <-c.conn.Done():
|
|
return c.conn.Get(), nil
|
|
}
|
|
}
|
|
|
|
func (c *connImpl) Close() {
|
|
c.initializationNotifier.Cancel()
|
|
c.initializationNotifier.BlockUntilFinish()
|
|
|
|
if c.conn.Ready() {
|
|
if err := c.conn.Get().Close(); err != nil {
|
|
log.Warn("close underlying grpc conn fail", zap.Error(err))
|
|
}
|
|
}
|
|
}
|