Merge pull request #1720 from ydb-platform/session-limit

* Improved the `ydb.WithSessionPoolSessionUsageLimit` option for allow `time.Duration` as argument type for limit max session time to live since create time

Author: asmyasnikov

CHANGELOG.md

@@ -1,3 +1,5 @@
+* Improved the `ydb.WithSessionPoolSessionUsageLimit` option for allow `time.Duration` as argument type for limit max session time to live since create time 
+
 ## v3.105.1
 * Changed the gRPC DNS balancer policy to `round_robin` for internal `discovery/ListEndpoints` call (reverted v3.90.2 changes)
 

internal/pool/pool.go

@@ -42,9 +42,11 @@ type (
 		closeItemFunc      func(ctx context.Context, item PT)
 		idleTimeToLive     time.Duration
 		itemUsageLimit     uint64
+		itemUsageTTL       time.Duration
 	}
 	itemInfo[PT ItemConstraint[T], T any] struct {
 		idle       *xlist.Element[PT]
+		created    time.Time
 		lastUsage  time.Time
 		useCounter *uint64
 	}
@@ -113,6 +115,12 @@ func WithItemUsageLimit[PT ItemConstraint[T], T any](itemUsageLimit uint64) Opti
 	}
 }
 
+func WithItemUsageTTL[PT ItemConstraint[T], T any](ttl time.Duration) Option[PT, T] {
+	return func(c *Config[PT, T]) {
+		c.itemUsageTTL = ttl
+	}
+}
+
 func WithTrace[PT ItemConstraint[T], T any](t *Trace) Option[PT, T] {
 	return func(c *Config[PT, T]) {
 		c.trace = t
@@ -235,9 +243,13 @@ func makeAsyncCreateItemFunc[PT ItemConstraint[T], T any]( //nolint:funlen
 			newItem, err := p.config.createItemFunc(createCtx)
 			if newItem != nil {
 				p.mu.WithLock(func() {
-					var useCounter uint64
+					var (
+						useCounter uint64
+						now        = p.config.clock.Now()
+					)
 					p.index[newItem] = itemInfo[PT, T]{
-						lastUsage:  p.config.clock.Now(),
+						created:    now,
+						lastUsage:  now,
 						useCounter: &useCounter,
 					}
 				})
@@ -695,6 +707,53 @@ func (p *Pool[PT, T]) pushIdle(item PT, now time.Time) {
 
 const maxAttempts = 100
 
+func needCloseItemByMaxUsage[PT ItemConstraint[T], T any](c *Config[PT, T], info itemInfo[PT, T]) bool {
+	if c.itemUsageLimit <= 0 {
+		return false
+	}
+	if *info.useCounter < c.itemUsageLimit {
+		return false
+	}
+
+	return true
+}
+
+func needCloseItemByTTL[PT ItemConstraint[T], T any](c *Config[PT, T], info itemInfo[PT, T]) bool {
+	if c.itemUsageTTL <= 0 {
+		return false
+	}
+	if c.clock.Since(info.created) < c.itemUsageTTL {
+		return false
+	}
+
+	return true
+}
+
+func needCloseItemByIdleTTL[PT ItemConstraint[T], T any](c *Config[PT, T], info itemInfo[PT, T]) bool {
+	if c.idleTimeToLive <= 0 {
+		return false
+	}
+	if c.clock.Since(info.lastUsage) < c.idleTimeToLive {
+		return false
+	}
+
+	return true
+}
+
+func needCloseItem[PT ItemConstraint[T], T any](c *Config[PT, T], info itemInfo[PT, T]) bool {
+	if needCloseItemByMaxUsage(c, info) {
+		return true
+	}
+	if needCloseItemByTTL(c, info) {
+		return true
+	}
+	if needCloseItemByIdleTTL(c, info) {
+		return true
+	}
+
+	return false
+}
+
 func (p *Pool[PT, T]) getItem(ctx context.Context) (item PT, finalErr error) { //nolint:funlen
 	var (
 		start   = p.config.clock.Now()
@@ -749,8 +808,7 @@ func (p *Pool[PT, T]) getItem(ctx context.Context) (item PT, finalErr error) { /
 					return info
 				})
 
-				if (p.config.itemUsageLimit > 0 && *info.useCounter > p.config.itemUsageLimit) ||
-					(p.config.idleTimeToLive > 0 && p.config.clock.Since(info.lastUsage) > p.config.idleTimeToLive) {
+				if needCloseItem(&p.config, info) {
 					p.closeItem(ctx, item,
 						closeItemWithLock(),
 						closeItemNotifyStats(),

internal/query/client.go

@@ -575,6 +575,7 @@ func New(ctx context.Context, cc grpc.ClientConnInterface, cfg *config.Config) *
 		pool: pool.New(ctx,
 			pool.WithLimit[*Session, Session](cfg.PoolLimit()),
 			pool.WithItemUsageLimit[*Session, Session](cfg.PoolSessionUsageLimit()),
+			pool.WithItemUsageTTL[*Session, Session](cfg.PoolSessionUsageTTL()),
 			pool.WithTrace[*Session, Session](poolTrace(cfg.Trace())),
 			pool.WithCreateItemTimeout[*Session, Session](cfg.SessionCreateTimeout()),
 			pool.WithCloseItemTimeout[*Session, Session](cfg.SessionDeleteTimeout()),

internal/query/config/config.go

@@ -19,6 +19,7 @@ type Config struct {
 
 	poolLimit             int
 	poolSessionUsageLimit uint64
+	poolSessionUsageTTL   time.Duration
 
 	sessionCreateTimeout   time.Duration
 	sessionDeleteTimeout   time.Duration
@@ -65,6 +66,10 @@ func (c *Config) PoolSessionUsageLimit() uint64 {
 	return c.poolSessionUsageLimit
 }
 
+func (c *Config) PoolSessionUsageTTL() time.Duration {
+	return c.poolSessionUsageTTL
+}
+
 // SessionCreateTimeout limits maximum time spent on Create session request
 func (c *Config) SessionCreateTimeout() time.Duration {
 	return c.sessionCreateTimeout

internal/query/config/options.go

@@ -34,9 +34,17 @@ func WithPoolLimit(size int) Option {
 	}
 }
 
-func WithPoolSessionUsageLimit(sessionUsageLimit uint64) Option {
+// WithSessionPoolSessionUsageLimit set pool session max usage:
+// - if argument type is uint64 - WithSessionPoolSessionUsageLimit limits max usage count of pool session
+// - if argument type is time.Duration - WithSessionPoolSessionUsageLimit limits max time to live of pool session
+func WithSessionPoolSessionUsageLimit[T interface{ uint64 | time.Duration }](limit T) Option {
 	return func(c *Config) {
-		c.poolSessionUsageLimit = sessionUsageLimit
+		switch v := any(limit).(type) {
+		case uint64:
+			c.poolSessionUsageLimit = v
+		case time.Duration:
+			c.poolSessionUsageTTL = v
+		}
 	}
 }
 

internal/table/client.go

@@ -42,6 +42,7 @@ func New(ctx context.Context, cc grpc.ClientConnInterface, config *config.Config
 		pool: pool.New[*Session, Session](ctx,
 			pool.WithLimit[*Session, Session](config.SizeLimit()),
 			pool.WithItemUsageLimit[*Session, Session](config.SessionUsageLimit()),
+			pool.WithItemUsageTTL[*Session, Session](config.SessionUsageTTL()),
 			pool.WithIdleTimeToLive[*Session, Session](config.IdleThreshold()),
 			pool.WithCreateItemTimeout[*Session, Session](config.CreateSessionTimeout()),
 			pool.WithCloseItemTimeout[*Session, Session](config.DeleteTimeout()),

internal/table/config/config.go

@@ -52,19 +52,27 @@ func With(config config.Common) Option {
 }
 
 // WithSizeLimit defines upper bound of pooled sessions.
-// If sizeLimit is less than or equal to zero then the
+// If poolLimit is less than or equal to zero then the
 // DefaultSessionPoolSizeLimit variable is used as a limit.
 func WithSizeLimit(sizeLimit int) Option {
 	return func(c *Config) {
 		if sizeLimit > 0 {
-			c.sizeLimit = sizeLimit
+			c.poolLimit = sizeLimit
 		}
 	}
 }
 
-func WithPoolSessionUsageLimit(sessionUsageLimit uint64) Option {
+// WithSessionPoolSessionUsageLimit set pool session max usage:
+// - if argument type is uint64 - WithSessionPoolSessionUsageLimit limits max usage count of pool session
+// - if argument type is time.Duration - WithSessionPoolSessionUsageLimit limits max time to live of pool session
+func WithSessionPoolSessionUsageLimit[T interface{ uint64 | time.Duration }](limit T) Option {
 	return func(c *Config) {
-		c.sessionUsageLimit = sessionUsageLimit
+		switch v := any(limit).(type) {
+		case uint64:
+			c.poolSessionUsageLimit = v
+		case time.Duration:
+			c.poolSessionUsageTTL = v
+		}
 	}
 }
 
@@ -182,8 +190,9 @@ func WithClock(clock clockwork.Clock) Option {
 type Config struct {
 	config.Common
 
-	sizeLimit         int
-	sessionUsageLimit uint64
+	poolLimit             int
+	poolSessionUsageLimit uint64
+	poolSessionUsageTTL   time.Duration
 
 	createSessionTimeout time.Duration
 	deleteTimeout        time.Duration
@@ -212,11 +221,15 @@ func (c *Config) Clock() clockwork.Clock {
 // If SizeLimit is less than or equal to zero then the
 // DefaultSessionPoolSizeLimit variable is used as a limit.
 func (c *Config) SizeLimit() int {
-	return c.sizeLimit
+	return c.poolLimit
 }
 
 func (c *Config) SessionUsageLimit() uint64 {
-	return c.sessionUsageLimit
+	return c.poolSessionUsageLimit
+}
+
+func (c *Config) SessionUsageTTL() time.Duration {
+	return c.poolSessionUsageTTL
 }
 
 // KeepAliveMinSize is a lower bound for sessions in the pool. If there are more sessions open, then
@@ -293,7 +306,7 @@ func (c *Config) DeleteTimeout() time.Duration {
 
 func defaults() *Config {
 	return &Config{
-		sizeLimit:            DefaultSessionPoolSizeLimit,
+		poolLimit:            DefaultSessionPoolSizeLimit,
 		createSessionTimeout: DefaultSessionPoolCreateSessionTimeout,
 		deleteTimeout:        DefaultSessionPoolDeleteTimeout,
 		idleThreshold:        DefaultSessionPoolIdleThreshold,

options.go

@@ -531,11 +531,13 @@ func WithSessionPoolSizeLimit(sizeLimit int) Option {
 	}
 }
 
-// WithSessionPoolSessionUsageLimit set max count for use session
-func WithSessionPoolSessionUsageLimit(sessionUsageLimit uint64) Option {
+// WithSessionPoolSessionUsageLimit set pool session max usage:
+// - if argument type is uint64 - WithSessionPoolSessionUsageLimit limits max usage count of pool session
+// - if argument type is time.Duration - WithSessionPoolSessionUsageLimit limits max time to live of pool session
+func WithSessionPoolSessionUsageLimit[T interface{ uint64 | time.Duration }](limit T) Option {
 	return func(ctx context.Context, d *Driver) error {
-		d.tableOptions = append(d.tableOptions, tableConfig.WithPoolSessionUsageLimit(sessionUsageLimit))
-		d.queryOptions = append(d.queryOptions, queryConfig.WithPoolSessionUsageLimit(sessionUsageLimit))
+		d.tableOptions = append(d.tableOptions, tableConfig.WithSessionPoolSessionUsageLimit(limit))
+		d.queryOptions = append(d.queryOptions, queryConfig.WithSessionPoolSessionUsageLimit(limit))
 
 		return nil
 	}