Skip to content

Commit fb252b8

Browse files
asmyasnikovasmyasnikov
authored
Merge pull request #1457 from ydb-platform/query-session
extracted query session controller logic into separated package
2 parents a06c3b5 + 802e7fb commit fb252b8

File tree

13 files changed

+466
-363
lines changed

13 files changed

+466
-363
lines changed

‎internal/pool/pool.go‎

Lines changed: 23 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -17,12 +17,15 @@ import (
1717
)
1818

1919
type (
20-
Item[T any] interface {
21-
*T
20+
Item interface {
2221
IsAlive() bool
2322
Close(ctx context.Context) error
2423
}
25-
Config[PT Item[T], T any] struct {
24+
ItemConstraint[T any] interface {
25+
*T
26+
Item
27+
}
28+
Config[PT ItemConstraint[T], T any] struct {
2629
trace *Trace
2730
clock clockwork.Clock
2831
limit int
@@ -32,15 +35,15 @@ type (
3235
closeItem func(ctx context.Context, item PT)
3336
idleThreshold time.Duration
3437
}
35-
itemInfo[PT Item[T], T any] struct {
38+
itemInfo[PT ItemConstraint[T], T any] struct {
3639
idle *xlist.Element[PT]
3740
touched time.Time
3841
}
39-
waitChPool[PT Item[T], T any] interface {
42+
waitChPool[PT ItemConstraint[T], T any] interface {
4043
GetOrNew() *chan PT
4144
Put(t *chan PT)
4245
}
43-
Pool[PT Item[T], T any] struct {
46+
Pool[PT ItemConstraint[T], T any] struct {
4447
config Config[PT, T]
4548

4649
createItem func(ctx context.Context) (PT, error)
@@ -55,62 +58,62 @@ type (
5558

5659
done chan struct{}
5760
}
58-
option[PT Item[T], T any] func(c *Config[PT, T])
61+
Option[PT ItemConstraint[T], T any] func(c *Config[PT, T])
5962
)
6063

61-
func WithCreateItemFunc[PT Item[T], T any](f func(ctx context.Context) (PT, error)) option[PT, T] {
64+
func WithCreateItemFunc[PT ItemConstraint[T], T any](f func(ctx context.Context) (PT, error)) Option[PT, T] {
6265
return func(c *Config[PT, T]) {
6366
c.createItem = f
6467
}
6568
}
6669

67-
func WithSyncCloseItem[PT Item[T], T any]() option[PT, T] {
70+
func WithSyncCloseItem[PT ItemConstraint[T], T any]() Option[PT, T] {
6871
return func(c *Config[PT, T]) {
6972
c.closeItem = func(ctx context.Context, item PT) {
7073
_ = item.Close(ctx)
7174
}
7275
}
7376
}
7477

75-
func WithCreateItemTimeout[PT Item[T], T any](t time.Duration) option[PT, T] {
78+
func WithCreateItemTimeout[PT ItemConstraint[T], T any](t time.Duration) Option[PT, T] {
7679
return func(c *Config[PT, T]) {
7780
c.createTimeout = t
7881
}
7982
}
8083

81-
func WithCloseItemTimeout[PT Item[T], T any](t time.Duration) option[PT, T] {
84+
func WithCloseItemTimeout[PT ItemConstraint[T], T any](t time.Duration) Option[PT, T] {
8285
return func(c *Config[PT, T]) {
8386
c.closeTimeout = t
8487
}
8588
}
8689

87-
func WithLimit[PT Item[T], T any](size int) option[PT, T] {
90+
func WithLimit[PT ItemConstraint[T], T any](size int) Option[PT, T] {
8891
return func(c *Config[PT, T]) {
8992
c.limit = size
9093
}
9194
}
9295

93-
func WithTrace[PT Item[T], T any](t *Trace) option[PT, T] {
96+
func WithTrace[PT ItemConstraint[T], T any](t *Trace) Option[PT, T] {
9497
return func(c *Config[PT, T]) {
9598
c.trace = t
9699
}
97100
}
98101

99-
func WithIdleThreshold[PT Item[T], T any](idleThreshold time.Duration) option[PT, T] {
102+
func WithIdleThreshold[PT ItemConstraint[T], T any](idleThreshold time.Duration) Option[PT, T] {
100103
return func(c *Config[PT, T]) {
101104
c.idleThreshold = idleThreshold
102105
}
103106
}
104107

105-
func WithClock[PT Item[T], T any](clock clockwork.Clock) option[PT, T] {
108+
func WithClock[PT ItemConstraint[T], T any](clock clockwork.Clock) Option[PT, T] {
106109
return func(c *Config[PT, T]) {
107110
c.clock = clock
108111
}
109112
}
110113

111-
func New[PT Item[T], T any](
114+
func New[PT ItemConstraint[T], T any](
112115
ctx context.Context,
113-
opts ...option[PT, T],
116+
opts ...Option[PT, T],
114117
) *Pool[PT, T] {
115118
p := &Pool[PT, T]{
116119
config: Config[PT, T]{
@@ -162,14 +165,14 @@ func New[PT Item[T], T any](
162165
}
163166

164167
// defaultCreateItem returns a new item
165-
func defaultCreateItem[T any, PT Item[T]](context.Context) (PT, error) {
168+
func defaultCreateItem[T any, PT ItemConstraint[T]](context.Context) (PT, error) {
166169
var item T
167170

168171
return &item, nil
169172
}
170173

171174
// makeAsyncCreateItemFunc wraps the createItem function with timeout handling
172-
func makeAsyncCreateItemFunc[PT Item[T], T any]( //nolint:funlen
175+
func makeAsyncCreateItemFunc[PT ItemConstraint[T], T any]( //nolint:funlen
173176
p *Pool[PT, T],
174177
) func(ctx context.Context) (PT, error) {
175178
return func(ctx context.Context) (PT, error) {
@@ -277,7 +280,7 @@ func (p *Pool[PT, T]) Stats() Stats {
277280
return p.stats()
278281
}
279282

280-
func makeAsyncCloseItemFunc[PT Item[T], T any](
283+
func makeAsyncCloseItemFunc[PT ItemConstraint[T], T any](
281284
p *Pool[PT, T],
282285
) func(ctx context.Context, item PT) {
283286
return func(ctx context.Context, item PT) {

‎internal/pool/pool_test.go‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -135,7 +135,7 @@ func caller() string {
135135
return fmt.Sprintf("%s:%d", path.Base(file), line)
136136
}
137137

138-
func mustGetItem[PT Item[T], T any](t testing.TB, p *Pool[PT, T]) PT {
138+
func mustGetItem[PT ItemConstraint[T], T any](t testing.TB, p *Pool[PT, T]) PT {
139139
s, err := p.getItem(context.Background())
140140
if err != nil {
141141
t.Helper()
@@ -145,7 +145,7 @@ func mustGetItem[PT Item[T], T any](t testing.TB, p *Pool[PT, T]) PT {
145145
return s
146146
}
147147

148-
func mustPutItem[PT Item[T], T any](t testing.TB, p *Pool[PT, T], item PT) {
148+
func mustPutItem[PT ItemConstraint[T], T any](t testing.TB, p *Pool[PT, T], item PT) {
149149
if err := p.putItem(context.Background(), item); err != nil {
150150
t.Helper()
151151
t.Fatalf("%s: %v", caller(), err)

‎internal/query/client.go‎

Lines changed: 24 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@ import (
1616
"github.com/ydb-platform/ydb-go-sdk/v3/internal/query/config"
1717
"github.com/ydb-platform/ydb-go-sdk/v3/internal/query/options"
1818
"github.com/ydb-platform/ydb-go-sdk/v3/internal/query/result"
19+
"github.com/ydb-platform/ydb-go-sdk/v3/internal/query/session"
1920
"github.com/ydb-platform/ydb-go-sdk/v3/internal/query/tx"
2021
"github.com/ydb-platform/ydb-go-sdk/v3/internal/stack"
2122
"github.com/ydb-platform/ydb-go-sdk/v3/internal/types"
@@ -41,9 +42,9 @@ type (
4142
With(ctx context.Context, f func(ctx context.Context, s *Session) error, opts ...retry.Option) error
4243
}
4344
Client struct {
44-
config *config.Config
45-
queryServiceClient Ydb_Query_V1.QueryServiceClient
46-
pool sessionPool
45+
config *config.Config
46+
client Ydb_Query_V1.QueryServiceClient
47+
pool sessionPool
4748

4849
done chan struct{}
4950
}
@@ -100,7 +101,7 @@ func (c *Client) FetchScriptResults(ctx context.Context,
100101
opID string, opts ...options.FetchScriptOption,
101102
) (*options.FetchScriptResult, error) {
102103
r, err := retry.RetryWithResult(ctx, func(ctx context.Context) (*options.FetchScriptResult, error) {
103-
r, err := fetchScriptResults(ctx, c.queryServiceClient, opID,
104+
r, err := fetchScriptResults(ctx, c.client, opID,
104105
append(opts, func(request *options.FetchScriptResultsRequest) {
105106
request.Trace = c.config.Trace()
106107
})...,
@@ -175,7 +176,7 @@ func (c *Client) ExecuteScript(
175176

176177
request, grpcOpts := executeQueryScriptRequest(a, q, settings)
177178

178-
op, err = executeScript(ctx, c.queryServiceClient, request, grpcOpts...)
179+
op, err = executeScript(ctx, c.client, request, grpcOpts...)
179180
if err != nil {
180181
return op, xerrors.WithStackTrace(err)
181182
}
@@ -200,18 +201,18 @@ func do(
200201
opts ...retry.Option,
201202
) (finalErr error) {
202203
err := pool.With(ctx, func(ctx context.Context, s *Session) error {
203-
s.setStatus(statusInUse)
204+
s.SetStatus(session.StatusInUse)
204205

205206
err := op(ctx, s)
206207
if err != nil {
207208
if xerrors.IsOperationError(err) {
208-
s.setStatus(statusClosed)
209+
s.SetStatus(session.StatusClosed)
209210
}
210211

211212
return xerrors.WithStackTrace(err)
212213
}
213214

214-
s.setStatus(statusIdle)
215+
s.SetStatus(session.StatusIdle)
215216

216217
return nil
217218
}, opts...)
@@ -337,7 +338,7 @@ func (c *Client) QueryRow(ctx context.Context, q string, opts ...options.Execute
337338
func clientExec(ctx context.Context, pool sessionPool, q string, opts ...options.Execute) (finalErr error) {
338339
settings := options.ExecuteSettings(opts...)
339340
err := do(ctx, pool, func(ctx context.Context, s *Session) (err error) {
340-
_, r, err := execute(ctx, s.id, s.queryServiceClient, q, settings, withTrace(s.cfg.Trace()))
341+
_, r, err := execute(ctx, s.ID(), s.client, q, settings, withTrace(s.trace))
341342
if err != nil {
342343
return xerrors.WithStackTrace(err)
343344
}
@@ -381,8 +382,8 @@ func clientQuery(ctx context.Context, pool sessionPool, q string, opts ...option
381382
) {
382383
settings := options.ExecuteSettings(opts...)
383384
err = do(ctx, pool, func(ctx context.Context, s *Session) (err error) {
384-
_, streamResult, err := execute(ctx, s.id, s.queryServiceClient, q,
385-
options.ExecuteSettings(opts...), withTrace(s.cfg.Trace()),
385+
_, streamResult, err := execute(ctx, s.ID(), s.client, q,
386+
options.ExecuteSettings(opts...), withTrace(s.trace),
386387
)
387388
if err != nil {
388389
return xerrors.WithStackTrace(err)
@@ -433,7 +434,7 @@ func clientQueryResultSet(
433434
ctx context.Context, pool sessionPool, q string, settings executeSettings, resultOpts ...resultOption,
434435
) (rs result.ClosableResultSet, finalErr error) {
435436
err := do(ctx, pool, func(ctx context.Context, s *Session) error {
436-
_, r, err := execute(ctx, s.id, s.queryServiceClient, q, settings, resultOpts...)
437+
_, r, err := execute(ctx, s.ID(), s.client, q, settings, resultOpts...)
437438
if err != nil {
438439
return xerrors.WithStackTrace(err)
439440
}
@@ -512,18 +513,18 @@ func (c *Client) DoTx(ctx context.Context, op query.TxOperation, opts ...options
512513
return nil
513514
}
514515

515-
func New(ctx context.Context, balancer grpc.ClientConnInterface, cfg *config.Config) *Client {
516+
func New(ctx context.Context, cc grpc.ClientConnInterface, cfg *config.Config) *Client {
516517
onDone := trace.QueryOnNew(cfg.Trace(), &ctx,
517518
stack.FunctionID("github.com/ydb-platform/ydb-go-sdk/v3/internal/query.New"),
518519
)
519520
defer onDone()
520521

521-
grpcClient := Ydb_Query_V1.NewQueryServiceClient(balancer)
522+
client := Ydb_Query_V1.NewQueryServiceClient(cc)
522523

523-
client := &Client{
524-
config: cfg,
525-
queryServiceClient: grpcClient,
526-
done: make(chan struct{}),
524+
return &Client{
525+
config: cfg,
526+
client: client,
527+
done: make(chan struct{}),
527528
pool: pool.New(ctx,
528529
pool.WithLimit[*Session, Session](cfg.PoolLimit()),
529530
pool.WithTrace[*Session, Session](poolTrace(cfg.Trace())),
@@ -541,7 +542,11 @@ func New(ctx context.Context, balancer grpc.ClientConnInterface, cfg *config.Con
541542
}
542543
defer cancelCreate()
543544

544-
s, err := createSession(createCtx, grpcClient, cfg)
545+
s, err := createSession(createCtx, client,
546+
session.WithConn(cc),
547+
session.WithDeleteTimeout(cfg.SessionDeleteTimeout()),
548+
session.WithTrace(cfg.Trace()),
549+
)
545550
if err != nil {
546551
return nil, xerrors.WithStackTrace(err)
547552
}
@@ -550,8 +555,6 @@ func New(ctx context.Context, balancer grpc.ClientConnInterface, cfg *config.Con
550555
}),
551556
),
552557
}
553-
554-
return client
555558
}
556559

557560
func poolTrace(t *trace.Query) *pool.Trace {

0 commit comments

Comments
 (0)