Merge pull request #163 from ydb-platform/stream-usages

v3.14.5: stream usages

Author: asmyasnikov

.github/workflows/tests.yml

@@ -67,7 +67,7 @@ jobs:
       - name: Checkout code
         uses: actions/checkout@v2
       - name: Test
-        run: go test -race -coverpkg=./... -coverprofile unit.txt -covermode atomic ./test/table_test.go
+        run: go test -race -coverpkg=./... -coverprofile table.txt -covermode atomic ./test/table_test.go
       - name: Upload coverage to Codecov
         uses: codecov/codecov-action@v2
         with:
@@ -106,7 +106,7 @@ jobs:
       - name: Checkout code
         uses: actions/checkout@v2
       - name: Test
-        run: go test -race -coverpkg=./... -coverprofile unit.txt -covermode atomic ./test/ratelimiter_test.go
+        run: go test -race -coverpkg=./... -coverprofile ratelimiter.txt -covermode atomic ./test/ratelimiter_test.go
         shell: bash
       - name: Upload coverage to Codecov
         uses: codecov/codecov-action@v2
@@ -146,13 +146,52 @@ jobs:
       - name: Checkout code
         uses: actions/checkout@v2
       - name: Test
-        run: go test -race -coverpkg=./... -coverprofile unit.txt -covermode atomic ./test/scripting_test.go
+        run: go test -race -coverpkg=./... -coverprofile scripting.txt -covermode atomic ./test/scripting_test.go
       - name: Upload coverage to Codecov
         uses: codecov/codecov-action@v2
         with:
           file: ./scripting.txt
           flags: scripting,e2e,integration,${{ matrix.os }},${{ matrix.go-version }}
           name: scripting
+  discovery:
+    strategy:
+      matrix:
+        go-version: [1.17.x]
+        os: [ubuntu-latest]
+    services:
+      ydb:
+        image: cr.yandex/yc/yandex-docker-local-ydb:latest
+        ports:
+          - 2135:2135
+          - 8765:8765
+        volumes:
+          - /tmp/ydb_certs:/ydb_certs
+        env:
+          YDB_LOCAL_SURVIVE_RESTART: true
+          YDB_USE_IN_MEMORY_PDISKS: true
+        options: '-h localhost'
+    env:
+      OS: ${{ matrix.os }}
+      GO: ${{ matrix.go-version }}
+      YDB_CONNECTION_STRING: grpcs://localhost:2135/?database=/local
+      YDB_SSL_ROOT_CERTIFICATES_FILE: /tmp/ydb_certs/ca.pem
+      YDB_ANONYMOUS_CREDENTIALS: 1
+    runs-on: ${{ matrix.os }}
+    steps:
+      - name: Install Go
+        uses: actions/setup-go@v2
+        with:
+          go-version: ${{ matrix.go-version }}
+      - name: Checkout code
+        uses: actions/checkout@v2
+      - name: Test
+        run: go test -race -coverpkg=./... -coverprofile discovery.txt -covermode atomic ./test/discovery_test.go
+      - name: Upload coverage to Codecov
+        uses: codecov/codecov-action@v2
+        with:
+          file: ./discovery.txt
+          flags: discovery,e2e,integration,${{ matrix.os }},${{ matrix.go-version }}
+          name: discovery
   connection:
     strategy:
       matrix:
@@ -185,7 +224,7 @@ jobs:
       - name: Checkout code
         uses: actions/checkout@v2
       - name: Test
-        run: go test -race -coverpkg=./... -coverprofile unit.txt -covermode atomic ./test/connection_test.go
+        run: go test -race -coverpkg=./... -coverprofile connection.txt -covermode atomic ./test/connection_test.go
       - name: Upload coverage to Codecov
         uses: codecov/codecov-action@v2
         with:

CHANGELOG.md

@@ -9,7 +9,10 @@
 * Refactored `table.CreateSession` as retry operation with options
 * Moved log level from root of repository to package `log`
 * Added details and address to transport error
-* Fixed `recursive` param in `ratelimiter.ListResource` 
+* Fixed `recursive` param in `ratelimiter.ListResource`
+* Added counting stream usages for exclude park connection if it in use
+* Added `trace.Driver` events about change stream usage and `conn.Release()` call
+* Fixed bug with non-applying meta headers on direct call `db.Discovery.Discover`
 
 ## 3.14.4
 * Implemented auto-removing `conn.Conn` from `conn.Pool` with counting usages of `conn.Conn`

internal/conn/config.go

@@ -5,6 +5,7 @@ import (
 
 	"google.golang.org/grpc"
 
+	"github.com/ydb-platform/ydb-go-sdk/v3/internal/meta"
 	"github.com/ydb-platform/ydb-go-sdk/v3/trace"
 )
 
@@ -13,4 +14,5 @@ type Config interface {
 	Trace() trace.Driver
 	ConnectionTTL() time.Duration
 	GrpcDialOptions() []grpc.DialOption
+	Meta() meta.Meta
 }

internal/conn/conn.go

@@ -46,34 +46,57 @@ func (c *conn) Address() string {
 }
 
 type conn struct {
-	sync.RWMutex
-	config    Config // ro access
-	cc        *grpc.ClientConn
-	done      chan struct{}
-	endpoint  endpoint.Endpoint // ro access
-	closed    bool
-	state     State
-	usages    int32
-	lastUsage time.Time
-	onClose   []func(*conn)
+	mtx          sync.RWMutex
+	config       Config // ro access
+	cc           *grpc.ClientConn
+	done         chan struct{}
+	endpoint     endpoint.Endpoint // ro access
+	closed       bool
+	state        State
+	usages       int32
+	streamUsages int32
+	lastUsage    time.Time
+	onClose      []func(*conn)
 }
 
 func (c *conn) Release(ctx context.Context) {
-	if c.decUsages() == 0 {
-		_ = c.Close(ctx)
+	var (
+		onDone = trace.DriverOnConnRelease(
+			c.config.Trace(),
+			&ctx,
+			c.endpoint.Copy(),
+		)
+		err error
+	)
+	defer func() {
+		onDone(err)
+	}()
+	var issues []error
+	if c.changeUsages(-1) == 0 {
+		if usages := atomic.LoadInt32(&c.streamUsages); usages > 0 {
+			issues = append(issues, fmt.Errorf("conn in stream use: usages=%d", usages))
+		}
+		if closeErr := c.Close(ctx); closeErr != nil {
+			issues = append(issues, closeErr)
+		}
+	}
+	if len(issues) > 0 {
+		err = errors.NewWithIssues("conn released with issues", issues...)
 	}
 }
 
 func (c *conn) LastUsage() time.Time {
-	c.RLock()
-	defer c.RUnlock()
-
+	if usages := atomic.LoadInt32(&c.streamUsages); usages > 0 {
+		return time.Now()
+	}
+	c.mtx.RLock()
+	defer c.mtx.RUnlock()
 	return c.lastUsage
 }
 
 func (c *conn) IsState(states ...State) bool {
-	c.RLock()
-	defer c.RUnlock()
+	c.mtx.RLock()
+	defer c.mtx.RUnlock()
 
 	for _, s := range states {
 		if s == c.state {
@@ -85,8 +108,8 @@ func (c *conn) IsState(states ...State) bool {
 }
 
 func (c *conn) park(ctx context.Context) (err error) {
-	c.Lock()
-	defer c.Unlock()
+	c.mtx.Lock()
+	defer c.mtx.Unlock()
 
 	if c.closed {
 		return nil
@@ -129,8 +152,8 @@ func (c *conn) Endpoint() endpoint.Endpoint {
 }
 
 func (c *conn) SetState(s State) State {
-	c.Lock()
-	defer c.Unlock()
+	c.mtx.Lock()
+	defer c.mtx.Unlock()
 	return c.setState(s)
 }
 
@@ -145,8 +168,8 @@ func (c *conn) setState(s State) State {
 }
 
 func (c *conn) GetState() (s State) {
-	c.RLock()
-	defer c.RUnlock()
+	c.mtx.RLock()
+	defer c.mtx.RUnlock()
 	return c.state
 }
 
@@ -168,8 +191,8 @@ func (c *conn) take(ctx context.Context) (cc *grpc.ClientConn, err error) {
 		)
 	}
 
-	c.Lock()
-	defer c.Unlock()
+	c.mtx.Lock()
+	defer c.mtx.Unlock()
 
 	if !isBroken(c.cc) {
 		return c.cc, nil
@@ -194,30 +217,44 @@ func (c *conn) take(ctx context.Context) (cc *grpc.ClientConn, err error) {
 	return c.cc, nil
 }
 
+func (c *conn) touchLastUsage() {
+	c.mtx.Lock()
+	defer c.mtx.Unlock()
+	c.lastUsage = time.Now()
+}
+
 func (c *conn) changeUsages(delta int32) int32 {
-	if usages := atomic.AddInt32(&c.usages, delta); usages < 0 {
+	defer c.touchLastUsage()
+
+	usages := atomic.AddInt32(&c.usages, delta)
+
+	if usages < 0 {
 		panic("negative usages" + strconv.Itoa(int(usages)))
-	} else {
-		trace.DriverOnConnUsagesChange(
-			c.config.Trace(),
-			c.endpoint.Copy(),
-			int(usages),
-		)
-		return usages
 	}
-}
 
-func (c *conn) incUsages() {
-	c.Lock()
-	defer c.Unlock()
-	c.lastUsage = time.Now()
-	c.changeUsages(1)
+	trace.DriverOnConnUsagesChange(
+		c.config.Trace(),
+		c.endpoint.Copy(),
+		int(usages),
+	)
+
+	return usages
 }
 
-func (c *conn) decUsages() int32 {
-	c.Lock()
-	defer c.Unlock()
-	return c.changeUsages(-1)
+func (c *conn) changeStreamUsages(delta int32) {
+	defer c.touchLastUsage()
+
+	usages := atomic.AddInt32(&c.streamUsages, delta)
+
+	if usages < 0 {
+		panic("negative stream usages" + strconv.Itoa(int(usages)))
+	}
+
+	trace.DriverOnConnStreamUsagesChange(
+		c.config.Trace(),
+		c.endpoint.Copy(),
+		int(usages),
+	)
 }
 
 func isBroken(raw *grpc.ClientConn) bool {
@@ -240,14 +277,14 @@ func (c *conn) close() (err error) {
 }
 
 func (c *conn) isClosed() bool {
-	c.RLock()
-	defer c.RUnlock()
+	c.mtx.RLock()
+	defer c.mtx.RUnlock()
 	return c.closed
 }
 
 func (c *conn) Close(ctx context.Context) (err error) {
-	c.Lock()
-	defer c.Unlock()
+	c.mtx.Lock()
+	defer c.mtx.Unlock()
 
 	if c.closed {
 		return nil
@@ -292,8 +329,17 @@ func (c *conn) invoke(
 		)
 	}
 
-	c.incUsages()
-	defer c.decUsages()
+	ctx, err = c.config.Meta().Meta(ctx)
+	if err != nil {
+		return errors.NewGrpcError(
+			codes.Unavailable,
+			errors.WithMsg("ydb driver conn apply meta failed"),
+			errors.WithErr(err),
+		)
+	}
+
+	c.changeUsages(1)
+	defer c.changeUsages(-1)
 
 	return cc.Invoke(ctx, method, req, res, opts...)
 }
@@ -370,8 +416,17 @@ func (c *conn) newStream(
 		)
 	}
 
-	c.incUsages()
-	defer c.decUsages()
+	ctx, err = c.config.Meta().Meta(ctx)
+	if err != nil {
+		return nil, errors.NewGrpcError(
+			codes.Unavailable,
+			errors.WithMsg("ydb driver conn apply meta failed"),
+			errors.WithErr(err),
+		)
+	}
+
+	c.changeStreamUsages(1)
+	defer c.changeStreamUsages(-1)
 
 	return cc.NewStream(ctx, desc, method, opts...)
 }

internal/conn/grpc_client_stream.go

@@ -47,8 +47,8 @@ func (s *grpcClientStream) Context() context.Context {
 }
 
 func (s *grpcClientStream) SendMsg(m interface{}) (err error) {
-	s.c.incUsages()
-	defer s.c.decUsages()
+	s.c.changeStreamUsages(1)
+	defer s.c.changeStreamUsages(-1)
 
 	err = s.s.SendMsg(m)
 	if err != nil && s.wrapping {
@@ -64,8 +64,8 @@ func (s *grpcClientStream) SendMsg(m interface{}) (err error) {
 }
 
 func (s *grpcClientStream) RecvMsg(m interface{}) (err error) {
-	s.c.incUsages()
-	defer s.c.decUsages()
+	s.c.changeStreamUsages(1)
+	defer s.c.changeStreamUsages(-1)
 
 	defer func() {
 		onDone := s.recv(errors.HideEOF(err))

internal/conn/pool.go

@@ -62,7 +62,7 @@ func (p *pool) Get(ctx context.Context, endpoint endpoint.Endpoint) Conn {
 	)
 
 	if cc, has = p.conns[address]; has {
-		cc.incUsages()
+		cc.changeUsages(1)
 		return cc
 	}