internal/pool async create and close item

Author: asmyasnikov

internal/pool/pool.go

@@ -56,7 +56,7 @@ func TestPool(t *testing.T) {
 		})
 		t.Run("WithLimit", func(t *testing.T) {
 			p := New[*testItem, testItem](rootCtx, WithLimit[*testItem, testItem](1))
-			require.EqualValues(t, 1, p.limit)
+			require.EqualValues(t, 1, p.config.limit)
 		})
 		t.Run("WithCreateFunc", func(t *testing.T) {
 			var newCounter int64
@@ -73,7 +73,7 @@ func TestPool(t *testing.T) {
 				return nil
 			})
 			require.NoError(t, err)
-			require.EqualValues(t, p.limit, atomic.LoadInt64(&newCounter))
+			require.EqualValues(t, p.config.limit, atomic.LoadInt64(&newCounter))
 		})
 	})
 	t.Run("Retry", func(t *testing.T) {
@@ -247,6 +247,9 @@ func TestPool(t *testing.T) {
 						return v, nil
 					}),
 				)
+				p.closeItem = func(ctx context.Context, item *testItem) {
+					_ = item.Close(ctx)
+				}
 				err := p.With(rootCtx, func(ctx context.Context, testItem *testItem) error {
 					if atomic.LoadInt64(&newItems) < 10 {
 						return expErr
@@ -265,7 +268,11 @@ func TestPool(t *testing.T) {
 	})
 	t.Run("Stress", func(t *testing.T) {
 		xtest.TestManyTimes(t, func(t testing.TB) {
-			p := New[*testItem, testItem](rootCtx)
+			trace := *defaultTrace
+			trace.OnChange = func(info ChangeInfo) {
+				require.GreaterOrEqual(t, info.Limit, info.Idle)
+			}
+			p := New[*testItem, testItem](rootCtx, WithTrace[*testItem, testItem](&trace))
 			var wg sync.WaitGroup
 			wg.Add(DefaultLimit*2 + 1)
 			for range make([]struct{}, DefaultLimit*2) {
@@ -290,7 +297,12 @@ func TestPool(t *testing.T) {
 	})
 	t.Run("ParallelCreation", func(t *testing.T) {
 		xtest.TestManyTimes(t, func(t testing.TB) {
-			p := New[*testItem, testItem](rootCtx)
+			trace := *defaultTrace
+			trace.OnChange = func(info ChangeInfo) {
+				require.Equal(t, DefaultLimit, info.Limit)
+				require.LessOrEqual(t, info.Idle, DefaultLimit)
+			}
+			p := New[*testItem, testItem](rootCtx, WithTrace[*testItem, testItem](&trace))
 			var wg sync.WaitGroup
 			for range make([]struct{}, DefaultLimit*10) {
 				wg.Add(1)
@@ -303,15 +315,11 @@ func TestPool(t *testing.T) {
 						t.Failed()
 					}
 					stats := p.Stats()
-					require.LessOrEqual(t, stats.Idle+stats.InUse, DefaultLimit)
+					require.LessOrEqual(t, stats.Idle, DefaultLimit)
 				}()
 			}
 
 			wg.Wait()
-			stats := p.Stats()
-			require.Equal(t, DefaultLimit, stats.Limit)
-			require.Equal(t, 0, stats.InUse)
-			require.LessOrEqual(t, stats.Idle, DefaultLimit)
 		}, xtest.StopAfter(14*time.Second))
 	})
 }

internal/pool/pool_test.go

@@ -1,9 +1,6 @@
 package pool
 
 type Stats struct {
-	Limit            int
-	Index            int
-	Idle             int
-	InUse            int
-	CreateInProgress int
+	Limit int
+	Idle  int
 }

internal/pool/stats.go

@@ -2,7 +2,6 @@ package query
 
 import (
 	"context"
-	"sync/atomic"
 	"time"
 
 	"github.com/ydb-platform/ydb-go-genproto/Ydb_Query_V1"
@@ -43,7 +42,6 @@ type (
 	}
 	poolStub struct {
 		createSession func(ctx context.Context) (*Session, error)
-		InUse         atomic.Int32
 	}
 	Client struct {
 		config             *config.Config
@@ -264,20 +262,13 @@ func (p *poolStub) Close(ctx context.Context) error {
 func (p *poolStub) Stats() pool.Stats {
 	return pool.Stats{
 		Limit: -1,
-		Index: 0,
 		Idle:  0,
-		InUse: int(p.InUse.Load()),
 	}
 }
 
 func (p *poolStub) With(
 	ctx context.Context, f func(ctx context.Context, s *Session) error, opts ...retry.Option,
 ) error {
-	p.InUse.Add(1)
-	defer func() {
-		p.InUse.Add(-1)
-	}()
-
 	err := retry.Retry(ctx, func(ctx context.Context) (err error) {
 		s, err := p.createSession(ctx)
 		if err != nil {
@@ -710,7 +701,7 @@ func poolTrace(t *trace.Query) *pool.Trace {
 			}
 		},
 		OnChange: func(info pool.ChangeInfo) {
-			trace.QueryOnPoolChange(t, info.Limit, info.Index, info.Idle, info.InUse)
+			trace.QueryOnPoolChange(t, info.Limit, info.Idle)
 		},
 	}
 }

internal/query/client.go

@@ -5,18 +5,18 @@ import (
 	"sync/atomic"
 )
 
-type Set[K comparable] struct {
+type Set[T comparable] struct {
 	m    sync.Map
 	size atomic.Int32
 }
 
-func (s *Set[K]) Has(key K) bool {
+func (s *Set[T]) Has(key T) bool {
 	_, ok := s.m.Load(key)
 
 	return ok
 }
 
-func (s *Set[K]) Add(key K) bool {
+func (s *Set[T]) Add(key T) bool {
 	_, exists := s.m.LoadOrStore(key, struct{}{})
 
 	if !exists {
@@ -26,17 +26,29 @@ func (s *Set[K]) Add(key K) bool {
 	return !exists
 }
 
-func (s *Set[K]) Size() int {
+func (s *Set[T]) Size() int {
 	return int(s.size.Load())
 }
 
-func (s *Set[K]) Range(f func(key K) bool) {
+func (s *Set[T]) Range(f func(key T) bool) {
 	s.m.Range(func(k, v any) bool {
-		return f(k.(K)) //nolint:forcetypeassert
+		return f(k.(T)) //nolint:forcetypeassert
 	})
 }
 
-func (s *Set[K]) Remove(key K) bool {
+func (s *Set[T]) Values() []T {
+	values := make([]T, 0, s.size.Load())
+
+	s.m.Range(func(k, v any) bool {
+		values = append(values, k.(T))
+
+		return true
+	})
+
+	return values
+}
+
+func (s *Set[T]) Remove(key T) bool {
 	_, exists := s.m.LoadAndDelete(key)
 
 	if exists {
@@ -46,7 +58,7 @@ func (s *Set[K]) Remove(key K) bool {
 	return exists
 }
 
-func (s *Set[K]) Clear() (removed int) {
+func (s *Set[T]) Clear() (removed int) {
 	s.m.Range(func(k, v any) bool {
 		removed++
 

internal/xsync/set.go

@@ -0,0 +1,23 @@
+package xsync
+
+type Value[T any] struct {
+	v  T
+	mu RWMutex
+}
+
+func NewValue[T any](initValue T) *Value[T] {
+	return &Value[T]{v: initValue}
+}
+
+func (v *Value[T]) Get() T {
+	v.mu.RLock()
+	defer v.mu.RUnlock()
+
+	return v.v
+}
+
+func (v *Value[T]) Change(change func(old T) T) {
+	v.mu.WithLock(func() {
+		v.v = change(v.v)
+	})
+}

internal/xsync/value.go

@@ -0,0 +1,16 @@
+package xsync
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestValue(t *testing.T) {
+	v := NewValue(5)
+	require.Equal(t, 5, v.Get())
+	v.Change(func(old int) int {
+		return 6
+	})
+	require.Equal(t, 6, v.Get())
+}

internal/xsync/value_test.go

@@ -42,8 +42,6 @@ func query(config Config) (t trace.Query) {
 			sizeConfig := poolConfig.WithSystem("size")
 			limit := sizeConfig.GaugeVec("limit")
 			idle := sizeConfig.GaugeVec("idle")
-			index := sizeConfig.GaugeVec("index")
-			inUse := sizeConfig.WithSystem("in").GaugeVec("use")
 
 			t.OnPoolChange = func(stats trace.QueryPoolChange) {
 				if sizeConfig.Details()&trace.QueryPoolEvents == 0 {
@@ -52,8 +50,6 @@ func query(config Config) (t trace.Query) {
 
 				limit.With(nil).Set(float64(stats.Limit))
 				idle.With(nil).Set(float64(stats.Idle))
-				inUse.With(nil).Set(float64(stats.InUse))
-				index.With(nil).Set(float64(stats.Index))
 			}
 		}
 	}

metrics/query.go

@@ -547,8 +547,6 @@ type (
 	// Internals: https://github.com/ydb-platform/ydb-go-sdk/blob/master/VERSIONING.md#internals
 	QueryPoolChange struct {
 		Limit int
-		Index int
 		Idle  int
-		InUse int
 	}
 )