Merge pull request #190 from ydb-platform/add-address-to-error-grom-take

* Added address to error description which thrown from `conn.take()`

Author: asmyasnikov

CHANGELOG.md

@@ -1,4 +1,7 @@
-* Renamed package `internal/db` to `internal/database` for exclude collisions with variable name `db`
+* Added `conn.Ping(ctx)` method for check availability of `conn.Conn`
+* Refactored `cluster.Cluster.Get(ctx)` to return only available connection (instead of returning any connection from balancer)
+* Added address to error description thrown from `conn.take()`
+* Renamed package `internal/db` to `internal/database` to exclude collisions with variable name `db`
 
 ## v3.18.0
 * Added `go1.18` to test matrix

internal/cluster/cluster.go

@@ -129,7 +129,7 @@ func parseOptions(opts ...crudOption) *crudOptionsHolder {
 
 type Getter interface {
 	// Get gets conn from cluster
-	Get(ctx context.Context, opts ...crudOption) (cc conn.Conn, err error)
+	Get(ctx context.Context) (cc conn.Conn, err error)
 }
 
 type Inserter interface {
@@ -247,17 +247,11 @@ func (c *cluster) Close(ctx context.Context) (err error) {
 
 // Get returns next available connection.
 // It returns error on given deadline cancellation or when cluster become closed.
-func (c *cluster) Get(ctx context.Context, opts ...crudOption) (cc conn.Conn, err error) {
+func (c *cluster) Get(ctx context.Context) (cc conn.Conn, err error) {
 	var cancel context.CancelFunc
 	ctx, cancel = context.WithTimeout(ctx, MaxGetConnTimeout)
 	defer cancel()
 
-	options := parseOptions(opts...)
-	if options.withLock {
-		c.mu.Lock()
-		defer c.mu.Unlock()
-	}
-
 	if c.closed {
 		return nil, errors.WithStackTrace(ErrClusterClosed)
 	}
@@ -272,7 +266,9 @@ func (c *cluster) Get(ctx context.Context, opts ...crudOption) (cc conn.Conn, er
 	}()
 
 	if e, ok := ContextEndpoint(ctx); ok {
+		c.mu.RLock()
 		cc, ok = c.endpoints[e.NodeID()]
+		c.mu.RUnlock()
 		if ok && cc.IsState(
 			conn.Created,
 			conn.Online,
@@ -282,15 +278,22 @@ func (c *cluster) Get(ctx context.Context, opts ...crudOption) (cc conn.Conn, er
 		}
 	}
 
-	c.balancerMtx.RLock()
-	defer c.balancerMtx.RUnlock()
-
-	cc = c.config.Balancer().Next()
-	if cc == nil {
-		return nil, errors.WithStackTrace(ErrClusterEmpty)
+	for {
+		select {
+		case <-ctx.Done():
+			return nil, errors.WithStackTrace(ctx.Err())
+		default:
+			c.balancerMtx.RLock()
+			cc = c.config.Balancer().Next()
+			c.balancerMtx.RUnlock()
+			if cc == nil {
+				return nil, errors.WithStackTrace(ErrClusterEmpty)
+			}
+			if err = cc.Ping(ctx); err == nil {
+				return cc, nil
+			}
+		}
 	}
-
-	return cc, nil
 }
 
 // Insert inserts new connection into the cluster.

internal/cluster/cluster_test.go

@@ -5,7 +5,6 @@ import (
 	"fmt"
 	"sort"
 	"testing"
-	"time"
 
 	"github.com/ydb-platform/ydb-go-sdk/v3/balancers"
 	"github.com/ydb-platform/ydb-go-sdk/v3/config"
@@ -17,53 +16,6 @@ import (
 	"github.com/ydb-platform/ydb-go-sdk/v3/trace"
 )
 
-func TestClusterFastRedial(t *testing.T) {
-	ctx, cancel := context.WithCancel(context.Background())
-	defer cancel()
-
-	l, b := stub.Balancer()
-	c := &cluster{
-		config: config.New(
-			config.WithBalancer(b),
-		),
-		index:     make(map[string]entry.Entry),
-		endpoints: make(map[uint32]conn.Conn),
-		pool:      conn.NewPool(ctx, config.New()),
-	}
-
-	pingConnects := func(size int) chan struct{} {
-		done := make(chan struct{})
-		go func() {
-			for i := 0; i < size*10; i++ {
-				cc, err := c.Get(context.Background())
-				// enforce close bad connects to track them
-				if err == nil && c != nil && cc.Endpoint().Address() == "bad:0" {
-					_ = c.Remove(ctx, cc.Endpoint())
-				}
-			}
-			close(done)
-		}()
-		return done
-	}
-
-	ne := []endpoint.Endpoint{
-		endpoint.New("foo:0"),
-		endpoint.New("bad:0"),
-	}
-	mergeEndpointIntoCluster(
-		ctx,
-		c,
-		[]endpoint.Endpoint{},
-		ne,
-	)
-	select {
-	case <-pingConnects(len(ne)):
-
-	case <-time.After(time.Second * 15):
-		t.Fatalf("Time limit exceeded while %d endpoints in balance. Wait channel used", len(*l))
-	}
-}
-
 func TestClusterMergeEndpoints(t *testing.T) {
 	ctx, cancel := context.WithCancel(context.Background())
 	defer cancel()

internal/cmd/gtrace/main.go

@@ -67,11 +67,15 @@ func main() {
 		// We should respect Go suffixes like `_linux.go`.
 		name, tags, ext := splitOSArchTags(&buildCtx, gofile)
 		openFile := func(name string) (*os.File, func()) {
-			p := filepath.Join(workDir, name)
 			var f *os.File
 			// nolint: gofumpt
 			// nolint: nolintlint
-			f, err = os.OpenFile(p, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0666)
+			// nolint: gosec
+			f, err = os.OpenFile(
+				filepath.Join(workDir, filepath.Clean(name)),
+				os.O_WRONLY|os.O_CREATE|os.O_TRUNC,
+				0600,
+			)
 			if err != nil {
 				log.Fatal(err)
 			}

internal/conn/conn.go

@@ -21,8 +21,12 @@ import (
 var (
 	// errOperationNotReady specified error when operation is not ready
 	errOperationNotReady = errors.New(fmt.Errorf("operation is not ready yet"))
+
 	// errClosedConnection specified error when connection are closed early
 	errClosedConnection = errors.New(fmt.Errorf("connection closed early"))
+
+	// errUnavailableConnection specified error when connection are closed early
+	errUnavailableConnection = errors.New(fmt.Errorf("connection unavailable"))
 )
 
 type Conn interface {
@@ -32,6 +36,7 @@ type Conn interface {
 
 	LastUsage() time.Time
 
+	Ping(ctx context.Context) error
 	IsState(states ...State) bool
 	GetState() State
 	SetState(State) State
@@ -57,6 +62,17 @@ type conn struct {
 	onClose      []func(*conn)
 }
 
+func (c *conn) Ping(ctx context.Context) error {
+	cc, err := c.take(ctx)
+	if err != nil {
+		return errors.WithStackTrace(err)
+	}
+	if !isAvailable(cc) {
+		return errors.WithStackTrace(errUnavailableConnection)
+	}
+	return nil
+}
+
 func (c *conn) Release(ctx context.Context) (err error) {
 	var (
 		onDone = trace.DriverOnConnRelease(
@@ -212,7 +228,7 @@ func (c *conn) take(ctx context.Context) (cc *grpc.ClientConn, err error) {
 
 	cc, err = grpc.DialContext(ctx, address, c.config.GrpcDialOptions()...)
 	if err != nil {
-		return nil, errors.WithStackTrace(err)
+		return nil, errors.WithStackTrace(fmt.Errorf("dial %s failed: %w", address, err))
 	}
 
 	c.cc = cc
@@ -269,6 +285,10 @@ func isBroken(raw *grpc.ClientConn) bool {
 	return s == connectivity.Shutdown || s == connectivity.TransientFailure
 }
 
+func isAvailable(raw *grpc.ClientConn) bool {
+	return raw != nil && raw.GetState() == connectivity.Ready
+}
+
 // conn must be locked
 func (c *conn) close() (err error) {
 	if c.cc == nil {